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.


翻译:混合语言模型(如Jamba)将注意力层与状态空间模型(SSM)相结合,产生两种具有相反特性的内存缓存类型:键值(KV)缓存随序列长度线性增长,而SSM状态在每层中保持固定大小。当前推理引擎对此处理不佳。统一内存池将SSM状态填充至注意力页面大小,造成高达7.3倍的容量浪费。静态双池架构无法适应请求间提示分布的动态变化。我们提出非对称虚拟内存分页(AVMP)技术。该分配器将两种缓存类型分离至统一虚拟地址空间背后的物理独立池中,并在某个池耗尽时动态迁移容量。迁移仅于分配失败时触发,确保行为确定性。我们在RTX 3060 12GB上对270个合成单元及60个ShareGPT轨迹回放单元进行AVMP评估。内存不足事件减少7.6%,合成工作负载的请求吞吐量提升1.83倍至13.3倍,ShareGPT场景下提升2.36倍。所有增益在配对自助法95%置信区间内保持显著。阶段时间分解揭示两种不同机制:在容量压力工作负载下缩短OOM恢复时间,在KV密集型工作负载下加速分配调用。当前实现采用纯Python,基于Triton的集成工作留待后续开展。

0
下载
关闭预览

相关内容

TransMLA:多头潜在注意力(MLA)即为所需
专知会员服务
23+阅读 · 2025年2月13日
《混合专家模型推理优化技术综述》
专知会员服务
46+阅读 · 2024年12月21日
非Transformer不可?最新《状态空间模型(SSM)》综述
专知会员服务
75+阅读 · 2024年4月16日
从Seq2seq到Attention模型到Self Attention(一)
量化投资与机器学习
76+阅读 · 2018年10月8日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
VIP会员
最新内容
学习数据的几何:形状空间分析数学综述
专知会员服务
7+阅读 · 6月17日
定向能反无人机系统最新发展动态
专知会员服务
8+阅读 · 6月17日
从燃煤战舰到算法战争:水面指挥的永恒要求
专知会员服务
6+阅读 · 6月17日
相关基金
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
Top
微信扫码咨询专知VIP会员