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

TL;DR

This work tackles the memory bottlenecks of long-context diffusion-based LLMs by reframing memory management from a KV-cache mindset to a global, activation-centric paradigm. It introduces Mosaic, a four-component system comprising a mask-only logits kernel, a graph registrar, a lazy chunking optimizer, and a global memory manager with virtual memory mapping, enabling offline graph templating and online bottleneck-driven memory planning. The approach yields substantial improvements in memory efficiency and context length, achieving an average memory peak-to-average reduction of 2.71$\times$ and up to 15.89–32.98$\times$ longer usable context, while reducing per-step latency by 4.12%–23.26% across multiple dLLMs. These gains are complemented by an ablation study showing the dominant contributions from global memory planning, mask-only computations, and adaptive chunking, collectively enabling scalable, long-context diffusion inference on commodity GPUs. Overall, Mosaic demonstrates that global memory planning and dynamic peak taming can unlock practical long-context diffusion LLMs without sacrificing speed or accuracy.

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$\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%.

Figures (18)

• Figure 1: dLLM inference pipeline.
• Figure 2: Memory breakdown of both LLMs.
• Figure 3: Dynamic memory bottleneck shift across varying mask ratios ($r_m$) at 128k context length. The memory peak toggles between Logits (at high $r_m$) and FFN (at low $r_m$).
• Figure 4: Reserved memory vs. theoretical peak.
• Figure 5: Latency w/ and w/o segment caching.