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

arXiv: https://arxiv.org/abs/2601.06562 · HTML: https://mlsanity.com/paper/2601.06562/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 32 references that are in the map.
- First posted on arXiv 10 Jan 2026.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Dream 7B: Diffusion Large Language Models](https://mlsanity.com/paper/2508.15487/) (2025)
- [Large Language Models: A Survey](https://mlsanity.com/paper/2402.06196/) (2024)
- [DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models](https://mlsanity.com/paper/2210.08933/) (2022)
- [Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding](https://mlsanity.com/paper/2505.22618/) (2025)
- [SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control](https://mlsanity.com/paper/2210.17432/) (2022)
- [dLLM-Cache: Accelerating Diffusion Large Language Models with Adaptive Caching](https://mlsanity.com/paper/2506.06295/) (2025)
- [Sparse GPU Kernels for Deep Learning](https://mlsanity.com/paper/2006.10901/) (2020)
- [A Survey on Large Language Model Acceleration based on KV Cache Management](https://mlsanity.com/paper/2412.19442/) (2024)
- [vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention](https://mlsanity.com/paper/2405.04437/) (2024)
- [Dream-Coder 7B: An Open Diffusion Language Model for Code](https://mlsanity.com/paper/2509.01142/) (2025)
- [FlashDLM: Accelerating Diffusion Language Model Inference via Efficient KV Caching and Guided Diffusion](https://mlsanity.com/paper/2505.21467/) (2025)
- [Discrete Diffusion in Large Language and Multimodal Models: A Survey](https://mlsanity.com/paper/2506.13759/) (2025)
- [LLaDA-MoE: A Sparse MoE Diffusion Language Model](https://mlsanity.com/paper/2509.24389/) (2025)
- [LongLLaDA: Unlocking Long Context Capabilities in Diffusion LLMs](https://mlsanity.com/paper/2506.14429/) (2025)
- [dInfer: An Efficient Inference Framework for Diffusion Language Models](https://mlsanity.com/paper/2510.08666/) (2025)
- [UltraLLaDA: Scaling the Context Length to 128K for Diffusion Large Language Models](https://mlsanity.com/paper/2510.10481/) (2025)
- [Efficient Memory Management for Deep Neural Net Inference](https://mlsanity.com/paper/2001.03288/) (2020)
- [PAT: Accelerating LLM Decoding via Prefix-Aware Attention with Resource Efficient Multi-Tile Kernel](https://mlsanity.com/paper/2511.22333/) (2025)

## Built on by (most cited first, 2 of 2)

- [DMax: Aggressive Parallel Decoding for dLLMs](https://mlsanity.com/paper/2604.08302/) (2026)
- [TACG: Trajectory-Aware Commit Gating for Diffusion Language Model Decoding](https://mlsanity.com/paper/2607.03236/) (2026)

## Abstract (arXiv)

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