Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

# Training Hybrid Block Diffusion Language Models with Partial Bidirectionality
Source: [https://arxiv.org/html/2607.02805](https://arxiv.org/html/2607.02805)
Pranshu Chaturvedi Stanford University pranshu@cs\.stanford\.edu&Parth Shroff Stanford University pshroff@cs\.stanford\.edu&Tarun Suresh Stanford University tsuresh@stanford\.edu&Hangoo Kang Stanford University hangook@stanford\.edu&Kaiyue Wen Stanford University kaiyuew@stanford\.edu

###### Abstract

High\-throughput long\-context generation is one of the central challenges for large language models\. Generation is typically memory\-bandwidth\-bound rather than compute\-bound: each decoding step must stream the accumulated key/value \(KV\) cache from memory, so bandwidth demand grows with context length while only one token is emitted\. Two parallel approaches have therefore emerged: reducing memory access with efficient attention variants and linear\-time mixers such as Mamba, or increasing parallel computation by generating blocks of tokens at once\. However, technical challenges arise when combining these two ideas\. Earlier hybrid diffusion models such as DiffuMamba use bidirectional Mamba mixing, including a reverse\-direction scan relative to causal generation\. This reverse scan needs to scan the entire sequence, so its states are not prefix\-only and cannot be precisely reused as a cache even when diffusion is performed block by block\. We propose a BDLM Mamba–attention hybrid that addresses this challenge by restricting the reverse Mamba scan to the active denoising block, which enables exact caching across blocks\. In an 87M\-parameter DCLM sweep, BDLM Mamba\-H achieves the best C4\-en validation perplexity compared to BDLM attention and full\-sequence baselines\. At 350M parameters, it remains competitive with BDLM attention\. For long\-context inference, BDLM Mamba\-H reaches 19\.7x the throughput of full\-sequence DiffuMamba\-H at 65K tokens and 3\.7x the throughput of BDLM attention at 262K, showing that Mamba hybrids are a potential long\-context diffusion architecture\.

## 1Introduction

Serving language models at long context remains largely a memory\-bandwidth problem\. In autoregressive decoding, every generated token reads the accumulated key/value \(KV\) cache, so bytes moved grow with context length while only one token is emitted\. Two lines of work address this pressure from different directions\. Hybrid autoregressive models reduce cache traffic by replacing many attention layers with linear\-time recurrent or state\-space mixers such as Mamba\(Gu and Dao,[2024](https://arxiv.org/html/2607.02805#bib.bib4); Dao and Gu,[2024](https://arxiv.org/html/2607.02805#bib.bib5); Qinet al\.,[2024](https://arxiv.org/html/2607.02805#bib.bib31); Kimi Teamet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib41); Li and others,[2025](https://arxiv.org/html/2607.02805#bib.bib33); Ling Teamet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib42)\)\. Masked diffusion language models instead generate many tokens in parallel by iterative denoising\(Austinet al\.,[2021a](https://arxiv.org/html/2607.02805#bib.bib36); Hoogeboomet al\.,[2021](https://arxiv.org/html/2607.02805#bib.bib37); Sahoo and others,[2024](https://arxiv.org/html/2607.02805#bib.bib38); Nieet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib39)\)\. Block diffusion language models \(BDLMs\) make diffusion cacheable by generating blocks left to right: completed prefix blocks become fixed conditioning, and the current block is denoised in parallel\(Arriola and others,[2025](https://arxiv.org/html/2607.02805#bib.bib40)\)\.

DiffuMamba combines diffusion with Mamba and shows strong long\-context inference throughput for full\-sequence denoisers\(Singhet al\.,[2026](https://arxiv.org/html/2607.02805#bib.bib13)\)\. Its fully bidirectional Mamba scan, however, includes a reverse\-direction pass over the whole denoising window\. Those reverse states depend on the active tokens and are therefore not reusable prefix states\. We study a native BDLM Mamba–attention hybrid, BDLM Mamba\-H, that restricts the reverse\-direction Mamba scan to the active denoising block\. Completed prefix blocks are stored as attention KV states for attention layers and as forward Mamba boundary states for Mamba layers\. The all\-block objective constructs these same prefix caches during training and lets downstream block losses train them directly, while the active block keeps local bidirectional mixing\.

In this work we explore whether cacheable BDLM Mamba hybrids can retain competitive validation quality while improving long\-context inference throughput\. We also describe timestep\-conditioning factorization that applies to only the current active block for BDLM Mamba\-H, so timestep modulation can be applied without making clean\-prefix caches depend on the reverse diffusion step \(Appendix[D](https://arxiv.org/html/2607.02805#A4)\)\. Our contributions are as follows\.

Clean prefixc\(k\)c^\{\(k\)\}CurrentblockFuture blocksmasked outb1⋆b\_\{1\}^\{\\star\}b2⋆b\_\{2\}^\{\\star\}b3⋆b\_\{3\}^\{\\star\}b~t\(k\)\\tilde\{b\}\_\{t\}^\{\(k\)\}bt\+1⋆b\_\{t\+1\}^\{\\star\}bt\+2⋆b\_\{t\+2\}^\{\\star\}Mamba PrefixscanCache⁡\(c\)\\operatorname\{Cache\}\(c\)Active blockdenoiserttCE lossGradient flow throughMamba Prefix states

Figure 1:BDLM Mamba\-H training\. The all\-block objective applies this active\-block computation at all block positions in parallel\. Solid arrows show forward computation, dashed green arrows show gradient flow, and dashed black arrows mark optional timestep input\. The denoiser receives a clean prefix cache, the noisy current block, and optional timestep conditioning while future blocks remain masked out\. For Mamba hybrid layers, the cached prefix is forward\-only and the reverse scan is local to the active block, so active\-block losses train the same prefix cache object used by the sampler\.- •Cache\-aligned training\.We formulate a BDLM all\-block objective whose cached objects match inference\-time block caches: attention KV states for attention layers and Mamba prefix states for recurrent layers\. For the BDLM Mamba\-H architecture, downstream block losses directly train the prefix scan that constructs the recurrent state reused at inference\.
- •Controlled validation and scale\-up\.We train full\-sequence attention, full\-sequence DiffuMamba\-H, BDLM attention, and BDLM Mamba\-H at 87M parameters on DCLM with tuned hyperparameters\. BDLM Mamba\-H reaches the best 87M C4\-en validation perplexity, 61\.6 compared with 76\.5 for BDLM attention and 83\.7 and 87\.6 for the full\-sequence baselines\. We scale BDLM attention and BDLM Mamba\-H to 350M parameters with a Complete\(d\)P\-style learning\-rate transfer rule\(Mlodozeniecet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib14)\); BDLM Mamba\-H remains comparable to BDLM attention in perplexity at this scale\.
- •Long\-context throughput gains\.We evaluate 350M training throughput at 8192\-token context on 8x A100\-80GB and 700M random\-initialized inference throughput on 1x A100\-80GB with batch size 1\. BDLM Mamba\-H overtakes BDLM attention \(and full\-sequence DiffuMamba\-H in throughput\) as generation length grows, reaching 19\.7x full\-sequence DiffuMamba\-H throughput at 65K and 3\.7x BDLM attention throughput at 262K\.

## 2Related Work

#### Diffusion language models\.

Discrete diffusion language models corrupt text with masking or discrete noise and train denoisers to reconstruct clean tokens\(Austinet al\.,[2021a](https://arxiv.org/html/2607.02805#bib.bib36); Hoogeboomet al\.,[2021](https://arxiv.org/html/2607.02805#bib.bib37); Sahoo and others,[2024](https://arxiv.org/html/2607.02805#bib.bib38); Nieet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib39)\)\. BDLMs add an autoregressive block factorization so each completed block becomes clean prefix context for later blocks\(Arriola and others,[2025](https://arxiv.org/html/2607.02805#bib.bib40)\)\. This gives block diffusion a native cache boundary, unlike full\-sequence MDLMs whose hidden states are tied to the current noisy sequence\. Training\-free cache methods retrofit reuse onto pretrained diffusion language models through approximate KV reuse or adaptive caching\(Wuet al\.,[2026](https://arxiv.org/html/2607.02805#bib.bib8); Maet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib9); Nguyen\-Triet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib11); Liuet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib10); Wanget al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib12)\)\. Our focus is instead to train a block\-diffusion architecture whose cache is part of both the objective and the sampler\.

#### Mamba hybrids for long\-context generation\.

Mamba replaces attention over token pairs with a selective state\-space scan whose transition, input, and output maps are functions of the current token\(Gu and Dao,[2024](https://arxiv.org/html/2607.02805#bib.bib4)\)\. In the notation used below, a Mamba layer carries a recurrent boundary statesiℓs\_\{i\}^\{\\ell\}after position or blockiiin layerℓ\\ell; this state includes the convolutional state and the selective\-SSM state needed to continue the forward scan without reprocessing the prefix\. Mamba\-2 recasts this family through structured state\-space duality, yielding more efficient scan kernels and a closer relationship to attention\(Dao and Gu,[2024](https://arxiv.org/html/2607.02805#bib.bib5)\)\. Recent autoregressive systems increasingly use hybrid stacks that preserve some full\-attention layers while replacing many layers with recurrent, linear\-attention, or state\-space modules: examples include Jamba’s Transformer–Mamba mixture\-of\-experts design, MiniMax\-01’s Lightning Attention and softmax\-attention mixture, Qwen3\-Next’s Gated DeltaNet/full\-attention hybrid, and Kimi Linear’s Kimi Delta Attention/MLA hybrid\(Lieber and others,[2024](https://arxiv.org/html/2607.02805#bib.bib43); Li and others,[2025](https://arxiv.org/html/2607.02805#bib.bib33); Qwen Team,[2025](https://arxiv.org/html/2607.02805#bib.bib44); Kimi Teamet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib41)\)\. These autoregressive models motivate hybridization as a practical long\-context scaling direction; our work studies the analogous question for block diffusion, where cache reuse must be aligned with the denoising objective rather than with a left\-to\-right next\-token likelihood\. DiffuMamba\-H interleaves attention and bidirectional Mamba layers for masked diffusion denoising and demonstrates that Mamba\-backed denoisers improve long\-context inference throughput\(Singhet al\.,[2026](https://arxiv.org/html/2607.02805#bib.bib13)\)\.

## 3Background

Mamba is a selective state\-space sequence mixer that replaces attention over tokens with a recurrent scan whose state can be continued across sequence segments\(Gu and Dao,[2024](https://arxiv.org/html/2607.02805#bib.bib4)\)\. Mamba\-2 improves this family with a structured state\-space duality and more efficient kernels\(Dao and Gu,[2024](https://arxiv.org/html/2607.02805#bib.bib5)\)\. In diffusion language models, DiffuMamba\-H uses a sparse hybrid schedule that interleaves attention and bidirectional Mamba layers, showing that full\-sequence diffusion denoisers can benefit from linear\-time mixers\(Singhet al\.,[2026](https://arxiv.org/html/2607.02805#bib.bib13)\)\. The present work uses the same broad hybrid motivation, but modifies the scan directionality so BDLM generation can reuse a prefix cache natively at training and inference time\. At layerℓ\\ell, we write a simplified selective state\-space update as

siℓ=Aiℓ​si−1ℓ\+Biℓ​xiℓ,yiℓ=Ciℓ​siℓ\+Dℓ​xiℓ,s\_\{i\}^\{\\ell\}=A\_\{i\}^\{\\ell\}s\_\{i\-1\}^\{\\ell\}\+B\_\{i\}^\{\\ell\}x\_\{i\}^\{\\ell\},\\qquad y\_\{i\}^\{\\ell\}=C\_\{i\}^\{\\ell\}s\_\{i\}^\{\\ell\}\+D^\{\\ell\}x\_\{i\}^\{\\ell\},\(1\)whereAiℓ,Biℓ,CiℓA\_\{i\}^\{\\ell\},B\_\{i\}^\{\\ell\},C\_\{i\}^\{\\ell\}are token\-dependent projections andsiℓs\_\{i\}^\{\\ell\}is the recurrent boundary state after token or blockii\. A prefix cache is the collection of these forward Mamba states, together with attention KV tensors for attention layers, at a block boundary\.

Masked discrete diffusion corrupts clean text by replacing tokens with masks and trains a denoiser to recover the original tokens\(Austinet al\.,[2021b](https://arxiv.org/html/2607.02805#bib.bib1); Sahooet al\.,[2024](https://arxiv.org/html/2607.02805#bib.bib2)\)\. Ifx~t\\tilde\{x\}\_\{t\}is the corrupted sequence at timesteptt, a full\-sequence denoiser optimizes

ℒfull​\(θ\)=𝔼x,t,Mt​\[1\|Mt\|​∑i∈MtCE⁡\(fθ​\(x~t,t\)i,xi\)\]\.\\mathcal\{L\}\_\{\\mathrm\{full\}\}\(\\theta\)=\\mathbb\{E\}\_\{x,t,M\_\{t\}\}\\left\[\\frac\{1\}\{\|M\_\{t\}\|\}\\sum\_\{i\\in M\_\{t\}\}\\operatorname\{CE\}\(f\_\{\\theta\}\(\\tilde\{x\}\_\{t\},t\)\_\{i\},x\_\{i\}\)\\right\]\.\(2\)This objective gives each masked token bidirectional context over the denoising window\. It also makes the entire window part of the timestep\-conditioned computation\. During generation, the hidden states are tied to the current noisy sequence, so a stable prefix has no reusable cache boundary and must be reprocessed at every reverse step\.

BDLMs introduce that boundary by partitioning a sequence intoKKblocksb⋆\(1\),…,b⋆\(K\)b^\{\\star\(1\)\},\\ldots,b^\{\\star\(K\)\}\. Following the BDLM formulation\(Arriola and others,[2025](https://arxiv.org/html/2607.02805#bib.bib40)\), the likelihood is autoregressive over blocks,

logpθ\(x\)=∑k=1Klogpθ\(b⋆\(k\)\|b⋆⁣\(<k\)\),\\log p\_\{\\theta\}\(x\)=\\sum\_\{k=1\}^\{K\}\\log p\_\{\\theta\}\\\!\\left\(b^\{\\star\(k\)\}\\,\\middle\|\\,b^\{\\star\(<k\)\}\\right\),\(3\)and each block conditional is represented by a discrete diffusion process over only the current block\(Arriola and others,[2025](https://arxiv.org/html/2607.02805#bib.bib40)\)\. Applying the diffusion NELBO to each term gives the block\-summed objective

−logpθ\(x\)≤ℒBD\(x;θ\):=∑k=1Kℒ\(b⋆\(k\),b⋆⁣\(<k\);θ\)\.\-\\log p\_\{\\theta\}\(x\)\\leq\\mathcal\{L\}\_\{\\mathrm\{BD\}\}\(x;\\theta\):=\\sum\_\{k=1\}^\{K\}\\mathcal\{L\}\\\!\\left\(b^\{\\star\(k\)\},b^\{\\star\(<k\)\};\\theta\\right\)\.\(4\)For the masked continuous\-time parameterization used in BDLM\-style training, this can be written as a sum of weighted current\-block denoising losses,

ℒBD​\(x;θ\)=∑k=1K𝔼t,q​\[w​\(t\)​CE⁡\(fθ​\(bt\(k\),b⋆⁣\(<k\),t\),b⋆\(k\)\)\],w​\(t\)=−αt′1−αt,\\mathcal\{L\}\_\{\\mathrm\{BD\}\}\(x;\\theta\)=\\sum\_\{k=1\}^\{K\}\\mathbb\{E\}\_\{t,q\}\\left\[w\(t\)\\,\\operatorname\{CE\}\\\!\\left\(f\_\{\\theta\}\(b\_\{t\}^\{\(k\)\},b^\{\\star\(<k\)\},t\),b^\{\\star\(k\)\}\\right\)\\right\],\\qquad w\(t\)=\\frac\{\-\\alpha^\{\\prime\}\_\{t\}\}\{1\-\\alpha\_\{t\}\},\(5\)wherebt\(k\)∼qt\(⋅∣b⋆\(k\)\)b\_\{t\}^\{\(k\)\}\\sim q\_\{t\}\(\\cdot\\mid b^\{\\star\(k\)\}\)is the noised current block,b⋆⁣\(<k\)b^\{\\star\(<k\)\}is a clean prefix, and future blocks are excluded\. The all\-block training objective evaluates this sum over block boundaries within the same training window, so every block is trained as a current denoising target conditioned on its clean prefix\.

This factorization enables caching for block diffusion\. For an attention BDLM, the denoiser for blockkkcan be written with the same cache interface used by autoregressive Transformers:

z\(k\),Kk,Vk←fθ\(k\)​\(bt\(k\),K1:k−1,V1:k−1\)≡fθ\(k\)​\(bt\(k\),b⋆⁣\(<k\)\),z^\{\(k\)\},K\_\{k\},V\_\{k\}\\leftarrow f\_\{\\theta\}^\{\(k\)\}\\\!\\left\(b\_\{t\}^\{\(k\)\},K\_\{1:k\-1\},V\_\{1:k\-1\}\\right\)\\equiv f\_\{\\theta\}^\{\(k\)\}\\\!\\left\(b\_\{t\}^\{\(k\)\},b^\{\\star\(<k\)\}\\right\),\(6\)wherez\(k\)z^\{\(k\)\}are clean\-token logits for the current block andK1:k−1,V1:k−1K\_\{1:k\-1\},V\_\{1:k\-1\}are key/value tensors computed from completed prefix blocks\. During sampling, the model denoises all tokens within the current block in parallel while reading this prefix cache; after the block is finalized, a clean forward pass over the new block appendsKk,VkK\_\{k\},V\_\{k\}for subsequent blocks\. Thus cache reuse follows from the block conditional structure itself: completed blocks are fixed inputs to later block denoisers, while reverse diffusion updates are confined to the active block\.

#### NELBO perplexity\.

For diffusion language models, exactlog⁡pθ​\(x\)\\log p\_\{\\theta\}\(x\)is generally intractable because likelihood marginalizes over reverse denoising trajectories\. We therefore evaluate a negative evidence lower bound \(NELBO\)\. For a validation corpus with token countsnjn\_\{j\}and sequence\-level NELBO estimatesℒ^NELBO​\(xj\)\\widehat\{\\mathcal\{L\}\}\_\{\\mathrm\{NELBO\}\}\(x\_\{j\}\), the reported token\-weighted perplexity is

PPLNELBO=exp⁡\(∑jℒ^NELBO​\(xj\)∑jnj\)\.\\operatorname\{PPL\}\_\{\\mathrm\{NELBO\}\}=\\exp\\\!\\left\(\\frac\{\\sum\_\{j\}\\widehat\{\\mathcal\{L\}\}\_\{\\mathrm\{NELBO\}\}\(x\_\{j\}\)\}\{\\sum\_\{j\}n\_\{j\}\}\\right\)\.\(7\)This is an upper\-bound\-derived perplexity metric rather than an exact autoregressive likelihood\. It lets full\-sequence MDLMs and BDLMs be compared under the same sequence\-token\-weighted diffusion validation protocol\.

#### Complete\(d\)P learning\-rate transfer\.

For scale\-up, we transfer learning rates from 87M to 350M using a Complete\(d\)P\-style rule\(Mlodozeniecet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib14)\)\. The motivation is that a larger run changes width, depth, batch size, and duration simultaneously\. We use the hidden\-weight transfer form

ηlarge=ηsmall​rwidthaw​rdepthad​rbatchab​rdurationaT,\\eta\_\{\\mathrm\{large\}\}=\\eta\_\{\\mathrm\{small\}\}\\,r\_\{\\mathrm\{width\}\}^\{a\_\{w\}\}r\_\{\\mathrm\{depth\}\}^\{a\_\{d\}\}r\_\{\\mathrm\{batch\}\}^\{a\_\{b\}\}r\_\{\\mathrm\{duration\}\}^\{a\_\{T\}\},\(8\)where eachrris the large\-to\-small ratio for that axis and the exponents are the Complete\(d\)P transfer exponents for hidden weights\. This gives scale factors of0\.310\.31for BDLM attention and0\.330\.33for BDLM Mamba\-H in our 87M to 350M transfer; the resulting learning rates are reported in Table[2](https://arxiv.org/html/2607.02805#S5.T2)\.

## 4Block Diffusion Language Model

We usec\(k\)=b⋆⁣\(<k\)c^\{\(k\)\}=b^\{\\star\(<k\)\}for the clean prefix before blockkk,b~t\(k\)\\tilde\{b\}^\{\(k\)\}\_\{t\}for the noised active block at reverse timesteptt, andfθf\_\{\\theta\}for the block denoiser\. For Mamba layerℓ\\ell,skℓs\_\{k\}^\{\\ell\}denotes the reusable forward recurrent boundary state produced by scanning the clean prefix\. The symbols→\\rightarrowand←\\leftarrowdenote forward and reverse Mamba scan directions, whileyk,t→,ℓy\_\{k,t\}^\{\\rightarrow,\\ell\}andyk,t←,ℓy\_\{k,t\}^\{\\leftarrow,\\ell\}are the corresponding active\-block hidden states\.

Let𝒞ℓ​\(c\)\\mathcal\{C\}^\{\\ell\}\(c\)denote the layer\-ℓ\\ellprefix cache\. Prefix\-cache reuse requires this cache to be a function of the clean prefix alone, so𝒞ℓ​\(c\)​must be independent of​t\\mathcal\{C\}^\{\\ell\}\(c\)\\ \\text\{must be independent of\}\\ t\.

Naively reusing a diffusion denoiser can violate this requirement because diffusion Transformers often inject timestep information through adaptive normalization applied to every token\(Peebles and Xie,[2023](https://arxiv.org/html/2607.02805#bib.bib7)\)\. If the same modulation touches prefix tokens, the cache becomes timestep\-specific and must be bespoke for every reverse step\. We therefore factor each layer into a timestep\-free prefix operator and a timestep\-conditioned active\-block continuation:

𝒞ℓ​\(c\)=Pℓ​\(c\),hb,tℓ\+1=Bℓ​\(hb,tℓ;𝒞ℓ​\(c\),et\)\.\\mathcal\{C\}^\{\\ell\}\(c\)=P^\{\\ell\}\(c\),\\qquad h\_\{b,t\}^\{\\ell\+1\}=B^\{\\ell\}\(h\_\{b,t\}^\{\\ell\};\\mathcal\{C\}^\{\\ell\}\(c\),e\_\{t\}\)\.\(9\)Appendix[E](https://arxiv.org/html/2607.02805#A5)gives the corresponding AdaLN factorization\.

Mamba directionality adds a second constraint\. A fully bidirectional Mamba denoiser over\[c;b~t\]\[c;\\tilde\{b\}\_\{t\}\]combines a forward scan with a reverse scan over the reversed prefix\-plus\-block sequence\. The reverse prefix states then depend on the active block, so they are not prefix\-only values and cannot be reused as a BDLM prefix cache\. Our BDLM Mamba hybrid instead runs the reverse\-direction Mamba only within the current denoising block, rather than over the full sequence as in DiffuMamba \(we denote this as a*block partial\-reverse*architecture\)\. BDLM Mamba caches only the forward boundary states,

skℓ=Fθ,→ℓ​\(c\(k\)\),s\_\{k\}^\{\\ell\}=F\_\{\\theta,\\rightarrow\}^\{\\ell\}\(c^\{\(k\)\}\),\(10\)and continues them into the active block:

yk,t→,ℓ\\displaystyle y\_\{k,t\}^\{\\rightarrow,\\ell\}=Fθ,→ℓ​\(b~t\(k\);skℓ\),\\displaystyle=F\_\{\\theta,\\rightarrow\}^\{\\ell\}\(\\tilde\{b\}\_\{t\}^\{\(k\)\};s\_\{k\}^\{\\ell\}\),\(11\)yk,t←,ℓ\\displaystyle y\_\{k,t\}^\{\\leftarrow,\\ell\}=rev⁡\(Fθ,←ℓ​\(rev⁡\(b~t\(k\)\)\)\),\\displaystyle=\\operatorname\{rev\}\\left\(F\_\{\\theta,\\leftarrow\}^\{\\ell\}\(\\operatorname\{rev\}\(\\tilde\{b\}\_\{t\}^\{\(k\)\}\)\)\\right\),\(12\)yk,tℓ\\displaystyle y\_\{k,t\}^\{\\ell\}=yk,t→,ℓ\+yk,t←,ℓ\.\\displaystyle=y\_\{k,t\}^\{\\rightarrow,\\ell\}\+y\_\{k,t\}^\{\\leftarrow,\\ell\}\.\(13\)The active block still receives both left\-to\-right and right\-to\-left mixing, but only within the current block\. The reusable prefix cache remains valid because the reverse scan is local to the active block\. This serves as a native BDLM analogue of generation with block\-caching: generated prefix blocks are summarized by forward recurrent state, while denoising inside the current block keeps local bidirectionality\.

The prefix\-cache objective trains the same forward cached representations which will be used at inference:

ℒprefix​\(θ\)=1K​∑k𝔼t​ℓθ​\(b⋆\(k\)∣sk1:L,b~t\(k\),t\)\.\\mathcal\{L\}\_\{\\mathrm\{prefix\}\}\(\\theta\)=\\frac\{1\}\{K\}\\sum\_\{k\}\\mathbb\{E\}\_\{t\}\\ell\_\{\\theta\}\(b^\{\\star\(k\)\}\\mid s\_\{k\}^\{1:L\},\\tilde\{b\}\_\{t\}^\{\(k\)\},t\)\.\(14\)For attention layers, the corresponding cache object is the usual prefix key/value tensor\. For Mamba layers, the cache object is the recurrent boundary state, including convolution and selective\-SSM state needed to continue the scan into the active block\.

### 4\.1Block Diffusion Mamba Hybrid Training

Consider a sequence split into blocksb1,b2,b3,b4b\_\{1\},b\_\{2\},b\_\{3\},b\_\{4\}, and letb4b\_\{4\}be the active diffusion block\. The model first scans the clean prefix to produce the Mamba prefix cache

C4=Fθ​\(b1,b2,b3\),C\_\{4\}=F\_\{\\theta\}\(b\_\{1\},b\_\{2\},b\_\{3\}\),\(15\)whereC4C\_\{4\}contains the convolution and state\-space states at the boundary beforeb4b\_\{4\}\. The active block is then denoised from that same cache object:

ℓ4=CE⁡\(Dθ​\(C4,b~4,t\),b4\)\.\\ell\_\{4\}=\\operatorname\{CE\}\\\!\\left\(D\_\{\\theta\}\(C\_\{4\},\\tilde\{b\}\_\{4\},t\),b\_\{4\}\\right\)\.\(16\)Across all\-block training, the same construction is evaluated for every block boundary in parallel\. Backpropagation follows the prefix\-cache construction path:

ℓ4\\displaystyle\\ell\_\{4\}→C4→Fθ​\(b1,b2,b3\)→θ\.\\displaystyle\\rightarrow C\_\{4\}\\rightarrow F\_\{\\theta\}\(b\_\{1\},b\_\{2\},b\_\{3\}\)\\rightarrow\\theta\.\(17\)The cache still contains the same prefix\-only Mamba states that are reused at inference\. Downstream block losses shape the representations stored in those states, aligning the training computation with the sampler’s cache interface: the state consumed by later blocks is produced by the same prefix scan that is used during native BDLM generation\.

## 5Experiments

### 5\.1Experiment Setup

We evaluate training and inference throughput as well as validation performance of our BDLM Mamba hybrid against full\-sequence attention, full\-sequence DiffuMamba\-H, and BDLM attention baselines\. The 87M sweep trains all four architectures on the following learning rates:\{5×10−4,10−3,2×10−3,4×10−3,8×10−3\}\\\{5\{\\times\}10^\{\-4\},\\,10^\{\-3\},\\,2\{\\times\}10^\{\-3\},\\,4\{\\times\}10^\{\-3\},\\,8\{\\times\}10^\{\-3\}\\\}\. All configurations use 5 billion tokens sampled from DCLM\(Liet al\.,[2024](https://arxiv.org/html/2607.02805#bib.bib21)\), 8192\-token training sequences, 16 diffusion steps, and no\-timestep modulation; BDLM configurations use block size 256\. The full\-sequence DiffuMamba\-H model uses a fully bidirectional Mamba hybrid denoiser over the denoising window, whereas BDLM Mamba\-H only runs the reverse Mamba on the current active denoising block to leverage prefix\-caching construction as described in Section[4\.1](https://arxiv.org/html/2607.02805#S4.SS1)\.

Likelihood validation uses the same BDLM/MDLM\-compatible diffusion\-NELBO evaluator for all models at 8192\-token context, with one Monte Carlo timestep/mask corruption sample per validation sequence\. We report C4\-en and Paloma\-C4 NELBO perplexity and bits per byte \(BPB\)\. NELBO perplexity is computed as the exponential of the token\-weighted average NELBO estimate, and BPB divides the same token\-weighted negative log\-likelihood in bits by the number of UTF\-8 source\-document bytes\. We report MCQA separately as pseudo\-likelihood answer\-choice accuracy over fixed candidate answers\. Masked\-token accuracy from the diffusion\-NELBO evaluator is included in the full 87M sweep in Table[4](https://arxiv.org/html/2607.02805#A1.T4)\.

For MCQA, we evaluate HellaSwag, PIQA, ARC\-Easy, ARC\-Challenge, BoolQ, and WinoGrande\(Zellerset al\.,[2019](https://arxiv.org/html/2607.02805#bib.bib25); Bisket al\.,[2020](https://arxiv.org/html/2607.02805#bib.bib26); Clarket al\.,[2018](https://arxiv.org/html/2607.02805#bib.bib27);[2019](https://arxiv.org/html/2607.02805#bib.bib28); Sakaguchiet al\.,[2021](https://arxiv.org/html/2607.02805#bib.bib29)\)and report macro\-average accuracy across tasks\. For each example, we score each provided answer choice under the model’s pseudo\-likelihood scoring rule and select the highest\-scoring choice\. We use fixed\-choice evaluation over evaluations that rely on free\-form generation and answer extraction, making it suitable for comparing base models with no post training \. The suite includes binary\-choice tasks \(PIQA, BoolQ, and WinoGrande\) and multi\-choice tasks with fixed or variable answer sets \(HellaSwag, ARC\-Easy, and ARC\-Challenge\)\.

Appendix[D](https://arxiv.org/html/2607.02805#A4)compares no\-timestep and timestep\-conditioned objectives\. We use the no\-timestep objective in the main experiments for two reasons\. First, it matches the original BDLM experimental convention\(Arriola and others,[2025](https://arxiv.org/html/2607.02805#bib.bib40)\)\. Second, it keeps clean\-prefix caches timestep\-invariant: if timestep modulation is applied to prefix tokens, the cached prefix states depend on the reverse diffusion step and cannot be reused across denoising steps\. Even when a timestep is held fixed, a technically cache\-aligned BDLM implementation should avoid learned timestep embeddings in the prefix\-cache construction, because completed prefix states are properties of clean prefix text rather than of a particular reverse step\. The AdaLN factorization in Appendix[E](https://arxiv.org/html/2607.02805#A5)describes how BDLM Mamba\-H can make use of prefix caching; timestep modulation is applied only to the active denoising block, while the clean\-prefix scan remains timestep\-free\. Empirically, the fixed learning rate comparison in Appendix[D](https://arxiv.org/html/2607.02805#A4)shows that timestep conditioning improves BDLM attention, but worsens BDLM Mamba\-H and both full\-sequence baselines on NELBO perplexity\. We therefore use the no\-timestep objective for the main model\-selection and scale\-up experiments\.

For the 350M scale\-up experiments, we transfer learning rates from the selected 87M configurations using a Complete\(d\)P\-style hidden\-weight rule\(Mlodozeniecet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib14)\)\. We use this rule as a principled way to choose scale\-up candidates because the 350M runs change width, depth, global batch size, and training horizon simultaneously\. Recent diffusion\-language\-model scaling studies that find learning rate and batch size transfer rules remain useful for DLMs\(Niet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib15)\), and with work using CompleteP for stable learning\-rate transfer across width and depth in discrete DLM scaling\(von Rütteet al\.,[2026](https://arxiv.org/html/2607.02805#bib.bib16)\)\. The resulting learning rates are shown in Table[2](https://arxiv.org/html/2607.02805#S5.T2)\. We select the best 87M models by C4\-en heldout perplexity:4×10−34\{\\times\}10^\{\-3\}for BDLM attention and8×10−38\{\\times\}10^\{\-3\}for BDLM Mamba\-H\. Both 350M models use the same warmup length, global batch, sequence length, data stream, and no\-timestep objective\. We choose model\-specific optimizer\-step counts so each run trains just above the same target token\-to\-parameter ratio, then vary only the architecture and Complete\(d\)P\-transferred learning rate\.

### 5\.2Language Modeling Quality

We select one checkpoint per architecture by the best C4\-enPPLNELBO\\mathrm\{PPL\}\_\{\\mathrm\{NELBO\}\}in the 87M\-class learning\-rate sweep\. Under this selection rule, BDLM Mamba\-H gives the lowest likelihood\-bound validation perplexity: 61\.6, compared with 76\.5 for BDLM attention, 83\.7 for full\-sequence DiffuMamba\-H, and 87\.6 for full\-sequence attention; the full 20\-row sweep is present in Appendix Table[4](https://arxiv.org/html/2607.02805#A1.T4)\. MCQA provides a complementary fixed\-choice preference evaluation separate to the NELBO\-perplexity ranking\.

At 350M scale, the two selected BDLM models remain comparable with BDLM attention obtaining 37\.0 C4\-enPPLNELBO\\mathrm\{PPL\}\_\{\\mathrm\{NELBO\}\}and 36\.1 Paloma\-C4PPLNELBO\\mathrm\{PPL\}\_\{\\mathrm\{NELBO\}\}, while BDLM Mamba\-H obtains 38\.2 and 36\.8, respectively\. The corresponding MCQA accuracies are 41\.3 for BDLM attention and 43\.3 for BDLM Mamba\-H\. Thus we show BDLM Mamba\-H remains in the same validation\-quality range while enabling native long\-context cached inference\.

Table 1:Best 87M\-class DCLM\(Liet al\.,[2024](https://arxiv.org/html/2607.02805#bib.bib21)\)no\-timestep validation configurations, selected by C4\-en NELBO perplexity\. PPL denotes the BDLM/MDLM\-compatible diffusion\-NELBO validation path with one timestep sample at 8192\-token context\. All models are trained on 5\.00B DCLM tokens\. BPB divides token\-weighted negative log\-likelihood in bits by the number of UTF\-8 bytes in the source documents\. MCQA is pseudo\-likelihood accuracy over fixed answer choices\. The full learning\-rate sweep is reported in Table[4](https://arxiv.org/html/2607.02805#A1.T4)\.Table 2:350M BDLM scale\-up configuration\. Learning rates are transferred from the best 87M C4\-en validation settings using a Complete\(d\)P\-style hidden\-weight transfer rule\(Mlodozeniecet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib14)\)\. Both models use no timestep conditioning, DCLM, the BDLM\(Arriola and others,[2025](https://arxiv.org/html/2607.02805#bib.bib40)\)objective, 8192\-token training sequences, global batch 64, 2000 warmup steps, the same optimizer\-step horizon, and 17\.65B DCLM training tokens\.Table 3:350M BDLM no\-timestep validation results\. C4 and Paloma C4 are evaluated at 8192\-token context with 256 examples per dataset\. BPB divides token\-weighted negative log\-likelihood in bits by the number of UTF\-8 bytes in the source documents\. MCQA is pseudo\-likelihood accuracy over fixed answer choices\.
### 5\.3Training and Inference Throughput

Training throughput is measured at 350M scale on 8x A100\-80GB with synthetic 8192\-token batches\. Inference throughput is measured with 700M random\-initialized models on 1x A100\-80GB with BF16 and batch size 1\. Both benchmarks use 5 warmup iterations and average over 15 measured iterations\. The context generation lengths follow the DiffuMamba long\-context protocol: full\-sequence denoisers are measured through the 65K\-token regime, while native BDLM models are extended along the block\-cached ladder up to 262k generated tokens\.

At 350M scale, all four architectures use the target global batch on 8x A100\-80GB and fit into memory without activation checkpointing\. The full\-sequence models are significantly faster in 8192\-token training throughput, reflecting the additional cost of evaluating the all\-block BDLM objective during training\. Among the BDLM models, BDLM Mamba\-H trains faster than BDLM attention and uses less peak reserved memory\. At inference time, the native BDLM models retain substantially higher throughput than full\-sequence denoisers as generation length grows\. BDLM attention is faster at short lengths in our implementation, while BDLM Mamba\-H overtakes it in the long\-context regime starting at approximately 16k generated tokens\.

Table 5:350M 8x A100\-80GB training throughput at 8192\-token context\. All models use synthetic batches, global batch 64, BF16, ZeRO\-2, tiled linear cross\-entropy, and FlashAttention kernels when available\.At 65K, BDLM Mamba\-H is 19\.7x faster than the full\-sequence DiffuMamba\-H baseline at the same length\. On the extended BDLM ladder, BDLM Mamba\-H is 3\.7x faster than BDLM attention at 262K\. This crossover is qualitatively consistent with DiffuMamba’s finding that Mamba\-backed block generation increasingly separates from attention\-backed block generation as sequence length grows\(Singhet al\.,[2026](https://arxiv.org/html/2607.02805#bib.bib13)\)\. Raw tokens/s, implementation details, and length\-by\-length speedup ratios are given in Appendix[B](https://arxiv.org/html/2607.02805#A2)\.

![Refer to caption](https://arxiv.org/html/2607.02805v1/figures/inference_throughput_350m.png)

\(a\) Throughput to 65K tokens\.

![Refer to caption](https://arxiv.org/html/2607.02805v1/figures/inference_throughput_350m_extended.png)

\(b\) BDLM throughput on the extended ladder\.

Figure 2:700M random\-initialized single\-A100 inference throughput\. Full\-sequence baselines use repeated full\-sequence denoising with CUDA graphs and are reported\. Native BDLM models use block size 256, BF16, cacheable prefix states, and optimized cached denoising; the extended ladder includes completed BDLM attention and BDLM Mamba\-H points through 262K generated tokens\.

## 6Limitations

While the block objective improves long\-context generation, we find that at 350M parameters on 8x A100\-80GB, BDLM attention trains 33\.9% slower than the full\-sequence attention baseline, and BDLM Mamba\-H trains 25\.5% slower than the full\-sequence DiffuMamba\-H baseline\. We leave it to future work to investigate optimal scaling laws for hyperparameter transfer for dLLMs and hybrid variants thereof\. Complete\(d\)P provided a parameterization\-based method of transferring learning rates across width, depth, batch, and duration\(Mlodozeniecet al\.,[2025](https://arxiv.org/html/2607.02805#bib.bib14)\), but did not explicitly derive these laws for dLLMs\. Masked and block diffusion objectives have different timestep sampling, loss weighting, and caching constraints, so the transferred 350M learning rates are strong candidates rather than a proof of optimality\.

## 7Conclusion

We presented a training recipe for block diffusion Mamba hybrids that trains the same prefix cache used during native BDLM generation\. The method keeps the prefix scan timestep\-free to optionally support valid AdaLN timestep conditioning for BDLM Mamba, confines reverse\-direction Mamba mixing to the active block rather than applying a full\-window reverse pass, and trains the prefix cache through downstream block losses\.

Our trained models preserve validation quality while improving the long\-context systems regime\. At 87M parameters, BDLM Mamba\-H reaches the best C4\-en PPL in the DCLM sweep, and at 350M parameters its PPL remains in the same range as BDLM attention\. In inference, the native BDLM models extend to lengths that full\-sequence denoisers do not practically cover, and BDLM Mamba\-H increasingly separates at long context, reaching 19\.7x the full\-sequence DiffuMamba\-H counterpart at 65K tokens and 3\.7x BDLM attention throughput at 262K\. We show that training block partial\-reverse BDLM Mamba hybrids yields a practical long\-context diffusion architecture worth further exploring\.

## Acknowledgements

We gratefully acknowledge Modal Labs for generously supporting this work with a Modal for Academics compute grant\.

## References

- Block diffusion: interpolating between autoregressive and diffusion language models\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2503.09573)Cited by:[Table 9](https://arxiv.org/html/2607.02805#A3.T9.15.15.18.2.2),[Appendix D](https://arxiv.org/html/2607.02805#A4.p2.1),[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1),[§3](https://arxiv.org/html/2607.02805#S3.p3.2),[§3](https://arxiv.org/html/2607.02805#S3.p3.5),[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p4.1),[Table 2](https://arxiv.org/html/2607.02805#S5.T2)\.
- J\. Austin, D\. D\. Johnson, J\. Ho, D\. Tarlow, and R\. van den Berg \(2021a\)Structured denoising diffusion models in discrete state\-spaces\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Austin, D\. D\. Johnson, J\. Ho, D\. Tarlow, and R\. van den Berg \(2021b\)Structured denoising diffusion models in discrete state\-spaces\.InAdvances in Neural Information Processing Systems,Cited by:[§3](https://arxiv.org/html/2607.02805#S3.p2.2)\.
- Y\. Bisk, R\. Zellers, J\. Gao, and Y\. Choi \(2020\)PIQA: reasoning about physical commonsense in natural language\.InAAAI Conference on Artificial Intelligence,Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p3.1)\.
- C\. Clark, K\. Lee, M\. Chang, T\. Kwiatkowski, M\. Collins, and K\. Toutanova \(2019\)BoolQ: exploring the surprising difficulty of natural yes/no questions\.InConference of the North American Chapter of the Association for Computational Linguistics,Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p3.1)\.
- P\. Clark, I\. Cowhey, O\. Etzioni, T\. Khot, A\. Sabharwal, C\. Schoenick, and O\. Tafjord \(2018\)Think you have solved question answering? try ARC, the AI2 reasoning challenge\.arXiv preprint arXiv:1803\.05457\.External Links:[Link](https://arxiv.org/abs/1803.05457)Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p3.1)\.
- T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. Ré \(2022\)FlashAttention: fast and memory\-efficient exact attention with io\-awareness\.InAdvances in Neural Information Processing Systems,Cited by:[Appendix B](https://arxiv.org/html/2607.02805#A2.p1.1)\.
- T\. Dao and A\. Gu \(2024\)Transformers are SSMs: generalized models and efficient algorithms through structured state space duality\.InInternational Conference on Machine Learning,External Links:[Link](https://arxiv.org/abs/2405.21060)Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px2.p1.3),[§3](https://arxiv.org/html/2607.02805#S3.p1.1)\.
- A\. Gu and T\. Dao \(2024\)Mamba: linear\-time sequence modeling with selective state spaces\.InConference on Language Modeling,External Links:[Link](https://openreview.net/forum?id=tEYskw1VY2)Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px2.p1.3),[§3](https://arxiv.org/html/2607.02805#S3.p1.1)\.
- E\. Hoogeboom, D\. Nielsen, P\. Jaini, P\. Forré, and M\. Welling \(2021\)Argmax flows and multinomial diffusion: learning categorical distributions\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- Kimi Team, Y\. Zhang, Z\. Lin, X\. Yao, J\. Hu, F\. Meng, C\. Liu, X\. Men, S\. Yang, Z\. Li, W\. Li, E\. Lu, W\. Liu, Y\. Chen, W\. Xu, L\. Yu, Y\. Wang, Y\. Fan, L\. Zhong, E\. Yuan, D\. Zhang, Y\. Zhang, T\. Y\. Liu, H\. Wang, S\. Fang, W\. He, S\. Liu, Y\. Li, J\. Su, J\. Qiu, B\. Pang, J\. Yan, Z\. Jiang, W\. Huang, B\. Yin, J\. You, C\. Wei, Z\. Wang, C\. Hong, Y\. Chen, G\. Chen, Y\. Wang, H\. Zheng, F\. Wang, Y\. Liu, M\. Dong, Z\. Zhang, S\. Pan, W\. Wu, Y\. Wu, L\. Guan, J\. Tao, G\. Fu, X\. Xu, Y\. Wang, G\. Lai, Y\. Wu, X\. Zhou, Z\. Yang, and Y\. Du \(2025\)Kimi linear: an expressive, efficient attention architecture\.Technical report\.External Links:[Link](https://github.com/MoonshotAI/Kimi-Linear)Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px2.p1.3)\.
- A\. Liet al\.\(2025\)MiniMax\-01: scaling foundation models with lightning attention\.Technical report\.External Links:[Link](https://github.com/MiniMax-AI/MiniMax-01)Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px2.p1.3)\.
- J\. Li, A\. Fang, G\. Smyrnis, M\. Ivgi, M\. Jordan, S\. Y\. Gadre, H\. Bansal, E\. Guha, S\. Keh, K\. Arora, S\. Garg, R\. Xin, N\. Muennighoff, R\. Heckel, J\. Mercat, M\. Chen, S\. Gururangan, M\. Wortsman, A\. Albalak, Y\. Bitton, M\. Nezhurina, A\. Abbas, C\. Hsieh, D\. Ghosh, J\. Gardner, M\. Kilian, H\. Zhang, R\. Shao, S\. Pratt, S\. Sanyal, G\. Ilharco, G\. Daras, K\. Marathe, A\. Gokaslan, J\. Zhang, K\. Chandu, T\. Nguyen, I\. Vasiljevic, S\. Kakade, S\. Song, S\. Sanghavi, F\. Faghri, S\. Oh, L\. Zettlemoyer, K\. Lo, A\. El\-Nouby, H\. Pouransari, A\. Toshev, S\. Wang, D\. Groeneveld, L\. Soldaini, P\. W\. Koh, J\. Jitsev, T\. Kollar, A\. G\. Dimakis, Y\. Carmon, A\. Dave, L\. Schmidt, and V\. Shankar \(2024\)DataComp\-LM: in search of the next generation of training sets for language models\.InAdvances in Neural Information Processing Systems Datasets and Benchmarks Track,External Links:[Link](https://openreview.net/forum?id=CNWdWn47IE)Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p1.1),[Table 1](https://arxiv.org/html/2607.02805#S5.T1)\.
- O\. Lieberet al\.\(2024\)Jamba: a hybrid transformer\-mamba language model\.OpenReview\.External Links:[Link](https://openreview.net/forum?id=JFPaD7lpBD)Cited by:[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px2.p1.3)\.
- Ling Team, B\. Han, C\. Tang, C\. Liang, D\. Zhang, F\. Yuan, F\. Zhu, J\. Gao, J\. Hu, L\. Li, M\. Li, M\. Zhang, P\. Jiang, P\. Jiao, Q\. Zhao, Q\. Yang, W\. Shen, X\. Yang, Y\. Zhang, Y\. Ren, Y\. Zhao, Y\. Cao, Y\. Sun, Y\. Zhang, Y\. Fang, Z\. Lin, Z\. Cheng, and J\. Zhou \(2025\)Every attention matters: an efficient hybrid architecture for long\-context reasoning\.arXiv preprint arXiv:2510\.19338\.Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1)\.
- Z\. Liu, Y\. Yang, Y\. Zhang, J\. Chen, C\. Zou, Q\. Wei, S\. Wang, and L\. Zhang \(2025\)dLLM\-Cache: accelerating diffusion large language models with adaptive caching\.arXiv preprint arXiv:2506\.06295\.External Links:[Link](https://arxiv.org/abs/2506.06295)Cited by:[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- X\. Ma, R\. Yu, G\. Fang, and X\. Wang \(2025\)dKV\-Cache: the cache for diffusion language models\.arXiv preprint arXiv:2505\.15781\.External Links:[Link](https://arxiv.org/abs/2505.15781)Cited by:[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- B\. Mlodozeniec, P\. Ablin, L\. Béthune, D\. Busbridge, M\. Klein, J\. Ramapuram, and M\. Cuturi \(2025\)Completed hyperparameter transfer across modules, width, depth, batch and duration\.arXiv preprint arXiv:2512\.22382\.External Links:2512\.22382,[Document](https://dx.doi.org/10.48550/arXiv.2512.22382),[Link](https://arxiv.org/abs/2512.22382)Cited by:[2nd item](https://arxiv.org/html/2607.02805#S1.I1.i2.p1.1),[§3](https://arxiv.org/html/2607.02805#S3.SS0.SSS0.Px2.p1.4),[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p5.2),[Table 2](https://arxiv.org/html/2607.02805#S5.T2),[§6](https://arxiv.org/html/2607.02805#S6.p1.1)\.
- Q\. Nguyen\-Tri, M\. Ranjan, and Z\. Shen \(2025\)Attention is all you need for KV cache in diffusion LLMs\.arXiv preprint\.Cited by:[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Ni, Q\. Liu, C\. Du, L\. Dou, H\. Yan, Z\. Wang, T\. Pang, and M\. Q\. Shieh \(2025\)Training optimal large diffusion language models\.arXiv preprint arXiv:2510\.03280\.External Links:[Link](https://arxiv.org/abs/2510.03280)Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p5.2)\.
- S\. Nie, F\. Zhu, Z\. You, X\. Zhang, J\. Ou, J\. Hu, J\. Zhou, Y\. Lin, J\. Wen, and C\. Li \(2025\)Large language diffusion models\.InAdvances in Neural Information Processing Systems,External Links:[Link](https://openreview.net/forum?id=KnqiC0znVF)Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- W\. Peebles and S\. Xie \(2023\)Scalable diffusion models with transformers\.InInternational Conference on Computer Vision,External Links:[Link](https://arxiv.org/abs/2212.09748)Cited by:[§4](https://arxiv.org/html/2607.02805#S4.p3.1)\.
- Z\. Qin, W\. Sun, D\. Li, X\. Shen, W\. Sun, and Y\. Zhong \(2024\)Various lengths, constant speed: efficient language modeling with lightning attention\.InProceedings of the 41st International Conference on Machine Learning,External Links:[Link](https://proceedings.mlr.press/v235/qin24c.html)Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1)\.
- Qwen Team \(2025\)Qwen3\-next: towards ultimate training efficiency\.Note:Qwen technical blogExternal Links:[Link](https://qwenlm.github.io/blog/qwen3-next/)Cited by:[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px2.p1.3)\.
- S\. S\. Sahoo, M\. Arriola, A\. Gokaslan, J\. T\. Chiu, A\. Rush, and V\. Kuleshov \(2024\)Simple and effective masked diffusion language models\.InAdvances in Neural Information Processing Systems,Cited by:[§3](https://arxiv.org/html/2607.02805#S3.p2.2)\.
- S\. S\. Sahooet al\.\(2024\)Simple and effective masked diffusion language models\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p1.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- K\. Sakaguchi, R\. L\. Bras, C\. Bhagavatula, and Y\. Choi \(2021\)WinoGrande: an adversarial winograd schema challenge at scale\.InAAAI Conference on Artificial Intelligence,Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p3.1)\.
- V\. Singh, O\. Ostapenko, P\. Noel, E\. Belilovsky, and T\. Scholak \(2026\)DiffuMamba: high\-throughput diffusion LMs with mamba backbone\.InInternational Conference on Machine Learning,External Links:[Link](https://arxiv.org/abs/2511.15927)Cited by:[§1](https://arxiv.org/html/2607.02805#S1.p2.1),[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px2.p1.3),[§3](https://arxiv.org/html/2607.02805#S3.p1.1),[§5\.3](https://arxiv.org/html/2607.02805#S5.SS3.p3.1)\.
- D\. von Rütte, J\. Fluri, O\. Pooladzandi, B\. Schölkopf, T\. Hofmann, and A\. Orvieto \(2026\)Scaling behavior of discrete diffusion language models\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=GDYaNzxt9T)Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p5.2)\.
- X\. Wang, C\. Xu, Y\. Jin, J\. Jin, H\. Zhang, and Z\. Deng \(2025\)Diffusion LLMs can do faster\-than\-AR inference via discrete diffusion forcing\.arXiv preprint arXiv:2508\.09192\.External Links:[Link](https://arxiv.org/abs/2508.09192)Cited by:[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- C\. Wu, H\. Zhang, S\. Xue, Z\. Liu, S\. Diao, L\. Zhu, P\. Luo, S\. Han, and E\. Xie \(2026\)Fast\-dLLM: training\-free acceleration of diffusion LLM by enabling KV cache and parallel decoding\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=3Z3Is6hnOT)Cited by:[§2](https://arxiv.org/html/2607.02805#S2.SS0.SSS0.Px1.p1.1)\.
- R\. Zellers, A\. Holtzman, Y\. Bisk, A\. Farhadi, and Y\. Choi \(2019\)HellaSwag: can a machine really finish your sentence?\.InAnnual Meeting of the Association for Computational Linguistics,Cited by:[§5\.1](https://arxiv.org/html/2607.02805#S5.SS1.p3.1)\.

## Appendix AFull 87M Validation Sweep

Table 4:Full 87M DCLM no\-timestep validation sweep\. All checkpoints are evaluated at step 19,074 with the same 8192\-token diffusion\-NELBO validation path used in Table[1](https://arxiv.org/html/2607.02805#S5.T1)\.
## Appendix BInference Throughput Protocol and Raw Results

All inference models use the same block size, denoising\-step ratio, precision, and measurement policy\. Full\-sequence baselines use fixed\-shape CUDA graphs for repeated full\-window denoising and FlashAttention kernels where available\(Daoet al\.,[2022](https://arxiv.org/html/2607.02805#bib.bib6)\)\. BDLM attention uses preallocated mutable prefix KV caches, tiled key\-cache append, BDLM\-style FlashAttention for prefix\-cache and active\-block attention, graph\-captured block denoising steps, and tiled LM\-head evaluation\. BDLM Mamba\-H uses the block partial\-reverse Mamba construction: native Mamba kernels, chunked prefix prefill, read\-only recurrent cache assumptions, compiled fixed\-shape Mamba block continuations with CUDA graphs enabled, graph\-captured block denoising steps, and the same chunked LM\-head path\.

Table[6](https://arxiv.org/html/2607.02805#A2.T6)gives the raw inference throughput numbers used in Figure[2](https://arxiv.org/html/2607.02805#S5.F2)\. In our implementation, BDLM attention is faster at short lengths, where the KV cache is still small and softmax attention remains efficient\. BDLM Mamba\-H overtakes BDLM attention at longer lengths: 1\.9x at 32K, 2\.6x at 65K, 3\.2x at 131K, 4\.0x at 200K, and 3\.7x at 262K relative to BDLM attention\. These absolute ratios are specific to the 700M random\-initialized A100 benchmark used here\.

Table 6:Raw 700M random\-initialized inference throughput in tokens/s\. Values are measured on 1x A100\-80GB with BF16 and batch size 1\. Full\-sequence baselines use CUDA graphs for fixed\-shape denoising\. Native BDLM models use block size 256 and optimized cached denoising with graph\-captured block steps where supported\.Model2565121K2K4K8K16K32K65K131K200K262KFull\-sequence attention1,0359151,4211,06276135112236––––Full\-sequence DiffuMamba\-H47054149846138742924010135–––BDLM attention1,9351,8871,9551,9271,7281,6631,3315082631429174Partially Reverse BDLM Mamba\-H1,1461,1671,0821,1581,0191,2811,234941694450364278

## Appendix CModel Architecture Details

Tables[7](https://arxiv.org/html/2607.02805#A3.T7)and[8](https://arxiv.org/html/2607.02805#A3.T8)give the model geometry used for the reported 87M sweep and 350M scale\-up configurations\. All configurations use an MLP ratio of44, timestep embedding dimension256256, Mamba convolution width44, and Mamba expansion factor22\. DiffuMamba\-H hybrid configurations use attention at layers0,6,12,…0,6,12,\\ldotsand Mamba\-2 mixers in the intervening layers\. The BDLM Mamba\-H architecture keeps that sparse\-attention hybrid schedule, but uses the block partial\-reverse construction described in Section[4\.1](https://arxiv.org/html/2607.02805#S4.SS1)rather than a fully bidirectional full\-sequence denoiser\.

Table 7:87M architecture details\. Parameter counts are total trainable parameters for the released model definitions\. Hybrid models use attention at layers0and66\.Table 8:350M architecture details\. Parameter counts are total trainable parameters for the released model definitions\. Hybrid models use attention at layers0,66, and1212; the BDLM Mamba\-H architecture uses a width\-adjusted 350M\-class geometry to match the BDLM attention scale\.Table 9:Key training and inference hyperparameters used across the reported DCLM experiments unless stated otherwise\. Learning\-rate values differ across the 87M sweep and 350M scale\-up configurations as shown in Tables[1](https://arxiv.org/html/2607.02805#S5.T1)and[2](https://arxiv.org/html/2607.02805#S5.T2)\.
## Appendix DNo\-Timestep Versus Timestep Conditioning

Table[10](https://arxiv.org/html/2607.02805#A4.T10)compares no\-timestep and original timestep conditioning at the 87M fixed\-learning\-rate slice used during recipe validation\. The relevant question for this paper is whether BDLM Mamba\-H can be trained with a timestep\-free prefix path without losing the validation signal needed for model selection\.

Table 10:87M DCLM conditioning comparison at LR4×10−34\{\\times\}10^\{\-3\}and 8192\-token validation context\. PPL is the BDLM/MDLM\-compatible diffusion\-NELBO perplexity\. No\-timestep values are copied from the main 87M sweep; timestep\-conditioned models use the DDiT\-normalization path and the same one\-sample BDLM/MDLM\-compatible validation protocol\.We find timestep conditioning improves BDLM attention at this learning rate, slightly worsens BDLM Mamba\-H, and worsens both full\-sequence baselines\. We therefore use the five\-learning\-rate no\-timestep sweep for main model selection, matching the original BDLM experimental convention\(Arriola and others,[2025](https://arxiv.org/html/2607.02805#bib.bib40)\)and preserving a timestep\-invariant prefix path\. Timestep conditioning remains architecturally supported through the active\-block\-only AdaLN factorization in Appendix[E](https://arxiv.org/html/2607.02805#A5), which keeps clean\-prefix caches independent of the reverse timestep\.

## Appendix EAdaLN Factorization

The architectural requirement for reusable prefix states is that they do not depend on the reverse timestep\. Standard diffusion Transformers often inject timestep conditioning into every layer through adaptive normalization\. If that modulation touches prefix tokens, then the layer cache becomes

𝒞ℓ​\(c,t\)=Fprefixℓ​\(c,et\),\\mathcal\{C\}^\{\\ell\}\(c,t\)=F\_\{\\mathrm\{prefix\}\}^\{\\ell\}\(c,e\_\{t\}\),\(18\)which changes at every reverse step and cannot be reused across the active block’s denoising trajectory\. Adaptive layer normalization commonly maps a timestep embeddingete\_\{t\}into shift, scale, and gate parameters and applies them before mixer and feed\-forward updates:

\(δ1ℓ,γ1ℓ,g1ℓ,δ2ℓ,γ2ℓ,g2ℓ\)\\displaystyle\(\\delta\_\{1\}^\{\\ell\},\\gamma\_\{1\}^\{\\ell\},g\_\{1\}^\{\\ell\},\\delta\_\{2\}^\{\\ell\},\\gamma\_\{2\}^\{\\ell\},g\_\{2\}^\{\\ell\}\)=Aℓ​\(et\),\\displaystyle=A^\{\\ell\}\(e\_\{t\}\),\(19\)uiℓ​\(t\)\\displaystyle u\_\{i\}^\{\\ell\}\(t\)=\(1\+γ1ℓ\)⊙LN⁡\(hiℓ\)\+δ1ℓ\.\\displaystyle=\(1\+\\gamma\_\{1\}^\{\\ell\}\)\\odot\\operatorname\{LN\}\(h\_\{i\}^\{\\ell\}\)\+\\delta\_\{1\}^\{\\ell\}\.\(20\)For a full\-sequence denoiser this modulation can touch every token\. For a BDLM prefix cache, applying it to prefix tokens would make cached states timestep\-dependent and force a different cache for every reverse step\. The cacheable layer therefore factors as

𝒞ℓ​\(c\)=Fprefixℓ​\(c\),hb,tℓ\+1=Factiveℓ​\(hb,tℓ,𝒞ℓ​\(c\),et\)\.\\mathcal\{C\}^\{\\ell\}\(c\)=F\_\{\\mathrm\{prefix\}\}^\{\\ell\}\(c\),\\qquad h\_\{b,t\}^\{\\ell\+1\}=F\_\{\\mathrm\{active\}\}^\{\\ell\}\\left\(h\_\{b,t\}^\{\\ell\},\\mathcal\{C\}^\{\\ell\}\(c\),e\_\{t\}\\right\)\.\(21\)The clean prefix path is evaluated once without timestep modulation; AdaLN and timestep gating are applied only to the active\-block continuation\. This remains the cleaner cache semantics even if an implementation uses a fixed timestep for all prefix tokens: the cached state should summarize completed text, not a learned timestep embedding attached to that text\.

TokensEmbeddingNoiseLevelttCacheableDiffusion BlockLayer NormLinear ReshapeLogitsN×N\\timesccPrefixbtb\_\{t\}Blockete\_\{t\}PrefixScanCache𝒞ℓ​\(c\)\\mathcal\{C\}^\{\\ell\}\(c\)RMS NormAdaLNAttn/MambaMixer \+ FFNMLP\+\+

Figure 3:Cacheable AdaLN factorization for block diffusion\. The clean prefix path is evaluated once without timestep modulation, producing a timestep\-invariant cache𝒞ℓ​\(c\)\\mathcal\{C\}^\{\\ell\}\(c\)containing attention keys/values or recurrent sequence state\. At each reverse step, AdaLN\(et\)\(e\_\{t\}\)modulates only the active\-block continuation; the top\+\+denotes the residual add after the MLP update\. This keeps prefix reuse aligned with the BDLM factorization, rather than a post\-hoc approximation to a full\-sequence denoiser\.This factorization is the reason the same prefix cache can be reused across reverse diffusion steps\. It is independent of whether the cached state contains attention keys and values, Mamba recurrent states, or another layer\-specific summary\.