混合 Mamba-Transformer 推理的非对称虚拟内存分页
机器学习
2026-05-22 v1 分布式、并行与集群计算
性能
摘要
混合语言模型如 Jamba 将注意力层与状态空间模型 (SSM) 混合，创建两种内存缓存类型：键值 (KV) 缓存随序列长度线性增长，而 SSM 状态在每层保持固定。当前推理引擎处理不当。统一的内存池将 SSM 状态填充至注意力页面大小，浪费最多 7.3 倍的容量。静态双池无法适应请求之间提示分布的变化。我们提出了非对称虚拟内存分页 (AVMP)。分配器将两种缓存类型分离到物理上不同的池后端，统一虚拟地址空间之下，当一种池耗尽时在池之间迁移容量。迁移仅在分配失败时触发，保持行为确定性。我们在 RTX 3060 12GB 上评估了 AVMP，针对 270 个合成单元外加 60 个 ShareGPT 轨迹重放。内存不足事件降低 7.6%，请求吞吐量在合成工作负载下提升 1.83 倍至 13.3 倍，在 ShareGPT 上提升 2.36 倍。所有收益在成对抽样 95% 置信区间内有效。阶段-时间分解揭示两种机制：在容量受压的工作负载上更短的 OOM 恢复时间，在 KV 重负载工作负载上更快的分配调用。实现为纯 Python；Triton 集成为未来工作。
引用
@article{arxiv.2605.22416,
title = {Asymmetric Virtual Memory Paging for Hybrid Mamba-Transformer Inference},
author = {An Xuan Nguyen},
journal= {arXiv preprint arXiv:2605.22416},
year = {2026}
}
备注
11 pages, 8 figures, 6 tables. Code and reproducibility artifacts at https://github.com/codepawl/cachepawl