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

• URL: http://arxiv.org/abs/2601.06562v1
• Date: Sat, 10 Jan 2026 13:17:08 GMT
• Title: Mosaic: Unlocking Long-Context Inference for Diffusion LLMs via Global Memory Planning and Dynamic Peak Taming
• Authors: Liang Zheng, Bowen Shi, Yitao Hu, Jiawei Zhang, Ruofan Li, Sheng Chen, Wenxin Li, Keqiu Li,
• Abstract summary: Diffusion-based large language models (dLLMs) have emerged as a promising paradigm, utilizing simultaneous denoising to enable global planning and iterative refinement.<n>Existing inference systems are ill-suited for this paradigm due to severe system inefficiencies.<n>We propose Mosaic, a memory-efficient inference system that shifts from local, static management to a global, dynamic paradigm.
• Score: 34.16016695663811
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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%.

Related papers

• AllMem: A Memory-centric Recipe for Efficient Long-context Modeling [32.025154452526856]
  Large Language Models (LLMs) encounter significant performance bottlenecks in long-sequence tasks.<n>We introduce textscAllMem, a novel and efficient hybrid architecture that integrates Sliding Window Attention (SWA) with non-linear Test-Time Training (TTT) memory networks.
  arXiv  Detail & Related papers  (2026-02-14T09:04:28Z)
• MemFly: On-the-Fly Memory Optimization via Information Bottleneck [35.420309099411874]
  Long-term memory enables large language model agents to tackle complex tasks through historical interactions.<n>Existing frameworks encounter a dilemma between compressing redundant information efficiently and maintaining precise retrieval for downstream tasks.<n>MemFly is a framework grounded in information bottleneck principles that facilitates on-the-fly memory evolution for LLMs.<n>MemFly substantially outperforms state-of-the-art baselines in memory coherence, response fidelity, and accuracy.
  arXiv  Detail & Related papers  (2026-02-08T09:37:25Z)
• Causal Autoregressive Diffusion Language Model [70.7353007255797]
  CARD reformulates the diffusion process within a strictly causal attention mask, enabling dense, per-token supervision in a single forward pass.<n>Our results demonstrate that CARD achieves ARM-level data efficiency while unlocking the latency benefits of parallel generation.
  arXiv  Detail & Related papers  (2026-01-29T17:38:29Z)
• SimpleMem: Efficient Lifelong Memory for LLM Agents [73.74399447715052]
  We introduce SimpleMem, an efficient memory framework based on semantic lossless compression.<n>We propose a three-stage pipeline designed to maximize information density and token utilization.<n> Experiments on benchmark datasets show that our method consistently outperforms baseline approaches in accuracy, retrieval efficiency, and inference cost.
  arXiv  Detail & Related papers  (2026-01-05T21:02:49Z)
• RevFFN: Memory-Efficient Full-Parameter Fine-Tuning of Mixture-of-Experts LLMs with Reversible Blocks [12.966077380225856]
  RevFFN is a memory efficient fine tuning paradigm for mixture of experts (MoE) LLMs.<n>RevFFN employs carefully designed reversible Transformer blocks that allow reconstruction of layer input activations from outputs during backpropagation.
  arXiv  Detail & Related papers  (2025-12-24T03:56:58Z)
• Mixture-of-Channels: Exploiting Sparse FFNs for Efficient LLMs Pre-Training and Inference [16.71963410333802]
  Large language models (LLMs) have demonstrated remarkable success across diverse artificial intelligence tasks.<n>MoC substantially reduces activation memory during pre-training.<n>MoC delivers significant memory savings and throughput gains while maintaining competitive model performance.
  arXiv  Detail & Related papers  (2025-11-12T13:30:57Z)
• Fast-dLLM v2: Efficient Block-Diffusion LLM [64.38006546510337]
  Fast-dLLM v2 is a block diffusion language model that adapts pretrained AR models into dLLMs for parallel text generation.<n>This represents a 500x reduction in training data compared to full-attention diffusion LLMs such as Dream (580B tokens)
  arXiv  Detail & Related papers  (2025-09-30T14:40:18Z)
• Sparse-dLLM: Accelerating Diffusion LLMs with Dynamic Cache Eviction [72.27673320976933]
  Diffusion Large Language Models (dLLMs) enable breakthroughs in reasoning and parallel decoding.<n>Current caching techniques accelerate decoding by storing full-layer states, yet impose substantial memory usage.<n>We propose Sparse-dLLM, the first training-free framework integrating dynamic cache eviction with sparse attention.
  arXiv  Detail & Related papers  (2025-08-04T16:14:03Z)
• mGRADE: Minimal Recurrent Gating Meets Delay Convolutions for Lightweight Sequence Modeling [0.5236468296934584]
  mGRADE is a hybrid-memory system that integrates a temporal 1D-convolution with learnable spacings followed by a minimal gated recurrent unit.<n>We demonstrate that mGRADE effectively separates and preserves multi-scale temporal features.<n>This highlights mGRADE's promise as an efficient solution for memory-constrained multi-scale temporal processing at the edge.
  arXiv  Detail & Related papers  (2025-07-02T15:44:35Z)
• Flexiffusion: Training-Free Segment-Wise Neural Architecture Search for Efficient Diffusion Models [50.260693393896716]
  Diffusion models (DMs) are powerful generative models capable of producing high-fidelity images but constrained by high computational costs.<n>We propose Flexiffusion, a training-free NAS framework that jointly optimize generation schedules and model architectures without modifying pre-trained parameters.<n>Our work pioneers a resource-efficient paradigm for searching high-speed DMs without sacrificing quality.
  arXiv  Detail & Related papers  (2025-06-03T06:02:50Z)
• dLLM-Cache: Accelerating Diffusion Large Language Models with Adaptive Caching [27.114862565164145]
  diffusion-based Large Language Models generate text by iteratively denoising masked segments.<n>dLLMs suffer from high inference latency.<n>Traditional ARM acceleration techniques are incompatible with dLLMs due to their bidirectional attention mechanism.<n>We propose dLLM-Cache, a training-free adaptive caching framework.
  arXiv  Detail & Related papers  (2025-05-17T15:50:46Z)
• Any Image Restoration via Efficient Spatial-Frequency Degradation Adaptation [158.37640586809187]
  Restoring any degraded image efficiently via just one model has become increasingly significant.<n>Our approach, termed AnyIR, takes a unified path that leverages inherent similarity across various degradations.<n>To fuse the degradation awareness and the contextualized attention, a spatial-frequency parallel fusion strategy is proposed.
  arXiv  Detail & Related papers  (2025-04-19T09:54:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.