用于混合 Mamba-Transformer 推理的非对称虚拟内存分页

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 集成是未来的工作。