[论文解读] Mosaic: Unlocking Long-Context Inference for Diffusion LLMs via Global Memory Planning and Dynamic Peak Taming

Liang Zheng, Bowen Shi+6 more|arXiv (Cornell University)|Jan 10, 2026

Parallel Computing and Optimization Techniques被引用 0

一句话总结

Mosaic 是一个用于扩散基础的大语言模型的内存高效推断系统，通过使用仅掩码的 logits、全局内存规划和惰性分块，降低内存峰值、极大扩展上下文长度，并在不损害准确性或速度的情况下改善延迟

ABSTRACT

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$ imes$ reduction in the memory peak-to-average ratio and increases the maximum inference sequence length supportable on identical hardware by 15.89-32.98$ imes$. This scalability is achieved without compromising accuracy and speed, and in fact reducing latency by 4.12%-23.26%.

研究动机与目标

• 识别在长上下文扩散型大模型（dLLMs）中内存成为瓶颈的原因，以及它与自回归大模型的差异。
• 设计一个面向 dLLMs 的内存高效推断系统，以缓解动态内存峰值和碎片化问题。
• 提出仅对掩码 token 计算 logits 的技术，以及在计算图中全局管理内存的方法。
• 在多种 dLLMs 上评估 Mosaic 对内存使用、最大可支持上下文长度、延迟和准确性的影响。

提出的方法

• 仅掩码 logits 的内核，通过 gather-GEMM 融合内核只对掩码 token 计算 logits。
• 图 registrar 用于定义带符号维度的参数化计算图，以实现全局可见性。
• 带在线瓶颈驱动搜索的惰性分块优化器，能自适应对内存密集算子进行分块。
• 具备单一全局复用计划和基于 VMM 的分配器的全局内存管理器，以消除碎片化。
• 离线图构造加上线运行时内存规划，以实现最小且足够的内存配置。
• 在有代表性的 dLLMs 上进行评估，以衡量内存、延迟和上下文长度的提升。

实验结果

研究问题

• RQ1 自回归大语言模型与扩散型大语言模型在长上下文场景下，内存瓶颈有何不同？
• RQ2 掩码仅 logits 和全局内存规划是否能在不损害延迟或准确性的前提下，减少 dLLMs 的内存峰值和碎片化？
• RQ3 动态内存峰值对最大可支持上下文长度有何影响，如何通过自适应分块来解决？
• RQ4 在相同硬件条件下，使用 Mosaic 能将上下文长度扩展到多少，及其对延迟的影响？

主要发现

• 平均内存峰值到平均值比降低 2.71×。
• 在相同硬件上，最大推断序列长度提升为 15.89–32.98×。
• 相较于基线，延迟平均降低 4.12%–23.26%。
• 在三种主流 dLLMs（LLaDA-8B、Dream-7B、LLaDA-MoE）上可以支持超出原生训练限制的上下文。
• 掩码仅 logits 和全局内存管理相对于先前的方法显著降低了内存膨胀和碎片化。