Fixed State, Long Reach: What a Constant-Size Cache Buys Block Diffusion at Scale

arXiv cs.LG Papers

Summary

This paper introduces a constant-size state cache for block diffusion models, demonstrating significant reductions in memory and latency compared to attention-based approaches, enabling efficient long-context generation without quality loss.

arXiv:2609.11998v1 Announce Type: new Abstract: 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.
Original Article
View Cached Full Text

Cached at: 09/14/26, 08:31 AM

# Fixed State, Long Reach: What a Constant-Size Cache Buys Block Diffusion at Scale
Source: [https://arxiv.org/html/2609.11998](https://arxiv.org/html/2609.11998)
Vaibhav Singh Pierre\-André Noël Torsten ScholakEugene Belilovsky Oleksiy OstapenkoMila Concordia University ServiceNow Research††thanks:Correspondence to: vaibhav\.singh@mila\.quebec

###### Abstract

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*:𝒪⁡\(L\)\\mathcal\{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\) on300300B tokens under one single\-frontier objective and decoding all three through a single cached interface\. Only the state\-space cache is𝒪⁡\(1\)\\mathcal\{O\}\(1\)in sequence length: its memory and per\-step latency stay*constant*at any context length, while an attention cache remains𝒪⁡\(L\)\\mathcal\{O\}\(L\)\. At 256k tokens \(where attention has grown to8282GB and2929ms/step\), the Mamba cache delivers4\.3×4\.3\\timeslower latency,𝟏𝟏×11\\timesless memory, and2\.6×2\.6\\timeshigher single\-stream throughput; and because that footprint is constant it scales with batch as well, reaching𝟏𝟒×14\\timesthe 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–𝟏𝟔×16\\timestheir training length, whereas attention’s retrieval collapses at𝟐×2\\times, at no measured quality cost\.

## 1Introduction

Diffusion Language Models \(DLMs\) have emerged as a compelling alternative to the autoregressive \(AR\) paradigm that underpins most large language models\[[5](https://arxiv.org/html/2609.11998#bib.bib5),[1](https://arxiv.org/html/2609.11998#bib.bib1),[8](https://arxiv.org/html/2609.11998#bib.bib8),[34](https://arxiv.org/html/2609.11998#bib.bib34)\]\. Masked DLMs such as LLaDA\[[24](https://arxiv.org/html/2609.11998#bib.bib24)\]and Dream\[[41](https://arxiv.org/html/2609.11998#bib.bib41)\]corrupt a sequence by masking tokens and learn to denoise it, decoding many positions in parallel rather than one token at a time\. This promises lower latency and richer per\-example supervision\[[23](https://arxiv.org/html/2609.11998#bib.bib23)\]\. In practice, however, DLM inference has often been*slower*than AR: the very bidirectional attention that makes denoising possible rules out the naive key–value \(KV\) cache that AR decoders rely on, so each denoising step recomputes attention over the entire sequence\[[39](https://arxiv.org/html/2609.11998#bib.bib39),[21](https://arxiv.org/html/2609.11998#bib.bib21)\]\.

*Block diffusion*\[[2](https://arxiv.org/html/2609.11998#bib.bib2)\]resolves the tension by interpolating between the two regimes: the sequence is partitioned into blocks that are generated autoregressively, while diffusion denoising runs*within*each block\. Because finalized blocks form a clean prefix, their representations can be cached and reused, restoring KV\-caching and enabling arbitrary\-length generation\. This has made block diffusion the foundation of recent fast\-decoding systems\[[39](https://arxiv.org/html/2609.11998#bib.bib39),[38](https://arxiv.org/html/2609.11998#bib.bib38),[37](https://arxiv.org/html/2609.11998#bib.bib37)\]\.

Two limitations remain\. First, the block caches deployed so far are built for an attention backbone\[[39](https://arxiv.org/html/2609.11998#bib.bib39),[21](https://arxiv.org/html/2609.11998#bib.bib21),[18](https://arxiv.org/html/2609.11998#bib.bib18),[22](https://arxiv.org/html/2609.11998#bib.bib22)\]; the resulting KV\-cache still grows linearly in context length and the underlying attention still costs quadratic compute, so the long\-context regime stays expensive\. Second, many of these methods are*training\-free*adaptations onto full\-attention models, which makes their block cache an*approximation*of the true bidirectional computation\[[39](https://arxiv.org/html/2609.11998#bib.bib39)\]\. Neither limitation is intrinsic to block diffusion, which only requires that finalized blocks summarize into a reusable state: any sequence mixer with such a state qualifies, and training with the block\-causal objective makes the resulting cache exact rather than approximate\[[2](https://arxiv.org/html/2609.11998#bib.bib2)\]\.

State\-space models \(SSMs\) such as Mamba\[[12](https://arxiv.org/html/2609.11998#bib.bib12),[9](https://arxiv.org/html/2609.11998#bib.bib9)\]are exactly such mixers\. They process a sequence with a linear\-time recurrence whose hidden state is a fixed\-size summary of the past, independent of how long the past is\. We build on the bidirectional\-Mamba DLM denoiser of DiffuMamba\[[32](https://arxiv.org/html/2609.11998#bib.bib32)\]and use its recurrent state as the block cache: when block diffusion finalizes a block, the forward SSM state already encodes the entire prefix in𝒪⁡\(1\)\\mathcal\{O\}\(1\)memory\. Concurrent work\[[7](https://arxiv.org/html/2609.11998#bib.bib7)\]has recently shown, at small scale \(8787M–350350M parameters\), that a Mamba–attention block\-diffusion hybrid admits exactly such a cache once the reverse Mamba scan is confined to the active block\. Whether this holds when pretrained at scale, how a*pure*Mamba denoiser fares, and what a constant\-size state delivers on long\-context*retrieval*rather than throughput alone, has not been studied\. We therefore expose the SSM state together with the attention KV\-cache and a hybrid of the two through a single decoding interface, and we*pretrain*all three models at 3B parameters under the same block\-causal objective, so that cached decoding reproduces training exactly\.

Concretely, we pretrain and analyze three 3B\-parameter block\-diffusion models on300300B tokens that share data, tokenizer, schedule, and decoding budget and differ only in the denoiser:Attn\(full attention\),Mamba\(bidirectional Mamba\-2\), andHybrid\(attention interleaved every five Mamba layers\)\. Because the Mamba mixer carries larger projections while the MLP is held fixed,MambaandHybridtechnically have13%13\\%and11%11\\%more parameters thanAttn;[Section5\.4](https://arxiv.org/html/2609.11998#S5.SS4)shows that these extra parameters add a fixed𝒪⁡\(d2\)\\mathcal\{O\}\(d^\{2\}\)cost per token at every context length, whereas attention’s score computation costs𝒪⁡\(L​d\)\\mathcal\{O\}\(Ld\)per token and grows with context\. Per\-token FLOPs are comparable at the training length, and the long\-context gap is due to the mixer rather than the parameter budget\. Our contributions are:

- •A controlled 3B pretraining study of cacheable block\-diffusion denoisers\([Sections3](https://arxiv.org/html/2609.11998#S3)and[4](https://arxiv.org/html/2609.11998#S4)\): attention, bidirectional Mamba\-2, and hybrid backbones, each pretrained on300300B tokens under the same single\-frontier block\-causal objective\. As in BD3LM\[[2](https://arxiv.org/html/2609.11998#bib.bib2)\]for attention and, concurrently,[Chaturvedi et al\. \[7\]](https://arxiv.org/html/2609.11998#bib.bib7)for Mamba hybrids, the cache is exact because the models are trained with the objective used at decode time; unlike training\-free retrofits, cached inference*is*the function the model learned\. To our knowledge this is the largest such study, and the only one that includes a*pure*state\-space denoiser\.
- •Constant memory and long context from a single linear\-state backbone\([Sections5\.1](https://arxiv.org/html/2609.11998#S5.SS1)and[5\.2](https://arxiv.org/html/2609.11998#S5.SS2)\)\. Measured on the trained checkpoints, the SSM cache is𝒪⁡\(1\)\\mathcal\{O\}\(1\)in length:Mambaholds∼\\sim7\.5 GB and 6\.8 ms/step at*every*length out to 256k, against attention’s 82 GB, reaching𝟏𝟒×\\mathbf\{14\\times\}the aggregate throughput at 256k\. The same backbone also generalizes far past its 1024\-token training length:Attn’s retrieval collapses just2×2\\timesout on NIAH and LongBench, whereasMamba/Hybridextrapolate to𝟖8–𝟏𝟔×16\\times\.
- •Long\-context evaluation at scale\([Section5](https://arxiv.org/html/2609.11998#S5)\)\. With everything but the denoiser held fixed, we cover cached\-decoding efficiency \(latency, memory, and throughput\) out to 256k tokens, NIAH and LongBench length extrapolation, an eight\-task downstream suite, MAUVE and generative perplexity, and a per\-layer FLOPs analysis that attributes the gains to the*architecture*rather than to parameter count\.

None of these properties is new to state\-space models\. Fixed\-size state\[[12](https://arxiv.org/html/2609.11998#bib.bib12),[9](https://arxiv.org/html/2609.11998#bib.bib9)\]and length extrapolation\[[10](https://arxiv.org/html/2609.11998#bib.bib10),[28](https://arxiv.org/html/2609.11998#bib.bib28)\]are established properties of*autoregressive*SSMs, with hybrids restoring the recall that pure recurrence lacks\[[35](https://arxiv.org/html/2609.11998#bib.bib35),[16](https://arxiv.org/html/2609.11998#bib.bib16)\]\. Whether they carry over to diffusion denoisers, which decode several tokens per pass, is open\. Our results are the block\-diffusion counterpart of that AR literature, and the combination matters because block diffusion reveals several tokens per forward pass rather than one, and a linear\-state backbone lets it do so on a state that never grows\.[Singh et al\. \[32\]](https://arxiv.org/html/2609.11998#bib.bib32)showed Mamba denoisers match attention on quality at the training length, and[Chaturvedi et al\. \[7\]](https://arxiv.org/html/2609.11998#bib.bib7)that they can be block\-cached like attention\[[2](https://arxiv.org/html/2609.11998#bib.bib2)\]at small scale; we show that the same backbones, pretrained at 3B,*decode*long contexts with a constant\-memory cache and retrieve far past their training length, extending the case for linear\-state diffusion LMs from training to inference and from throughput to long\-context retrieval\.

## 2Background and Related Work

#### Masked diffusion language models\.

Discrete diffusion models define a forward process that progressively corrupts a token sequence and learn a parameterized reverse process that denoises\[[3](https://arxiv.org/html/2609.11998#bib.bib3),[6](https://arxiv.org/html/2609.11998#bib.bib6)\]\. The*masked*\(absorbing\-state\) instantiation has proven the most effective for language: SEDD\[[20](https://arxiv.org/html/2609.11998#bib.bib20)\], MDLM\[[29](https://arxiv.org/html/2609.11998#bib.bib29)\], and their refinements\[[31](https://arxiv.org/html/2609.11998#bib.bib31),[30](https://arxiv.org/html/2609.11998#bib.bib30)\]train a bidirectional network to recover masked tokens under a noise schedule, and LLaDA\[[24](https://arxiv.org/html/2609.11998#bib.bib24)\]and Dream\[[41](https://arxiv.org/html/2609.11998#bib.bib41)\]scale this recipe to billions of parameters\. All decode by iteratively unmasking positions in parallel; none admits a KV\-cache natively, because every denoising step attends over the full \(partially masked\) sequence in both directions\.

#### Block diffusion\.

BD3LM\[[2](https://arxiv.org/html/2609.11998#bib.bib2)\]interpolates between AR and diffusion: tokens are grouped into blocks generated left\-to\-right, with intra\-block diffusion\. Blocks already produced act as a clean prefix, so their keys and values can be cached between blocks, and sequences of arbitrary length can be generated\. Our attention model follows this scheme; we use the nameAttnrather than “BD3LM” to keep our trained model distinct from the cited training recipe\. Subsequent systems build fast decoders on this foundation\[[39](https://arxiv.org/html/2609.11998#bib.bib39),[38](https://arxiv.org/html/2609.11998#bib.bib38),[37](https://arxiv.org/html/2609.11998#bib.bib37)\]\.

\(a\)Mamba/Hybrid: a single recurrent statehhof fixed shapeℝH×P×N\\mathbb\{R\}^\{H\\times P\\times N\}, carried and updated in place \(ht=A¯t​ht−1\+B¯t​xth\_\{t\}=\\bar\{A\}\_\{t\}\\,h\_\{t\-1\}\+\\bar\{B\}\_\{t\}\\,x\_\{t\}\),𝒪⁡\(1\)\\mathcal\{O\}\(1\)in sequence length\.\(b\)Attn: the key–value memory, which appends one entry per block \(KVt−2→KVt\+1\\mathrm\{KV\}\_\{t\-2\}\\\!\\to\\\!\\mathrm\{KV\}\_\{t\+1\}\) and therefore grows as𝒪⁡\(L\)\\mathcal\{O\}\(L\)\.
Figure 1:The unified block cache\.Block diffusion decodes one block at a time, left to right: finalized blocks \(xt−2,xt−1x\_\{t\-2\},x\_\{t\-1\}\) are clean, the frontier blockxtx\_\{t\}is being denoised \(its first token is revealed, the rest are\[MASK\]\), and future blocks \(xt\+1x\_\{t\+1\}\) stay fully masked\. Once a block is finalized, every layer writes its state into a cache that the next block reads, so the clean prefix is never recomputed\. The cached object depends on the backbone\.
#### Caching for diffusion LMs\.

Because the bidirectional denoiser blocks the standard KV\-cache, a line of work restores caching for attention DLMs\. Fast\-dLLM\[[39](https://arxiv.org/html/2609.11998#bib.bib39)\]introduces a block\-wise*approximate*KV\-cache plus confidence\-aware parallel decoding, reporting large speedups on LLaDA/Dream by caching the fixed context and refreshing it at block boundaries; dKV\-Cache\[[21](https://arxiv.org/html/2609.11998#bib.bib21)\]proposes a delayed KV\-cache compatible with bidirectional attention; and dLLM\-Cache\[[18](https://arxiv.org/html/2609.11998#bib.bib18)\]and elastic\-cache variants\[[22](https://arxiv.org/html/2609.11998#bib.bib22)\]add adaptive eviction and reuse\. A more recent wave targets the*long\-context*regime specifically: Prefilling\-dLLM\[[40](https://arxiv.org/html/2609.11998#bib.bib40)\]caches chunked prefix KV and keeps the top\-KKrelevant chunks, and Focus\-dLLM\[[19](https://arxiv.org/html/2609.11998#bib.bib19)\]exploits attention sparsity for dynamic cache eviction, reporting99–28×28\\timesspeedups at88–3232k\. A separate line extends the*usable*context of attention DLMs by rescaling positional encodings: LongLLaDA\[[17](https://arxiv.org/html/2609.11998#bib.bib17)\]applies training\-free NTK\-RoPE \(and documents a “local perception” sliding\-window bias\), while UltraLLaDA\[[14](https://arxiv.org/html/2609.11998#bib.bib14)\]*trains*to a128128k window\. All of these caches are attention\-only and𝒪⁡\(L\)\\mathcal\{O\}\(L\)in memory, and, being largely training\-free retrofits, only*approximate*the full\-attention computation they replace\. The exception is training\-time block\-causal schemes, BD3LM\[[2](https://arxiv.org/html/2609.11998#bib.bib2)\]for attention and, concurrently with this work,[Chaturvedi et al\. \[7\]](https://arxiv.org/html/2609.11998#bib.bib7)for Mamba–attention hybrids, whose caches are exact because the model is trained on the same clean\-prefix computation it decodes with; our models belong to this family\.

#### State\-space models and hybrids\.

SSMs model sequences with a linear\-time selective recurrence: S4\[[13](https://arxiv.org/html/2609.11998#bib.bib13)\], Mamba\[[12](https://arxiv.org/html/2609.11998#bib.bib12)\], and Mamba\-2\[[9](https://arxiv.org/html/2609.11998#bib.bib9)\]match or beat attention on language while carrying a fixed\-size state, and related structured operators\[[27](https://arxiv.org/html/2609.11998#bib.bib27),[11](https://arxiv.org/html/2609.11998#bib.bib11)\]share the linear\-time property\. Hybrid AR models interleave attention with linear recurrence to combine global recall with cheap long\-range mixing\[[16](https://arxiv.org/html/2609.11998#bib.bib16),[10](https://arxiv.org/html/2609.11998#bib.bib10),[36](https://arxiv.org/html/2609.11998#bib.bib36)\]\. DiffuMamba\[[32](https://arxiv.org/html/2609.11998#bib.bib32)\]brought bidirectional Mamba and Mamba–attention hybrids to*diffusion*denoisers\.[Chaturvedi et al\. \[7\]](https://arxiv.org/html/2609.11998#bib.bib7)trained block cache diffusion models at a small scale with Mamba–attention hybrids of8787M and350350M parameters\.

What this study adds is scale and evidence: pretraining at 3B parameters and300300B tokens with all three backbones, including a*pure*Mamba denoiser absent from that work; a single\-frontier objective; latency, memory, and throughput measured on the trained checkpoints, single\-stream and batched; and, most importantly, long\-context*retrieval*and length extrapolation \(NIAH, LongBench\) alongside downstream and generation quality\. Our claim is therefore not the mechanism but its behavior at scale: constant\-memory block decoding survives pretraining at 3B and comes with a marked long\-context retrieval advantage over the attention baseline\.

## 3Method

Algorithm 1Cached block\-wise generation1:prefix, block size

GG, steps per block

SS, blocks

KK, mask id

2:

𝑐𝑎𝑐ℎ𝑒←∅\\mathit\{cache\}\\leftarrow\\varnothing; warm

𝑐𝑎𝑐ℎ𝑒\\mathit\{cache\}on the clean prompt blocks

3:for

i←ifirsti\\leftarrow i\_\{\\text\{first\}\}to

K−1K\-1do

4:

s←Clone​\(𝑐𝑎𝑐ℎ𝑒\)s\\leftarrow\\textsc\{Clone\}\(\\mathit\{cache\}\)⊳\\trianglerightentry state, reused across denoising steps

5:

m←m\\leftarrow\# masked tokens in block

ii
6:for

t←1t\\leftarrow 1to

SSdo

7:

ℓ←ForwardBlockCached​\(xi,s,posi\)\\ell\\leftarrow\\textsc\{ForwardBlockCached\}\(x\_\{i\},\\,s,\\,\\mathrm\{pos\}\_\{i\}\)
8:reveal the

⌈m/S⌉\\lceil m/S\\rceilhighest\-confidence masked positions of

xix\_\{i\}using

ℓ\\ell
9:ifblock

iifully revealedthen

10:break

11:endif

12:endfor

13:

\_,𝑐𝑎𝑐ℎ𝑒←ForwardBlockCached​\(xi,s,posi\)\\\_,\\ \\mathit\{cache\}\\leftarrow\\textsc\{ForwardBlockCached\}\(x\_\{i\},\\,s,\\,\\mathrm\{pos\}\_\{i\}\)⊳\\trianglerightfold finalized block into cache

14:endfor

15:return

xx

### 3\.1Block\-diffusion preliminaries

Letx=\(x1,…,xL\)x=\(x\_\{1\},\\dots,x\_\{L\}\)be a token sequence\. A masked diffusion model defines a forward process that, at timet∈\(0,1\]t\\in\(0,1\], independently replaces each token by a special\[MASK\]symbol with probability1−αt1\-\\alpha\_\{t\}, whereαt\\alpha\_\{t\}follows a monotone schedule\. The denoiserfθf\_\{\\theta\}is trained to predict the original tokens from the masked sequence, minimizing a re\-weighted cross\-entropy that forms a negative ELBO on the data log\-likelihood\[[29](https://arxiv.org/html/2609.11998#bib.bib29),[24](https://arxiv.org/html/2609.11998#bib.bib24)\]\.

*Block diffusion*factorizes the sequence intoK=L/GK=L/Gcontiguous blocks of sizeGGand is autoregressive across blocks: blockiiis generated conditioned on the clean blocks0,…,i−10,\\dots,i\{\-\}1, while the tokens*inside*blockiiare produced by diffusion\. We train with asingle\-frontierobjective \([Figure1](https://arxiv.org/html/2609.11998#S2.F1)\): for each example we sample one frontier blockii, keep blocks<i<iclean, mask a time\-dependent random subset of blockii, and mask blocks\>i\>i*entirely*; the loss is the re\-weighted cross\-entropy on the masked positions of blockiionly\. This exactly matches the information available during cached generation, where blockiisees a clean prefix and all\-masked future\. The frontier structure is enforced differently per backbone\.*Attention*layers use ablock\-causalmask: a query in blockiiattends to all keys in blocks≤i\\leq i\(bidirectionally within its own block\) and to none in blocks\>i\>i\.*Mamba*layers run a forward \(left\-to\-right\) recurrence over the whole prefix, and a backward \(right\-to\-left\) recurrence whose support is*restricted to the frontier block*, so that no information leaks backward from the masked future \(the block\-restricted reverse scan that[Chaturvedi et al\. \[7\]](https://arxiv.org/html/2609.11998#bib.bib7)term*partial bidirectionality*\)\. The hybrid applies both rules, layer by layer\.

### 3\.2A unified block cache

At inference the blocks are produced strictly left\-to\-right, so once blockiiis finalized its contribution to every future block is fixed\. We expose a single per\-layer interface that consumes only theGGtokens of the current block plus a per\-layer cache summarizing blocks0,…,i−10,\\dots,i\{\-\}1, and returns the block logits and an updated cache\. The cache type is backbone\-specific \([Table1](https://arxiv.org/html/2609.11998#S4.T1)\)\. In the attention layers, for theGGnew tokens, we compute queries, keys, and values, and then concatenate the new keys/values with the cached\(K,V\)\(K,V\)of all finalized blocks\. The attended context is the clean prefix plus the current block; the cache grows as𝒪⁡\(L\)\\mathcal\{O\}\(L\)and the per\-step cost is𝒪⁡\(G​L\)\\mathcal\{O\}\(GL\)\.

A Mamba\-2 layer maintains two states that together summarize the prefix: a short causal\-convolution state \(the lastk−1k\{\-\}1inputs, wherekkis the conv width\) and the selective\-SSM recurrent state of shape\(heads×head\_dim×dstate\)\(\\text\{heads\}\\times\\text\{head\\\_dim\}\\times d\_\{\\text\{state\}\}\)\. The forward recurrence consumes the cached states, emits the block output, and returns updated states via a chunked scan with initial states; both states are*independent ofLL*\. The backward \(bidirectional\) recurrence is re\-run from a zero state*within the current block only*, reproducing the frontier restriction used in training; it is never carried across blocks\. Thus the Mamba cache is𝒪⁡\(1\)\\mathcal\{O\}\(1\)in length and the per\-step cost is𝒪⁡\(G2\)\\mathcal\{O\}\(G^\{2\}\), independent of how much text precedes the block\. Lastly the hybrid denoiser interleaves the two layer types, so its cache is a heterogeneous list: attention layers store\(K,V\)\(K,V\)tensors and Mamba layers store\(conv,ssm\)\(\\text\{conv\},\\text\{ssm\}\)states\. The overall memory is dominated by the \(few\) attention layers, giving a cost between those of the two pure backbones\.

Figure 2:Cached\-decoding efficiency vs\. context length \(3B, single H100 80 GB,bf16, cudagraph, batch 1\)\.\(a\) per\-step decode latency, \(b\) end\-to\-end decode throughput \(S=16S\{=\}16\), \(c\) peak memory\.Mambais flat at every length to 256k \(constant SSM state\), whileAttngrows as𝒪⁡\(L\)\\mathcal\{O\}\(L\): by 256kMambais2\.6×\\mathbf\{2\.6\\times\}faster \(296 vs\. 114 tok/s\) at𝟏𝟏×\\mathbf\{11\\times\}less memory \(7\.5 vs\. 82 GB\);Hybridinterpolates\.[Figure4](https://arxiv.org/html/2609.11998#A3.F4)sweeps batch11–88\.
### 3\.3Inference via cached generation

Generation proceeds block by block \([Algorithm1](https://arxiv.org/html/2609.11998#alg1)\)\. Entering blockii, we clone the layer cache so its entry state can be reused across the block’s denoising steps \(each cached forward returns fresh tensors and never mutates its inputs\)\. We then runSSdenoising steps: at each step a cached forward pass scores the current block, and we reveal the⌈m/S⌉\\lceil m/S\\rceilhighest\-confidence still\-masked positions, wheremmis the number of masked tokens at block entry\. Once the block is fully revealed, one final cached pass folds it into the cache that blocki\+1i\{\+\}1reads\. WithSSthe only quality–speed knob, the total number of forward passes is𝒪⁡\(K​S\)=𝒪⁡\(\(L/G\)​S\)\\mathcal\{O\}\(KS\)=\\mathcal\{O\}\(\(L/G\)S\), and for the SSM backbone each pass touches only theGGcurrent tokens regardless ofLL\.

### 3\.4Complexity and asymptotic throughput

[Table1](https://arxiv.org/html/2609.11998#S4.T1)summarizes the cost of one cached denoising step and the state it carries\. Decoding the frontier block ofGGtokens against a prefix of lengthLLcosts𝒪⁡\(G​L\)\\mathcal\{O\}\(GL\)for attention – each new query attends over the growing KV\-cache – so its per\-step latency grows linearly inLLand the decode throughput scales asT=𝒪⁡\(1/L\)T=\\mathcal\{O\}\(1/L\)\. The Mamba model consumes a fixed\-size SSM state and touches only theGGcurrent tokens, giving anLL\-independent per\-step cost of𝒪⁡\(G2\)\\mathcal\{O\}\(G^\{2\}\)and hence constant throughput\. The hybrid inherits the attention scaling with a small constant \(5/285/28of its layers carry attention\)\.[Section5\.1](https://arxiv.org/html/2609.11998#S5.SS1)measures exactly these three curves\.

## 4Experimental Setup

#### Models\.

We compare three 3B\-parameter denoisers that are identical except for the sequence mixer \([Table1](https://arxiv.org/html/2609.11998#S4.T1)\):Attn\(28 attention layers\),Mamba\(28 bidirectional Mamba\-2 layers\), andHybrid\(5 attention layers interleaved with 23 Mamba layers, following[Singh et al\. \[32\]](https://arxiv.org/html/2609.11998#bib.bib32)\)\. All models usedmodel=2560d\_\{\\text\{model\}\}=2560, block sizeG=32G=32and are trained on Nemotron\-CC\[[33](https://arxiv.org/html/2609.11998#bib.bib33)\]with300300B token budget under the single\-frontier block\-diffusion objective \([Section3\.1](https://arxiv.org/html/2609.11998#S3.SS1)\) which masks a random subset of a single sampled frontier block per example, keeping earlier blocks clean and later blocks fully masked, so training sees exactly the clean\-prefix, all\-masked\-future context each block meets at cached decode time\. We trained with sequence length of10241024tokens, and the global batch being40964096sequences \(1024×4096≈4\.21024\\times 4096\\approx 4\.2M tokens per step\) with a71 52571\\,525\-step schedule covering1024×4096×71 525≈3001024\\times 4096\\times 71\\,525\\approx 300B tokens\. Optimization is AdamW inbf16\(learning rate1×10−41\\times 10^\{\-4\}, cosine schedule,5%5\\%warmup, weight decay0\.10\.1, gradient\-norm clip1\.01\.0\); the complete set of training and model hyperparameters is in[AppendixA](https://arxiv.org/html/2609.11998#A1)\.

Table 1:Three denoisers, one variable: the mixer\.The 3B models sharedmodeld\_\{\\text\{model\}\}, depth, block sizeGG, and differ only in the sequence mixer\. The right block gives each backbone’s per\-step decode cost over a block ofGGtokens against a prefix of lengthLL: only the SSM cache is𝒪⁡\(1\)\\mathcal\{O\}\(1\)inLL\(conv\+SSM state vs\. an𝒪⁡\(L\)\\mathcal\{O\}\(L\)KV\-cache\), the root ofMamba’s flat memory and latency;∗small constant \(5/285/28layers carry attention\)\.
#### Hardware and timing\.

All efficiency measurements use a single NVIDIA H100 80 GB inbf16with PyTorch SDPA attention kernels\. We time one denoising step by CUDA\-graph capture which removes kernel launch overhead, and report the mean of 50 replays under CUDA events\. End\-to\-end decode throughput is the ratio of total tokens generated to the total time to decode them, adding up each block’s latency at its true cache depth \(blockiiis decoded against a prefix of lengthi⋅Gi\\cdot G\)\.

#### Evaluation protocols\.

We measured*Efficiency*by sweeping context lengthLLfrom6464to256256k tokens at batch sizes11–88, and report per\-step decode latency, peak memory, and end\-to\-end decode throughput atS=16S\{=\}16denoising steps per block\. For*Long Context Performance*we measure retrieval with a RULER\-style\[[15](https://arxiv.org/html/2609.11998#bib.bib15)\]needle\-in\-a\-haystack \(NIAH\) passkey task, where a short key is hidden at one of five depths inside a long distractor context and must be recovered \(2020trials per \(length, depth\),LLfrom512512to1616k\)\. We further evaluated the models on LongBench\[[4](https://arxiv.org/html/2609.11998#bib.bib4)\], a1616\-task English\-and\-code long\-context suite \(middle\-truncated toLL,2020examples per task,LLfrom22k to1616k\)\. For the attention backbones we additionally report an NTK\-RoPE variant\[[25](https://arxiv.org/html/2609.11998#bib.bib25),[17](https://arxiv.org/html/2609.11998#bib.bib17)\]: a training\-free rescaling of RoPE’s rotation frequencies that folds the longer positions seen at inference back toward the range covered during training, extending the usable window without any fine\-tuning\. For*Downstream Performance*we score eight common\-sense / reasoning tasks with the block\-diffusion likelihood harness, and for evaluating*Generation Quality*we report generative perplexity \(scored by GPT\-2 large\) alongside MAUVE\[[26](https://arxiv.org/html/2609.11998#bib.bib26)\], which quantifies the distributional gap between generated and reference text\.

## 5Results

### 5\.1Efficiency: the state\-space cache is constant in length

[Figure2](https://arxiv.org/html/2609.11998#S3.F2)and[Table9](https://arxiv.org/html/2609.11998#A3.T9)report per\-step latency, peak memory, and decode throughput against context length at batch 1\. The Mamba block cache delivers exactly𝒪⁡\(1\)\\mathcal\{O\}\(1\)complexity across the whole6464–256256k\-token range, i\.e\.Mambaholds a*flat*6\.86\.8ms per step,7\.57\.5GB of memory, and∼\\sim296296tok/s, its state never growing with the prefix\.Attninstead pays the𝒪⁡\(L\)\\mathcal\{O\}\(L\)KV\-cache in every panel: from6464to256256k tokens its latency climbs→295\.5\\\!\\to\\\!29ms \(5\.3×5\.3\\times\), its decoding throughput falls→114346\\\!\\to\\\!114tok/s \(3×3\\times\), and its memory balloons→826\.5\\\!\\to\\\!82GB, essentially the whole GPU\. By256256k,Mambathus beatsAttnon all three axes at once:4\.3×4\.3\\timeslower latency \(6\.86\.8vs\.2929ms\),𝟏𝟏×11\\timesless memory \(7\.57\.5vs\.8282GB\), and2\.6×2\.6\\timeshigher throughput \(297297vs\.114114tok/s\)\.Hybridis the pragmatic sweet spot: with only55of its2828layers attentional it tracks the attention curves at5/285/28of the slope, so at256256k it keeps full attention for global recall while still beatingAttnby2\.7×2\.7\\timeson latency \(10\.710\.7vs\.2929ms\),3\.8×3\.8\\timeson memory \(21\.521\.5vs\.8282GB\), and2\.0×2\.0\\timeson throughput \(231231vs\.114114tok/s\)\.

The constant footprint pays off most under batching as shown in[Figure4](https://arxiv.org/html/2609.11998#A3.F4): becauseMamba’s per\-stream memory does not grow, many sequences fit on one accelerator, and at256256k it sustains15931593tok/s at batch88\(a flat∼\\sim88GB\), whereasAttn’s𝒪⁡\(B⋅L\)\\mathcal\{O\}\(B\\cdot L\)KV\-cache is already out of memory at batch22and capped at114114tok/s; a𝟏𝟒×\\mathbf\{14\\times\}aggregate\-throughput advantage \([Appendix C](https://arxiv.org/html/2609.11998#A3)\)\. This dominance is a long\-context effect: below∼\\sim1616k the three backbones sit within a few percent \(withAttnmarginally ahead on its cheap short\-context blocks\), but where long\-context decoding actually lives, at tens to hundreds of thousands of tokens,MambaandHybridwin on latency, memory, and throughput at once\.

### 5\.2Long\-context generalization: linear state extrapolates, attention collapses

The position\-free recurrence that makes the SSM cache cheap also makes it length robust\. The mechanism is the block\-causal mask: regardless of block size, a query still attends over the*entire*clean prefix, so pastL=1024L\{=\}1024the attention layers meet absolute RoPE positions and prefix lengths never seen in training \(an out\-of\-distribution regime\), whereas the recurrent state carries no positions and simply folds a longer prefix into the same fixed\-size memory\.[Table2](https://arxiv.org/html/2609.11998#S5.T2)bears this out on NIAH passkey retrieval\. Within the10241024training length all three models retrieve near\-perfectly; at2×2\\times\(L=2L\{=\}2k\)Attn*collapses*to12%12\\%whileHybridholds53%53\\%andMamba76%76\\%, and by16×16\\times\(L=16L\{=\}16k\)Attnsits at0%0\\%whileMambastill recovers22%22\\%of needles\.Mambaowns the highest floor at the extremes \(2×2\\timesand16×16\\times\) andHybridthe middle of the range \(4×4\\times, and8×8\\timeswith NTK\-RoPE\)\.

Table 2:Attention collapses at2×2\\timestrain length; linear state extrapolates to88–16×16\\times\.Needle\-in\-a\-haystack passkey retrieval accuracy \(%\) vs\. context lengthLL\(3B; trained atL=1024L\{=\}1024, soL≥2L\{\\geq\}2k is extrapolation\)\. “\+NTK” adds NTK\-RoPE to the attention\-bearing models; best per lengthshaded\.Attndrops to12%12\\%at2×2\\timesand0%0\\%by8×8\\times, whileMamba/Hybridstill retrieve needles at16×16\\times\.![Refer to caption](https://arxiv.org/html/2609.11998v1/niah_heatmap_3b.png)Figure 3:NIAH retrieval by needle depth×\\timescontext length \(3B\)\.Token accuracy at five depths \(0%0\\%= oldest,100%100\\%= most recent\) across lengths\.Attncollapses almost uniformly past the10241024\-token training length;Mamba/Hybriddegrade*gracefully*and retain near\-perfect retrieval of*recent*needles \(rightmost column:Mamba99%99\\%at16×16\\timestraining length; per\-depth numbers in[AppendixB](https://arxiv.org/html/2609.11998#A2)\)\. This matches the “local perception” / sliding\-window behavior reported for diffusion LMs by LongLLaDA\[[17](https://arxiv.org/html/2609.11998#bib.bib17)\], which appears here as a property of the linear\-state backbone rather than of an approximate cache\.The same pattern holds on*realistic*tasks \([Table3](https://arxiv.org/html/2609.11998#S5.T3)\)\.Mambais the strongest backbone and barely degrades with length, holding essentially constant at10\.210\.2–11\.211\.2from22k to1616k and ahead ofHybrid\(→9\.110\.8\\\!\\to\\\!9\.1\) at every length\.Attn, by contrast, falls by a third by44k \(→4\.77\.1\\\!\\to\\\!4\.7\) and roughly halves by1616k \(→3\.6\\to\\\!3\.6\), so the gap widens with context: at1616kMambanearly triples andHybridmore than doublesAttn\(10\.210\.2and9\.19\.1vs\.3\.63\.6\), just as a position\-agnostic recurrence that needs no length extrapolation predicts\.

Table 3:On realistic long\-context tasks, linear state stays flat while attention slides\.LongBench macro\-average \(%\) over 16 English\+code tasks vs\. lengthLL\(3B, cached greedy decoding, middle\-truncation, 20 examples/task; every rowL≥2L\{\\geq\}2k is extrapolation from the10241024training length\)\. “\+NTK”==NTK\-RoPE on the attention\-bearing models; best per lengthshaded\.Mambaholds10\.210\.2–11\.211\.2throughout;Attnmore than halves by44k and reaches3\.63\.6at1616k\.Table 4:Constant\-memory decoding at no downstream cost\.Accuracy on eight common\-sense/reasoning tasks \(3B\) via the block\-diffusion likelihood harness \(block\-decomposed MC\-ELBO, 128 samples, CFG0\.50\.5\);Avgis the mean, bestshaded\.AttnandHybridtie \(0\.4340\.434/0\.4320\.432\) andMambatrails by∼\\sim1 point\.Table 5:One dial,SS, trades quality for speed\.Per\-block denoising budgetSS\(3B, cached generation,64→12864\\to 128\-token continuations, 512 samples\): largerSSmonotonically lowers Gen\-PPL and broadly raises MAUVE, at55–6×6\\timeslower throughput \(best per columnshaded\)\. MAUVE \(kk\-means over GPT\-2\-large features\) is noisy for small base models, so Gen\-PPL is the cleaner signal\.
### 5\.3Quality parity: efficiency at no cost

The efficiency and long\-context gains cost essentially nothing in quality\. The three models sit within∼\\sim0\.03 nats on validation NLL, withHybridbest \([Table1](https://arxiv.org/html/2609.11998#S4.T1)\), and on an eight\-task downstream suite scored by the block\-diffusion likelihood harness \([Table4](https://arxiv.org/html/2609.11998#S5.T4)\)AttnandHybridare level \(macro0\.4340\.434vs\.0\.4320\.432\) whileMambatrails by about a point\. Thus the hybrid delivers its88–16×16\\timeslength generalization and steeply reduced decode memory at no measured quality cost, and the pure\-SSMMambagives up only∼\\sim1 downstream point for fully constant\-memory, length\-independent inference\.

#### Quality and throughput trade on a single dial\.

The per\-block denoising budgetSStrades quality for speed as shown in[Table5](https://arxiv.org/html/2609.11998#S5.T5)\. For all three backbones, raisingSSfrom44to3232*monotonically*lowers generative perplexity \(15\.3→6\.415\.3\\to 6\.4forAttn,17\.1→6\.517\.1\\to 6\.5forHybrid,18\.5→6\.518\.5\\to 6\.5forMamba\) and broadly improves MAUVE, at55–6×6\\timeslower throughput, so a deployment can choose its operating point by turning one knob\. One caveat: the absolute MAUVE values are low and noisy \(they peak nearS=16S\{=\}16rather than rising strictly\), because these are small*base*models producing short continuations, so Gen\-PPL is the cleaner quality signal here\.

### 5\.4The gains are architectural, not parameter count

A natural worry is that Mamba’s extra projection parameters, rather than its linear\-time mixing, drive the results\.[Table6](https://arxiv.org/html/2609.11998#S5.T6)shows the advantage is a property of the*architecture*, not the parameter count\.MambaandHybriddo carry13%13\\%and11%11\\%*more*parameters thanAttn, but those parameters live in𝒪⁡\(d2\)\\mathcal\{O\}\(d^\{2\}\)projections that are independent ofLLand add no long\-context cost: their forward FLOPs/token stay flat with context \(Mambaat6\.26\.2G for*every*length\), whereas attention’s parameter\-free𝒪⁡\(L​d\)\\mathcal\{O\}\(Ld\)term grows to24\.224\.2G\. The three are comparable at the10241024training length, then diverge sharply: atL=64L\{=\}64kMambaandHybridspend just0\.26×0\.26\\timesand0\.39×0\.39\\timesthe attention FLOPs\. With quality already at parity \([Section5\.3](https://arxiv.org/html/2609.11998#S5.SS3)\), the long\-context win rests on the linear\-time mixer, not on the parameter budget:MambaandHybridwin decisively at long context\.

We calculate forward\-pass FLOPs per token*analytically*from the architecture \(two FLOPs per multiply–accumulate\), summing every layer’s mixer and gated FFN plus the untied output head over the vocabulary; kernels are not profiled\. Per token, an attention layer costs8​d28d^\{2\}for theQ,K,V,OQ,K,V,Oprojections plus4​L​d4Ldfor the score and value\-aggregation matmuls, the*only*LL\-dependent term\. A bidirectional Mamba\-2 layer costs2​\[2​d​\(2​din\+2​dstate\+nh\)\+2​din​d\+2​din​dstate\+2​dconv​\(din\+2​dstate\)\]2\\big\[\\,2d\(2d\_\{\\text\{in\}\}\{\+\}2d\_\{\\text\{state\}\}\{\+\}n\_\{h\}\)\+2d\_\{\\text\{in\}\}d\+2d\_\{\\text\{in\}\}d\_\{\\text\{state\}\}\+2d\_\{\\text\{conv\}\}\(d\_\{\\text\{in\}\}\{\+\}2d\_\{\\text\{state\}\}\)\\,\\big\]withdin=dd\_\{\\text\{in\}\}\{=\}dandnhn\_\{h\}Mamba heads \(in/out projections, the SSD scan, and the depthwise convolution, doubled for the two directions\), all independent ofLL; the output head adds2​d​V2dV\. Unlike the smaller\-scale DiffuMamba hybrids\[[32](https://arxiv.org/html/2609.11998#bib.bib32)\], the 3B Mamba/hybrid variants keep the*same*MLP as attention, so their extra mixer parameters are not offset\.[Table6](https://arxiv.org/html/2609.11998#S5.T6)evaluates this forward count for the 3B models \(d=2560d\{=\}2560,2828layers, hybrid=5=5attention\+23\+\\,23Mamba, gated FFN with hidden size76807680,dstate=64d\_\{\\text\{state\}\}\{=\}64, expand11,dconv=4d\_\{\\text\{conv\}\}\{=\}4, and an untied vocabulary ofV=126 464V\{=\}126\\,464\) atL=1024L\{=\}1024andL=64L\{=\}64k; parameter counts are read exactly from the trained checkpoints\.

Table 6:More parameters, but the compute win is long\-context\.Exact parameters and analytic forward FLOPs/token for the 3B models \([Section5\.4](https://arxiv.org/html/2609.11998#S5.SS4)\)\.MambaandHybridcarry1111–13%13\\%*more*parameters thanAttn, but these live inLL\-independent𝒪⁡\(d2\)\\mathcal\{O\}\(d^\{2\}\)projections that add no long\-context cost: per\-token FLOPs are comparable at short context, while only the SSM cost stays flat with length \(shaded\), so atL=64L\{=\}64kMambaspends0\.26×0\.26\\timesthe attention FLOPs whileAttn’s𝒪⁡\(L​d\)\\mathcal\{O\}\(Ld\)term dominates\.

## 6Conclusion

We pretrained attention, bidirectional\-Mamba, and hybrid block\-diffusion denoisers at 3B parameters and300300B tokens under one block\-causal frontier objective, and decoded all three through a single,*exact*cached interface whose state\-space variant is𝒪⁡\(1\)\\mathcal\{O\}\(1\)in sequence length\. Concurrent small\-scale work\[[7](https://arxiv.org/html/2609.11998#bib.bib7)\]established that such a cache is possible; our results show what it delivers at scale\. At 256k tokensMambadecodes at a flat6\.86\.8ms/step and7\.57\.5GB whereAttnneeds2929ms/step and8282GB, and because that per\-stream footprint never grows it scales with batch to14×14\\timesthe aggregate throughput\. The same linear\-state bias makes decoding length\-robust:MambaandHybridretrieve needles and answer LongBench queries at88–16×16\\timestheir training length whileAttn’s retrieval collapses at2×2\\times, and the hybrid matchesAttnon likelihood and downstream accuracy, so these gains come at no measured quality cost\.

#### Limitations and future work\.

Our models are trained at a10241024\-token context, so the long\-context evaluations probe*extrapolation*rather than trained long\-context ability, which explains the modest absolute accuracies far past training length; closing this gap \(via RoPE\-scaling recipes such as LongLLaDA/UltraLLaDA\[[17](https://arxiv.org/html/2609.11998#bib.bib17),[14](https://arxiv.org/html/2609.11998#bib.bib14)\], or simply training at longer contexts\) is orthogonal to, and compatible with, our cache\. Likewise, we train with a single\-frontier objective at a10241024\-token context, whereas[Chaturvedi et al\. \[7\]](https://arxiv.org/html/2609.11998#bib.bib7)use an all\-block objective at88k tokens; disentangling the effect of objective and training length on cached long\-context behavior is a natural ablation we leave to future work\. Natural next steps are to combine our exact cache with confidence\-aware parallel decoding\[[39](https://arxiv.org/html/2609.11998#bib.bib39)\]inside the cached generator and to push block\-cached hybrids to larger scales\. We believe constant\-memory, length\-generalizing block decoding is a promising foundation for efficient long\-context diffusion language models\.

## References

- \[1\]J\. Achiam, S\. Adler, S\. Agarwal, L\. Ahmad, I\. Akkaya, F\. L\. Aleman, D\. Almeida, J\. Altenschmidt, S\. Altman, S\. Anadkat, et al\.Gpt\-4 technical report\.*arXiv preprint arXiv:2303\.08774*, 2023\.
- \[2\]M\. Arriola, A\. Gokaslan, J\. Chiu, Z\. Yang, Z\. Qi, J\. Han, S\. Sahoo, and V\. Kuleshov\.Block diffusion: Interpolating between autoregressive and diffusion language models\.In*International Conference on Learning Representations*, volume 2025, pages 50726–50753, 2025\.
- \[3\]J\. Austin, D\. D\. Johnson, J\. Ho, D\. Tarlow, and R\. Van Den Berg\.Structured denoising diffusion models in discrete state\-spaces\.*Advances in neural information processing systems*, 34:17981–17993, 2021\.
- \[4\]Y\. Bai, X\. Lv, J\. Zhang, H\. Lyu, J\. Tang, Z\. Huang, Z\. Du, X\. Liu, A\. Zeng, L\. Hou, Y\. Dong, J\. Tang, and J\. Li\.LongBench: A bilingual, multitask benchmark for long context understanding\.In*Annual Meeting of the Association for Computational Linguistics \(ACL\)*, 2024\.
- \[5\]T\. Brown, B\. Mann, N\. Ryder, M\. Subbiah, J\. D\. Kaplan, P\. Dhariwal, A\. Neelakantan, P\. Shyam, G\. Sastry, A\. Askell, et al\.Language models are few\-shot learners\.*Advances in neural information processing systems*, 33:1877–1901, 2020\.
- \[6\]A\. Campbell, J\. Benton, V\. De Bortoli, T\. Rainforth, G\. Deligiannidis, and A\. Doucet\.A continuous time framework for discrete denoising models\.*Advances in Neural Information Processing Systems*, 35:28266–28279, 2022\.
- \[7\]P\. Chaturvedi, P\. Shroff, T\. Suresh, H\. Kang, and K\. Wen\.Training hybrid block diffusion language models with partial bidirectionality\.*arXiv preprint arXiv:2607\.02805*, 2026\.
- \[8\]A\. Chowdhery, S\. Narang, J\. Devlin, M\. Bosma, G\. Mishra, A\. Roberts, P\. Barham, H\. W\. Chung, C\. Sutton, S\. Gehrmann, et al\.Palm: Scaling language modeling with pathways\.*Journal of machine learning research*, 24\(240\):1–113, 2023\.
- \[9\]T\. Dao and A\. Gu\.Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality\.In*Forty\-first International Conference on Machine Learning*, 2024\.URL[https://openreview\.net/forum?id=ztn8FCR1td](https://openreview.net/forum?id=ztn8FCR1td)\.
- \[10\]S\. De, S\. L\. Smith, A\. Fernando, A\. Botev, G\. Cristian\-Muraru, A\. Gu, R\. Haroun, L\. Berrada, Y\. Chen, S\. Srinivasan, et al\.Griffin: Mixing gated linear recurrences with local attention for efficient language models\.*arXiv preprint arXiv:2402\.19427*, 2024\.
- \[11\]D\. Y\. Fu, T\. Dao, K\. K\. Saab, A\. W\. Thomas, A\. Rudra, and C\. Re\.Hungry hungry hippos: Towards language modeling with state space models\.In*The Eleventh International Conference on Learning Representations*, 2023\.URL[https://openreview\.net/forum?id=COZDy0WYGg](https://openreview.net/forum?id=COZDy0WYGg)\.
- \[12\]A\. Gu and T\. Dao\.Mamba: Linear\-time sequence modeling with selective state spaces\.In*First Conference on Language Modeling*, 2024\.URL[https://openreview\.net/forum?id=tEYskw1VY2](https://openreview.net/forum?id=tEYskw1VY2)\.
- \[13\]A\. Gu, K\. Goel, and C\. Re\.Efficiently modeling long sequences with structured state spaces\.In*International Conference on Learning Representations*, 2022\.URL[https://openreview\.net/forum?id=uYLFoz1vlAC](https://openreview.net/forum?id=uYLFoz1vlAC)\.
- \[14\]G\. He, S\. Nie, F\. Zhu, Y\. Zhao, T\. Bai, R\. Yan, J\. Fu, C\. Li, and B\. Yuan\.Ultrallada: Scaling the context length to 128k for diffusion large language models\.*arXiv preprint arXiv:2510\.10481*, 2025\.
- \[15\]C\.\-P\. Hsieh, S\. Sun, S\. Kriman, S\. Acharya, D\. Rekesh, F\. Jia, Y\. Zhang, and B\. Ginsburg\.Ruler: What’s the real context size of your long\-context language models?*arXiv preprint arXiv:2404\.06654*, 2024\.COLM 2024\.
- \[16\]O\. Lieber, B\. Lenz, H\. Bata, G\. Cohen, J\. Osin, I\. Dalmedigos, E\. Safahi, S\. Meirom, Y\. Belinkov, S\. Shalev\-Shwartz, et al\.Jamba: A hybrid transformer\-mamba language model\.*arXiv preprint arXiv:2403\.19887*, 2024\.
- \[17\]X\. Liu, Y\. Song, Z\. Liu, Z\. Huang, Q\. Guo, Z\. He, and X\. Qiu\.Longllada: Unlocking long context capabilities in diffusion llms\.*arXiv preprint arXiv:2506\.14429*, 2025a\.
- \[18\]Z\. Liu, Y\. Yang, Y\. Zhang, J\. Chen, C\. Zou, Q\. Wei, S\. Wang, Y\. Zhu, and L\. Zhang\.dllm\-cache: Accelerating diffusion large language models with adaptive caching\.*arXiv preprint arXiv:2506\.06295*, 2025b\.
- \[19\]L\. Long, Y\. Huang, S\. Bai, R\. Gong, J\. Zhang, A\. Zhou, and J\. Yang\.Focus\-dllm: Accelerating long\-context diffusion llm inference via confidence\-guided context focusing\.*arXiv preprint arXiv:2602\.02159*, 2026\.
- \[20\]A\. Lou, C\. Meng, and S\. Ermon\.Discrete diffusion modeling by estimating the ratios of the data distribution\.In*Forty\-first International Conference on Machine Learning*, 2024\.URL[https://openreview\.net/forum?id=CNicRIVIPA](https://openreview.net/forum?id=CNicRIVIPA)\.
- \[21\]X\. Ma, R\. Yu, G\. Fang, and X\. Wang\.dkv\-cache: The cache for diffusion language models\.*Advances in Neural Information Processing Systems*, 38:149009–149033, 2025\.
- \[22\]Q\. Nguyen\-Tri, M\. Ranjan, and Z\. Shen\.Attention is all you need for kv cache in diffusion llms\.*arXiv preprint arXiv:2510\.14973*, 2025\.
- \[23\]J\. Ni, Q\. Liu, L\. Dou, C\. Du, Z\. Wang, H\. Yan, T\. Pang, and M\. Q\. Shieh\.Diffusion language models are super data learners\.*arXiv preprint arXiv:2511\.03276*, 2025\.
- \[24\]S\. Nie, F\. Zhu, Z\. You, X\. Zhang, J\. Ou, J\. Hu, J\. Zhou, Y\. Lin, J\.\-R\. Wen, and C\. Li\.Large language diffusion models\.*Advances in Neural Information Processing Systems*, 38:50608–50646, 2025\.
- \[25\]B\. Peng, J\. Quesnelle, H\. Fan, and E\. Shippole\.YaRN: Efficient context window extension of large language models\.In*International Conference on Learning Representations \(ICLR\)*, 2024\.
- \[26\]K\. Pillutla, S\. Swayamdipta, R\. Zellers, J\. Thickstun, S\. Welleck, Y\. Choi, and Z\. Harchaoui\.MAUVE: Measuring the gap between neural text and human text using divergence frontiers\.In*Advances in Neural Information Processing Systems*, 2021\.
- \[27\]M\. Poli, S\. Massaroli, E\. Nguyen, D\. Y\. Fu, T\. Dao, S\. Baccus, Y\. Bengio, S\. Ermon, and C\. Ré\.Hyena hierarchy: Towards larger convolutional language models\.In*International Conference on Machine Learning*, pages 28043–28078\. PMLR, 2023\.
- \[28\]L\. Ren, Y\. Liu, Y\. Lu, Y\. Shen, C\. Liang, and W\. Chen\.Samba: Simple hybrid state space models for efficient unlimited context language modeling\.In*International Conference on Learning Representations*, 2025\.
- \[29\]S\. S\. Sahoo, M\. Arriola, Y\. Schiff, A\. Gokaslan, E\. Marroquin, J\. T\. Chiu, A\. Rush, and V\. Kuleshov\.Simple and effective masked diffusion language models\.*Advances in Neural Information Processing Systems*, 37:130136–130184, 2024\.
- \[30\]S\. S\. Sahoo, J\. Deschenaux, A\. Gokaslan, G\. Wang, J\. T\. Chiu, and V\. Kuleshov\.The diffusion duality\.In*Forty\-second International Conference on Machine Learning*, 2025\.URL[https://openreview\.net/forum?id=9P9Y8FOSOk](https://openreview.net/forum?id=9P9Y8FOSOk)\.
- \[31\]J\. Shi, K\. Han, Z\. Wang, A\. Doucet, and M\. Titsias\.Simplified and generalized masked diffusion for discrete data\.In A\. Globerson, L\. Mackey, D\. Belgrave, A\. Fan, U\. Paquet, J\. Tomczak, and C\. Zhang, editors,*Advances in Neural Information Processing Systems*, volume 37, pages 103131–103167\. Curran Associates, Inc\., 2024\.doi:10\.52202/079017\-3277\.
- \[32\]V\. Singh, O\. Ostapenko, P\.\-A\. Noël, E\. Belilovsky, and T\. Scholak\.Diffumamba: High\-throughput diffusion lms with mamba backbone\.In*International Conference on Machine Learning*, 2026\.
- \[33\]D\. Su, K\. Kong, Y\. Lin, J\. Jennings, B\. Norick, M\. Kliegl, M\. Patwary, M\. Shoeybi, and B\. Catanzaro\.Nemotron\-CC: Transforming Common Crawl into a refined long\-horizon pretraining dataset\.*arXiv preprint arXiv:2412\.02595*, 2024\.ACL 2025\.
- \[34\]H\. Touvron, T\. Lavril, G\. Izacard, X\. Martinet, M\.\-A\. Lachaux, T\. Lacroix, B\. Rozière, N\. Goyal, E\. Hambro, F\. Azhar, et al\.Llama: Open and efficient foundation language models\.*arXiv preprint arXiv:2302\.13971*, 2023\.
- \[35\]R\. Waleffe, W\. Byeon, D\. Riach, B\. Norick, V\. Korthikanti, T\. Dao, A\. Gu, A\. Hatamizadeh, S\. Singh, D\. Narayanan, G\. Kulshreshtha, V\. Singh, J\. Casper, J\. Kautz, M\. Shoeybi, and B\. Catanzaro\.An empirical study of mamba\-based language models\.*arXiv preprint arXiv:2406\.07887*, 2024\.
- \[36\]D\. Wang, R\.\-J\. Zhu, S\. Abreu, Y\. Shan, T\. Kergan, Y\. Pan, Y\. Chou, Z\. Li, G\. Zhang, W\. Huang, et al\.A systematic analysis of hybrid linear attention\.*arXiv preprint arXiv:2507\.06457*, 2025a\.
- \[37\]X\. Wang, C\. Xu, Y\. Jin, J\. Jin, H\. Zhang, and Z\. Deng\.Diffusion llms can do faster\-than\-ar inference via discrete diffusion forcing\.*arXiv preprint arXiv:2508\.09192*, 2025b\.
- \[38\]C\. Wu, H\. Zhang, S\. Xue, S\. Diao, Y\. Fu, Z\. Liu, P\. Molchanov, P\. Luo, S\. Han, and E\. Xie\.Fast\-dllm v2: Efficient block\-diffusion LLM\.*arXiv preprint arXiv:2509\.26328*, 2025\.
- \[39\]C\. Wu, H\. Zhang, S\. Xue, Z\. Liu, S\. Diao, L\. Zhu, P\. Luo, S\. Han, and E\. Xie\.Fast\-dLLM: Training\-free acceleration of diffusion LLM by enabling KV cache and parallel decoding\.In*International Conference on Learning Representations \(ICLR\)*, 2026\.
- \[40\]J\. Xiong, Q\. Han, S\. Gong, Y\. Hsieh, C\. Wu, C\. Tao, C\. Zhao, and N\. Wong\.Prefilling\-dllm: Predictive prefilling for long\-context inference in diffusion language models\.*arXiv preprint arXiv:2606\.10537*, 2026\.
- \[41\]J\. Ye, Z\. Xie, L\. Zheng, J\. Gao, Z\. Wu, X\. Jiang, Z\. Li, and L\. Kong\.Dream 7b: Diffusion large language models\.*arXiv preprint arXiv:2508\.15487*, 2025\.

## Appendix ATraining and model hyperparameters

[Table7](https://arxiv.org/html/2609.11998#A1.T7)lists the full training and architecture configuration\. The three denoisers are identical except for the sequence mixer \([Table1](https://arxiv.org/html/2609.11998#S4.T1)\); every other setting is shared\.

Table 7:Full training and model hyperparameters, shared acrossAttn,Hybrid, andMamba\(the sequence mixer is the only difference;[Table1](https://arxiv.org/html/2609.11998#S4.T1)\)\.
## Appendix BPer\-depth needle\-in\-a\-haystack

[Table8](https://arxiv.org/html/2609.11998#A2.T8)reports the retrieval accuracy behind[Figure3](https://arxiv.org/html/2609.11998#S5.F3)by needle depth \(0 = far from the end\-of\-sequence query, 1\.0 = adjacent\) at the two lengths just past training length\. All models show the recency/sliding\-window effect, but only the linear\-state backbones retain any retrieval at4×4\\times\.

Table 8:The linear\-state advantage is recency\.Per\-depth NIAH retrieval \(%\) atL=2L\{=\}2k \(2×2\\times\) andL=4L\{=\}4k \(4×4\\times\) \(3B; depth00==oldest,1\.01\.0==adjacent to the query\)\. Every backbone peaks on*recent*needles, but onlyMamba/Hybridkeep retrieving older ones past training length;Attnis near\-zero at all but the last depth\.
## Appendix CFull batch\-size efficiency sweep

[Figure2](https://arxiv.org/html/2609.11998#S3.F2)plots batch 1 \(exact numbers in[Table9](https://arxiv.org/html/2609.11998#A3.T9)\);[Figure4](https://arxiv.org/html/2609.11998#A3.F4)gives the full sweep over batch sizes1/2/4/81/2/4/8\.Mambathroughput*scales with batch and stays flat in length*\(1593 tok/s at batch 8 for*every*length to 256k,∼\\sim5\.4×5\.4\\timesthe batch\-1 rate, at a constant∼\\sim8 GB\), whileAttn’s𝒪⁡\(B⋅L\)\\mathcal\{O\}\(B\\\!\\cdot\\\!L\)KV\-cache exhausts the H100 \(OOM\) at progressively shorter lengths as batch grows \(last feasible: 256k@B1, 128k@B2, 64k@B4, 32k@B8\)\. At 256k,Mambaat batch 8 delivers 1593 tok/s whileAttncannot run even batch 2; its best feasible aggregate is batch 1 at 114 tok/s, the∼\\sim14×14\\timesgap quoted in[Section5\.1](https://arxiv.org/html/2609.11998#S5.SS1)\.

Table 9:Flat vs\.𝒪⁡\(L\)\\mathcal\{O\}\(L\), in numbers\.Batch\-1 per\-step latency \(ms\), peak memory \(GB\), and end\-to\-end decode throughput \(tok/s,S=16S\{=\}16\) vs\. context length \(3B, H100 80 GB, cudagraph\), the exact values behind[Figure2](https://arxiv.org/html/2609.11998#S3.F2)\.Mambais constant \(∼\\sim6\.86\.8ms,∼\\sim7\.5 GB\) out to 256k whileAttngrows to2929ms and8282GB; best per lengthshaded\.Figure 4:Efficiency vs\. context length across batch sizes \(3B, single H100 80 GB,bf16, cudagraph\)\.Top to bottom: per\-step decode latency, peak memory, aggregate decode throughput \(S=16S\{=\}16\); within each row the four panels are batch1/2/4/81/2/4/8\(left to right\)\.Mambastays∼\\sim8 GB at every \(batch, length\);Attn’s last\-feasible length shrinks with batch \(256k@B1, 128k@B2, 64k@B4, 32k@B8\), markedOOM\(×\\times\) beyond\.

Similar Articles

From Local Mismatch to Global Impact: Optimizing Cache Reuse Policy for Efficient Diffusion

arXiv cs.AI

This paper proposes Global-ImpactCache (GCache), a bilevel optimization framework that learns cache reuse policies for diffusion models by aligning error weighting with final generation quality, instead of relying on local similarity heuristics. It achieves significant speedups and quality improvements on image and video generation tasks, including a 2.17x speedup on Wan2.1 with lower LPIPS.

Affix Cache for Diffusion Large Language Models

arXiv cs.CL

ACache introduces a caching mechanism for Diffusion Large Language Models that selectively recomputes critical tokens to improve inference efficiency without losing accuracy.