Source-linked AI summary
Fixed State, Long Reach: What a Constant-Size Cache Buys Block Diffusion at Scale
Vaibhav Singh, Pierre-André Noël, Torsten Scholak, Eugene Belilovsky, Oleksiy Ostapenko
TL;DR
Bidirectional denoising prevents masked diffusion language models from using native KV caching, and existing block caches can grow with context or approximate the model’s computation. The paper trains attention, Mamba, and hybrid denoisers with an exact cached-decoding objective; at scale, Mamba maintains constant cache cost and Mamba/Hybrid retrieve beyond their training length.
Problem
Bidirectional attention in diffusion language models prevents native KV caching, while existing attention-based block caches grow with context and training-free retrofits can approximate the model’s computation.
Source: §1 · pp. 1–2Method
The authors pretrain 3B attention, Mamba, and hybrid block-diffusion denoisers under a block-causal objective matched to a unified cached-decoding interface.
Source: §1 · p. 2, §3.1 · p. 4Results
Takeaways & Limitations
The results extend the case for linear-state diffusion language models from training to inference and from throughput to long-context retrieval.
Source: §1 · p. 2Models were trained at 1024-token context, so long-context evaluations test extrapolation rather than trained long-context ability; objective and training-length effects remain disentangled future work.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Diffusion language models decode tokens in parallel, but their bidirectional denoiser rules out the naive key--value (KV) cache behind fast autoregressive inference. Block diffusion restores caching by decoding block-by-block, and the block caches deployed on it so far are tied to attention: O(L)in memory and, if used as training-free retrofits, only an approximation of the model's computation. Both constraints can be overcome: sequence mixers that summarize finalized blocks into a reusable state support block caching, and the corresponding block-causal training objective makes the cache exact. We study this recipe at scale, pretraining three 3B block-diffusion denoisers (attention, mamba, and hybrid) on 300B tokens under one single-frontier objective and decoding all three through a single cached interface. Only the state-space cache is O(1) in sequence length: its memory and per-step latency stay constant at any context length, while an attention cache remains O(L). At 256k tokens (where attention has grown to 82GB and 29 ms/step), the Mamba cache delivers 4.3x lower latency, 11x less memory, and 2.6x higher single-stream throughput; and because that footprint is constant it scales with batch as well, reaching 14x the aggregate throughput, where attention cannot run beyond a single stream. The same linear-state bias lets the Mamba and hybrid backbones keep retrieving out to 8-16x their training length, whereas attention's retrieval collapses at 2x, at no measured quality cost.
1 Introduction
Bidirectional denoising prevents masked diffusion language models from using the usual autoregressive KV cache, while attention-based block caches remain costly and can approximate the model’s computation. The paper trains attention, Mamba, and hybrid block-diffusion denoisers with an objective matched to cached decoding, then studies their efficiency and long-context retrieval at scale.
1 Introduction
- Diffusion language models denoise masked tokens in parallel, but bidirectional attention prevents native KV caching and forces recomputation over the sequence at each step [39, 21].This limits the practical speed advantage expected from parallel decoding.Source: §1 · p. 1
- Block diffusion enables caching by generating blocks left-to-right with diffusion inside each block, but deployed attention caches remain O(L) and training-free retrofits approximate bidirectional computation [39] [21] [18] [22].A block-causal objective can instead make cached computation exact [2].Source: §1 · pp. 1–2
- The study exposes attention KV, Mamba recurrent state, and hybrid caches through one interface and pretrains all three 3B denoisers under the same block-causal objective.The study uses 300B tokens and a single-frontier objective; cached decoding reproduces the computation used in training.Source: §1 · p. 2
- At 256k tokens, Mamba uses ∼7.5 GB and 6.8 ms/step, reaches 14× aggregate throughput, and Mamba/Hybrid retrieve at 8–16× training length while Attn retrieval collapses at 2×.The models were trained at 1024 tokens; these longer-context evaluations therefore probe extrapolation.Source: §1 · p. 2
- The work extends prior state-space-model findings on fixed-size state and length extrapolation to diffusion denoisers, where block decoding reveals several tokens per pass [12] [9] [10] [28] [35] [16].Prior work had shown Mamba denoiser quality at training length [32] and small-scale hybrid block caching [7].Source: §1 · p. 2
2 Background and Related Work
The related work covers masked diffusion models, block diffusion, attention-cache retrofits, and state-space sequence mixers. This study builds on those lines by testing exact block-causal caching at 3B scale, including a pure Mamba denoiser, with long-context retrieval evaluation.
Masked diffusion language models
- Masked diffusion models train bidirectional networks to recover masked tokens, and LLaDA and Dream [41] scale this approach to billions of parameters; their full-sequence denoising does not natively admit KV caching.SEDD [20], MDLM [29], and refinements [31] [30] are cited as earlier masked-diffusion approaches.Source: §2 · p. 3
Block diffusion
Caching for diffusion LMs
State-space models and hybrids
- Prior exact training-time block-causal caches include BD3LM [2] and concurrent Mamba–attention hybrids [7], while earlier hybrid diffusion models were trained at 87M and 350M parameters [7].The cited exactness comes from training on the same clean-prefix computation used at decoding.Source: §2 · p. 4
Study contribution
- This study adds 3B-parameter, 300B-token comparisons across attention, hybrid, and pure Mamba models, measuring cached efficiency and long-context retrieval alongside quality.Its contribution is evidence about cache behavior at scale, rather than a new caching mechanism.Source: §2 · p. 4
3 Method
The method trains block-diffusion denoisers to use a clean prefix and masked future, then caches each backbone’s representation of finalized blocks. Decoding reuses that cache across denoising steps, with costs that depend on the backbone.
3.1 Block-diffusion preliminaries
- The single-frontier objective masks part of one sampled block, keeps earlier blocks clean, masks later blocks fully, and computes loss only on masked frontier tokens.Attention is block-causal; Mamba’s forward scan covers the prefix while its reverse scan is restricted to the frontier, preventing information leakage from masked future blocks [7].Source: §3.1 · pp. 4–5
3.2 A unified block cache
- The unified interface takes the current block and a per-layer cache, then returns logits and an updated cache; attention stores growing KV tensors while Mamba carries fixed-size convolution and SSM states.Mamba reruns its backward recurrence within the current block, and the hybrid combines both cache types; its memory lies between the pure backbones.Source: §3.2 · p. 5
3.3 Inference via cached generation
- Each block reuses a cloned entry cache across S denoising steps, reveals the highest-confidence masked positions, then folds the finalized block into the cache for the next block.The number of forward passes is O((L/G)S); each SSM pass touches only the G current tokens.Source: §3.3 · p. 5
3.4 Complexity and asymptotic throughput
- Per-step attention cost is O(GL) against the growing prefix, whereas Mamba costs O(G^2) independently of L; the hybrid retains attention scaling with a smaller constant.Only the SSM cache is O(1) in L; five of the hybrid’s 28 layers use attention.Source: §3.4 · p. 5, Table 1 · p. 6
4 Experimental Setup
The experiments compare three 3B denoisers trained on 300B tokens under a shared single-frontier objective, then evaluate decoding efficiency, long-context performance, downstream tasks, and generation quality.
Experimental setup
- Three 3B models differing only in sequence mixer were trained on 300B tokens at 1024-token context and evaluated for cached efficiency, retrieval, LongBench, downstream tasks, and generation quality.Efficiency was measured on one H100 across lengths of 64–256k and batches 1–8; retrieval and LongBench extended to 16k.Source: §4 · p. 6
5 Results
The state-space cache keeps decoding cost constant as context grows, while Mamba and Hybrid retain long-context retrieval that collapses for attention. These gains come with near-parity quality and reflect the mixer architecture rather than parameter count.
5.1 Efficiency: the state-space cache is constant in length
- At 256k, Mamba delivers 4.3× lower latency (6.8 vs. 29 ms), 11× less memory (7.5 vs. 82 GB), and 2.6× higher throughput (297 vs. 114 tok/s).At the same length, Hybrid beats Attn by 2.7× on latency, 3.8× on memory, and 2.0× on throughput.Source: §5.1 · pp. 6–7
5.2 Long-context generalization: linear state extrapolates, attention collapses
- At 2× training length, Attn retrieves 12% of needles versus Mamba’s 76%; by 16×, Attn is at 0% while Mamba recovers 22%.Within the 1024-token training length, all three models retrieve near-perfectly.Source: Table 2 · p. 7, §5.2 · p. 7
5.1 Efficiency: the state-space cache is constant in length
- At 256k and batch 8, Mamba sustains 1593 tok/s in about 8 GB, while Attn is out of memory at batch 2.The resulting aggregate-throughput advantage is 14×; below about 16k tokens the backbones are within a few percent.Source: §5.1 · p. 7
5.2 Long-context generalization: linear state extrapolates, attention collapses
- On LongBench, Mamba stays at 10.2–11.2 from 2k to 16k, while Attn drops to 3.6 at 16k, where Mamba nearly triples Attn.The task scores are macroaverages over 16 English and code tasks, and all lengths from 2k onward exceed the 1024-token training length.Source: §5.2 · pp. 7–8, Table 3 · p. 8
5.3 Quality parity: efficiency at no cost
- At S=4→32, Gen-PPL falls from 15.3→6.4 for Attn, 17.1→6.5 for Hybrid, and 18.5→6.5 for Mamba, at 5–6× lower throughput.MAUVE broadly improves but is noisy for these small base models, making Gen-PPL the cleaner quality signal.Source: Table 5 · p. 8, §5.3 · p. 8
- Validation NLL differs by about 0.03 nats; Attn and Hybrid score 0.434/0.432 on the downstream suite, while Mamba trails by about one point.The reported downstream score is the average across eight common-sense and reasoning tasks.Source: §5.3 · p. 8, Table 4 · p. 8
5.4 The gains are architectural, not parameter count
- At 64k, Mamba uses 0.26× Attn’s FLOPs despite 13% more parameters; its per-token FLOPs stay flat as context grows.The extra Mamba and Hybrid parameters are in O(d^2) projections independent of context length, while attention’s O(Ld) term grows with length.Source: Table 6 · p. 9, §5.4 · p. 9
6 Conclusion
The study presents exact cached block-diffusion decoding with a constant-memory state-space cache and reports long-context retrieval by Mamba and Hybrid. Its long-context tests extrapolate from 1024-token training, and the effects of objective and training length remain unresolved.
6 Conclusion
- At 256k tokens, Mamba uses 7.5 GB and 6.8 ms/step versus Attn’s 82 GB and 29 ms/step, while batch scaling reaches 14× aggregate throughput.Mamba and Hybrid retrieve at 8–16× their 1024-token training length while Attn retrieval collapses at 2×; the hybrid matches Attn on likelihood and downstream accuracy.Source: §6 · p. 9
- Long-context evaluations extrapolate from 1024-token training, and disentangling training length from the single-frontier objective is left for future work.The authors report modest absolute accuracies far past training length and note that longer-context training is compatible with their cache.Source: §6 · p. 9
A Training and model hyperparameters
The three denoisers share all training and architecture settings except for their sequence mixer.
A Training and model hyperparameters
- The three denoisers are identical except for the sequence mixer; every other training and architecture setting is shared.The configuration is listed in Table 7, and the mixers in Table 1.Source: A Training and model hyperparameters · p. 13
B Per-depth needle-in-a-haystack
Needle retrieval exhibits recency effects across backbones, but only Mamba and Hybrid retain retrieval of older needles beyond training length.
B Per-depth needle-in-a-haystack
- At 4× training length, only the linear-state backbones retain any retrieval, while all models show recency effects.Table 8 reports retrieval accuracy by needle depth at 2× and 4× training length; only Mamba and Hybrid keep retrieving older needles past training length, while Attn is near-zero except at the last depth.Source: B Per-depth needle-in-a-haystack · p. 13, Table 8 · p. 14
C Full batch-size efficiency sweep
Across batch sizes and context lengths, Mamba’s cache maintains a roughly constant footprint while attention runs out of memory at progressively shorter lengths as batch grows.
C Full batch-size efficiency sweep
- At 256k, Mamba at batch 8 delivers 1593 tok/s while Attn cannot run even batch 2; Attn’s best feasible aggregate is batch 1 at 114 tok/s, a ∼14× gap.Mamba throughput stays at 1593 tok/s at batch 8 across lengths to 256k, with roughly 8 GB memory; attention’s last feasible lengths are 256k, 128k, 64k, and 32k at batches 1, 2, 4, and 8.Source: C Full batch-size efficiency sweep · p. 14, Figure 4 · p. 15