arXiv:2601.06562cs.LG2026-01被引 3

解决扩散模型长文本生成的内存瓶颈,提升推理效率与上下文长度。

Mosaic: Unlocking Long-Context Inference for Diffusion LLMs via Global Memory Planning and Dynamic Peak Taming

  • 采用全局动态内存管理,通过掩码优化和在线搜索自适应缓解内存峰值。
  • 内存峰值与平均值比降低2.71倍,支持序列长度提升15.89至32.98倍。
  • 适合需要长上下文生成且对内存敏感的场景,如复杂推理与文档理解。

基于扩散的大型语言模型(dLLMs)通过并行去噪实现全局规划与迭代优化,特别适用于长上下文生成。然而,其部署面临严重的内存容量障碍,源于系统效率低下。现有推理系统不适用于此范式:与受累积键值缓存限制的自回归模型不同,dLLMs受限于每步重新计算的瞬时激活。此外,通用内存复用机制缺乏全局可见性,无法适应dLLMs在输出概率和前馈网络间动态切换的内存高峰。为此,我们提出Mosaic,一种从局部静态管理转向全局动态管理的高效推理系统。Mosaic集成仅掩码的输出层内核以消除冗余,基于在线启发式搜索的懒惰分块优化器以自适应缓解动态峰值,并通过虚拟寻址实现全局内存管理以解决碎片化问题。大量实验表明,Mosaic使内存峰值与平均值之比平均降低2.71倍,相同硬件上支持的最大推理序列长度提升15.89–32.98倍。该可扩展性在不牺牲准确率和速度的前提下实现,甚至将延迟降低4.12%–23.26%。

原文摘要 · Abstract (English)

Diffusion-based large language models (dLLMs) have emerged as a promising paradigm, utilizing simultaneous denoising to enable global planning and iterative refinement. While these capabilities are particularly advantageous for long-context generation, deploying such models faces a prohibitive memory capacity barrier stemming from severe system inefficiencies. We identify that existing inference systems are ill-suited for this paradigm: unlike autoregressive models constrained by the cumulative KV-cache, dLLMs are bottlenecked by transient activations recomputed at every step. Furthermore, general-purpose memory reuse mechanisms lack the global visibility to adapt to dLLMs' dynamic memory peaks, which toggle between logits and FFNs. To address these mismatches, we propose Mosaic, a memory-efficient inference system that shifts from local, static management to a global, dynamic paradigm. Mosaic integrates a mask-only logits kernel to eliminate redundancy, a lazy chunking optimizer driven by an online heuristic search to adaptively mitigate dynamic peaks, and a global memory manager to resolve fragmentation via virtual addressing. Extensive evaluations demonstrate that Mosaic achieves an average 2.71$\times$ reduction in the memory peak-to-average ratio and increases the maximum inference sequence length supportable on identical hardware by 15.89-32.98$\times$. This scalability is achieved without compromising accuracy and speed, and in fact reducing latency by 4.12%-23.26%.

扩散模型长上下文内存优化推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。