针对混合模型推理内存瓶颈，提出异构页管理方案提升利用率
Asymmetric Virtual Memory Paging for Hybrid Mamba-Transformer Inference
- 将两类缓存分置物理池，虚拟地址统一，按需动态迁移容量
- 在合成与真实数据上，内存溢出减少7.6%，吞吐最高提升13.3倍
- 适合处理提示分布多变的混合注意力-状态空间模型推理场景
混合语言模型如Jamba将注意力层与状态空间模型（SSM）结合，产生两种内存缓存：键值（KV）缓存随序列长度线性增长，而SSM状态每层固定。当前推理引擎处理不佳：统一池需对齐页大小，使SSM状态浪费达7.3倍容量；静态双池无法适应请求间提示分布变化。本文提出异构虚拟内存分页（AVMP），将两类缓存置于物理分离池中，通过统一虚拟地址空间实现容量动态迁移，仅在分配失败时触发，保证行为确定性。在RTX 3060 12GB上评估，270个合成任务及60个ShareGPT trace回放显示，内存溢出下降7.6%，合成工作负载吞吐提升1.83–13.3倍，真实场景提升2.36倍，所有增益均通过配对自举95%置信区间验证。相位时间分析揭示双重机制：容量压力下缩短OOM恢复时间，KV密集负载下加速分配调用。实现纯Python，Triton集成为后续工作。
原文摘要 · Abstract (English)
Hybrid language models like Jamba mix attention layers with State Space Models (SSMs), creating two memory cache types with opposite profiles: Key-Value (KV) caches grow linearly with sequence length, while SSM states stay fixed per layer. Current inference engines handle this poorly. Unified pools pad SSM states to attention page sizes, wasting up to 7.3x capacity. Static dual pools cannot adapt when prompt distributions shift between requests. We present Asymmetric Virtual Memory Paging (AVMP). The allocator separates the two cache types into physically distinct pools behind a unified virtual address space, and migrates capacity between pools when one runs out. Migration triggers only on allocation failure, keeping behavior deterministic. We evaluate AVMP across 270 synthetic cells plus 60 cells of ShareGPT trace replay on an RTX 3060 12GB. Out-of-Memory events drop 7.6% and request throughput improves 1.83x to 13.3x across synthetic workloads and 2.36x on ShareGPT. All gains hold under paired-bootstrap 95% confidence intervals. A phase-time breakdown reveals two distinct mechanisms: shorter OOM recovery on capacity-pressured workloads, and faster allocation calls on KV-heavy workloads. Implementation is pure Python; Triton integration is future work.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。