MOM让长文本模型在单卡上运行更长上下文，内存占用减半。
MOM: Memory-Efficient Offloaded Mini-Sequence Inference for Long Context Language Models
- 将关键层拆成小序列，结合缓存卸载降低内存压力。
- 在A100上使上下文长度从15.5万扩展至45.5万token，内存降超50%。
- 适合部署长文本生成的开发者，尤其关注推理效率与显存优化者。
长上下文语言模型表现优异，但推理时高显存需求制约其部署。本文提出内存高效卸载小序列推理（MOM），将关键层划分为更小的“小序列”，并无缝集成到KV缓存卸载中。在Llama、Qwen和Mistral等模型上的实验表明，MOM平均降低峰值显存使用量超过50%。在Meta-Llama-3.2-8B上，单张A100 80GB GPU下，最大上下文长度由155k扩展至455k token，输出完全一致且精度无损。MOM通过极低计算开销和高效的最后一层处理，保持高吞吐。相比传统分块预填充方法，其上下文扩展能力提升35%。更重要的是，该方法大幅减少预填充阶段内存消耗，彻底消除其长期作为主要瓶颈的问题。这一突破促使研究重点从预填充优化转向解码阶段残差KV缓存效率提升。
原文摘要 · Abstract (English)
Long-context language models exhibit impressive performance but remain challenging to deploy due to high GPU memory demands during inference. We propose Memory-efficient Offloaded Mini-sequence Inference (MOM), a method that partitions critical layers into smaller "mini-sequences" and integrates seamlessly with KV cache offloading. Experiments on various Llama, Qwen, and Mistral models demonstrate that MOM reduces peak memory usage by over 50\% on average. On Meta-Llama-3.2-8B, MOM extends the maximum context length from 155k to 455k tokens on a single A100 80GB GPU, while keeping outputs identical and not compromising accuracy. MOM also maintains highly competitive throughput due to minimal computational overhead and efficient last-layer processing. Compared to traditional chunked prefill methods, MOM achieves a 35\% greater context length extension. More importantly, our method drastically reduces prefill memory consumption, eliminating it as the longstanding dominant memory bottleneck during inference. This breakthrough fundamentally changes research priorities, redirecting future efforts from prefill-stage optimizations to improving decode-stage residual KV cache efficiency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。