Asymmetric Virtual Memory Paging for Hybrid Mamba-Transformer Inference

📄 arXiv: 2605.22416v1 📥 PDF

作者: An Xuan Nguyen

分类: cs.LG, cs.DC, cs.PF

发布日期: 2026-05-21

备注: 11 pages, 8 figures, 6 tables. Code and reproducibility artifacts at https://github.com/codepawl/cachepawl

💡 一句话要点

提出非对称虚拟内存分页AVMP，优化混合Mamba-Transformer模型推理的内存管理。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 混合Transformer 内存管理 推理优化 虚拟内存分页 状态空间模型

📋 核心要点

  1. 现有混合Transformer架构在推理时，KV缓存和SSM状态的内存需求差异大，传统内存管理方法效率低。
  2. 提出AVMP，将KV缓存和SSM状态分配到独立的物理内存池，并根据需求动态迁移容量，优化内存利用率。
  3. 实验表明，AVMP能显著减少OOM事件，提高请求吞吐量，尤其在容量受限和KV密集型场景下效果显著。

📝 摘要（中文）

混合语言模型（如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集成是未来的工作。

🔬 方法详解

问题定义：现有混合Transformer模型（如Jamba）在推理时，KV缓存随着序列长度线性增长，而SSM状态则相对固定。传统的内存管理方法，如统一内存池或静态双内存池，无法有效处理这种非对称的内存需求。统一内存池会造成大量内存浪费，而静态双内存池则无法适应不同请求之间的动态变化，容易导致内存溢出（OOM）等问题。

核心思路：AVMP的核心思路是针对KV缓存和SSM状态的内存需求差异，采用非对称的内存管理策略。具体来说，将两种缓存类型分配到物理上独立的内存池中，并利用虚拟内存技术，将这些物理池映射到统一的虚拟地址空间。这样，可以避免统一内存池的浪费，同时提供灵活的内存分配和迁移机制。

技术框架：AVMP的整体架构包含以下几个主要模块：1) 内存池管理器：负责创建和管理KV缓存池和SSM状态池。2) 虚拟地址空间管理器：负责将物理内存池映射到统一的虚拟地址空间。3) 内存分配器：根据请求类型（KV缓存或SSM状态）从相应的内存池中分配内存。4) 内存迁移器：当某个内存池耗尽时，负责将另一个内存池中的部分内存迁移到该池中。迁移器只在分配失败时触发，以保证行为的确定性。

关键创新：AVMP最重要的技术创新点在于其非对称的内存管理策略和动态内存迁移机制。与传统的统一内存池或静态双内存池相比，AVMP能够更好地适应混合Transformer模型中KV缓存和SSM状态的内存需求差异，从而提高内存利用率和推理性能。

关键设计：AVMP的关键设计包括：1) 独立的KV缓存池和SSM状态池，针对性地管理不同类型的内存需求。2) 基于分配失败触发的内存迁移机制，避免频繁的内存迁移操作，保证性能稳定。3) 统一的虚拟地址空间，简化内存管理和访问。论文中未提及具体的参数设置、损失函数或网络结构等细节，因为AVMP主要关注内存管理层面，与模型本身的结构关系不大。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在RTX 3060 12GB上的实验结果表明，AVMP能够显著减少OOM事件（降低7.6%），并提高请求吞吐量。在合成工作负载下，吞吐量提升了1.83倍至13.3倍，在ShareGPT数据集上提升了2.36倍。这些提升在95%置信区间下具有统计显著性。阶段时间分解显示，AVMP通过缩短OOM恢复时间和加速内存分配调用来提高性能。

🎯 应用场景

AVMP可应用于各种混合Transformer架构的推理加速，尤其是在资源受限的边缘设备或云服务器上。通过优化内存管理，AVMP能够提高模型的推理吞吐量，降低延迟，并减少OOM事件的发生，从而提升用户体验和系统稳定性。该技术还有潜力应用于其他具有类似非对称内存需求的AI模型。

📄 摘要（原文）

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.