PaperlayerView exact revision on arXiv

Source-linked AI summary

One Block Ahead: Promoting Sampler State in Diffusion Language Models

Yangyi Shen, Pin-Jui Ku, Nicolai Oswald, Jean-Marie Lemercier, Yonggan Fu, Stefano Ermon, Jure Leskovec, Boris Ginsburg, Ante Jukić

arXiv:2610.19510v1cs.CL

TL;DR

Diffusion LMs denoise blocks in parallel but still process blocks sequentially, leaving cross-block acceleration open. OBA denoises a warm-up block alongside the current block and promotes its sampler state; on DiffusionGemma, it raises TPF by 42.3% and single-request throughput by 16.3% across ten benchmarks, with a small accuracy regression.

  • Problem

    Diffusion LMs parallelize denoising within blocks, but blocks are still denoised sequentially, leaving concurrent decoding of successive blocks an open question.

    Source: §1 · p. 1

  • Method

    OBA concurrently denoises the next block and promotes it with its tokens, self-conditioning signal, and denoising time when the current block commits.

    Source: §1 · p. 2, §3 · p. 4, §3.2 · p. 5

  • Results

    Across ten DiffusionGemma benchmarks, OBA raises TPF from 15.64 to 22.25 (+42.3%) and single-request denoising throughput by 16.3%, with a small average-accuracy reduction.

    Source: §1 · p. 3, §6 · p. 11

  • Takeaways & Limitations

    OBA can improve throughput on DiffusionGemma when sampler state is promoted with the warm-up block, rather than promoting tokens alone.

    Source: §4.3 · p. 9, §6 · p. 11

    Evaluation covers batch size one, one GPU type, and one model; the promotion mechanism is specific to uniform-state, block-autoregressive diffusion.

    Source: §7 · p. 11

Abstract

from arXiv · show

Block-diffusion language models denoise a whole block of token positions in parallel so each forward pass can emit many tokens. However, blocks are still denoised in strict sequence, a cost that within-block acceleration leaves untouched. Here, we investigate decoding blocks concurrently and propose One Block Ahead (OBA), a training-free decoding algorithm that denoises the next block alongside the current one. We demonstrate that for uniform-state diffusion LMs, where a block's denoising progress is not visible in its intermediate tokens, concurrent decoding is advantageous if the next block's tokens are promoted together with the corresponding sampler state, including the self-conditioning signal and the denoising time. We also derive a break-even criterion that relates the benefit of OBA to the relative cost of a paired-block forward pass. On DiffusionGemma, a state-of-the-art diffusion LM, OBA increases the tokens per forward pass by 42.3% (from 15.64 to 22.25) and the single-request denoising throughput by 16.3%, averaged over ten benchmarks, with a small accuracy regression.

1. Introduction

Diffusion language models gain parallel decoding within blocks, but DiffusionGemma still denoises blocks sequentially. OBA overlaps adjacent blocks, carrying warm-up sampler state forward, and improves throughput on DiffusionGemma with a small accuracy reduction.

1. Introduction

  • Diffusion decoding parallelizes positions within a block, but accepting more positions per step trades quality for latency and does not remove sequential decoding across blocks.DiffusionGemma offers around 20 tokens per forward pass and roughly 1,500 output tokens per second on one H100 GPU, yet still decodes blocks sequentially.Source: §1 · p. 1
  • Concurrent decoding for DiffusionGemma remains unestablished despite prior multi-block methods, whose reported results use 32-token blocks and baselines averaging 3.47 or 4.62 TPF.For DiffusionGemma, self-conditioning and denoising time must survive block transitions, while doubling its native 256-token blocks to 512 positions may reduce hardware utilization.Source: §1 · pp. 1–2
  • OBA denoises the next block alongside the current block, promoting the warm-up block with its accumulated tokens and sampler state when the current block commits.The paired-block layout uses a block-causal attention mask so the warm-up block can attend to the current block, but not vice versa.Source: §1 · p. 2, Figure 1 · p. 2
  • OBA tests how to structure attention between current and warm-up blocks and which warm-up components, including tokens, self-conditioning, and denoising time, to retain at promotion.The study also optimizes denoising-step components for single- and paired-block layouts to compare how tokens per forward pass translate into throughput.Source: §1 · pp. 2–3
  • TPF rises from 15.64 to 22.25 (+42.3%) and single-request denoising throughput from 2,375 to 2,761 TPS (+16.3%) across ten DiffusionGemma benchmarks, with a small average-accuracy reduction.By contrast, doubling the block to 512 positions lowers throughput by 41.4% and reduces accuracy further.Source: §1 · p. 3

2. Background

The background distinguishes diffusion processes and explains DiffusionGemma’s block-based generation and sampling. It also defines the inference-efficiency metrics used to evaluate decoding.

2.1. Discrete Diffusion Language Models

  • Masked diffusion reveals positions monotonically, whereas uniform-state diffusion keeps real tokens revisable until commitment, so tokens alone do not show a block’s denoising progress.DiffusionGemma uses uniform-state diffusion; block diffusion adds causal attention across blocks and bidirectional attention within the current block, with committed blocks appended to the cache once.Source: §2.1 · p. 3

2.2. DiffusionGemma

  • DiffusionGemma uses an entropy-bounded sampler that accepts predictions in increasing-entropy order up to a budget, then uniformly corrupts other positions and can stop early when entropy and argmax stability criteria are met.The model is an instruction-tuned mixture-of-experts diffusion LM with 26 billion total and 4 billion active parameters per token.Source: §2.2 · pp. 3–4

2.3. Inference Efficiency Metrics

  • TPF measures useful output tokens per request-serving forward pass, while decode TPS measures useful output tokens per second of denoising GPU time, excluding prefill, cache appends, and host overhead.Adaptive stopping makes denoising-step counts variable, so DiffusionGemma measures denoising efficiency empirically.Source: §2.3 · p. 4

3. Method

OBA pairs the current block with a warm-up block, promotes accumulated sampler state when the current block commits, and raises throughput when saved denoising passes outweigh the paired-pass cost.

3.1. Paired-Block Layout

  • OBA runs the decoder over current and warm-up blocks together, using block-causal attention so the current block’s logits match those produced when denoised alone in exact arithmetic.The warm-up block attends to the current block, while the current block cannot attend to the warm-up block; grouped mixture-of-experts computation may nevertheless cause floating-point differences from single-block decoding.Source: §3.1 · pp. 4–5

3.2. State Inheritance at Promotion

  • State inheritance means retaining a warm-up block’s progress at promotion; the sampler state comprises self-conditioning, denoising time, and adaptive-stopping bookkeeping.OBA promotes the tokens, self-conditioning signal, and denoising time, but resets stopping bookkeeping so the promoted block does not use a stability judgment formed before its context was finalized.Source: §3.2 · p. 5
  • Under uniform-state diffusion, tokens alone cannot reveal denoising progress, so the sampler state cannot be reconstructed from the block tokens.This makes inheriting state beyond tokens material to promotion in this setting.Source: §3.2 · p. 5

3.3. Break-Even Criterion

  • The break-even criterion predicts higher throughput when saved current-block passes per warm-up step, α, exceed the paired-pass cost relative to Native, r−1.The model assumes a constant linear saving α across blocks and passes; under those assumptions, expected passes approach a fixed point and throughput changes by (1 + α)/r.Source: §3.3 · pp. 5–6

4. Experiments

Across ten DiffusionGemma benchmarks, OBA improves decoding efficiency over the native 256-token configuration, while accuracy and gains vary across comparisons. The experiments examine sampler-state inheritance, attention, and implementation choices that affect whether paired-block decoding pays off.

4.1. Experimental Setup

  • The evaluation uses DiffusionGemma at batch size one on a single GB300 across ten benchmarks, comparing native 256-token decoding with a 512-token block and concurrent FlowBlock decoding.FlowBlock (W=2) provides a comparison for when the warm-up block starts; the study also ablates inherited state and attention.Source: §4.1 · p. 6

4.2. Decoding Efficiency

  • 42.3% higher TPF (22.25 vs. 15.64) and 16.3% higher TPS (2,761 vs. 2,375) over Native (C=256), with average accuracy down 1.37 percentage points.OBA leads both metrics on all ten benchmarks; results average three generation seeds.Source: Table 1 · p. 7, §4.2 · p. 7
  • FlowBlock (W=2) raises TPS 4.5% over Native (C=256), versus 16.3% for OBA; Native (C=512) raises TPF from 15.64 to 16.67 but cuts throughput 41.4%.FlowBlock's readiness gate retains 0.73 accuracy points over OBA, while its margin α−(r−1) is 0.06 versus 0.28 for OBA.Source: §4.2 · pp. 7–8
  • 39.9% of promoted positions already hold their committed tokens on the first current-block pass, versus 5.1% for newly initialized Native (C=256) blocks.The fractions of first-pass acceptances retained until commitment are similar: 97% under OBA and 94% under Native.Source: §4.2 · p. 8, Figure 2 · p. 8

4.3. Ablation Study

  • Promoting tokens alone yields α=0.25, below the break-even r−1=0.26; adding self-conditioning raises α to 0.41, and adding denoising time raises it to 0.54.Withholding self-conditioning lowers TPS 19.9% below Native (C=256), while allowing positions to settle before promotion makes little difference once the signal is propagated.Source: §4.3 · p. 9, Table 2 · p. 9
  • Dense Mask reaches TPF 20.05 versus 15.64 for Native (C=256), but its 1.70× pass cost requires α>0.70; at α=0.38, throughput is 20.9% lower.Its average accuracy is 80.29 versus 82.70 for OBA, as the current block conditions on one still near noise.Source: §4.3 · pp. 9–10

4.4. Kernel Implementation

  • Denoising passes take 7.08 ms with Native (C=256), 8.94 ms with OBA, and 13.06 ms with Native (C=512); averaged over benchmarks, paired passes cost r=1.26 versus r=1.84 for dense passes over 512 positions.Source: §4.4 · p. 10

5. Related Work

Prior work accelerates denoising within a block or reduces serial decoding across blocks. OBA differs by carrying sampler state across blocks in uniform-state diffusion, where progress is not visible in the tokens alone.

Within-block acceleration

Multi-block decoding

  • D2F (Wang et al., 2026) and MBD-LM (Jin et al., 2026) post-train masked diffusion models for bounded windows, while FlowBlock (Tian et al., 2026) admits a downstream block without training and uses editing to repair stale-prefix drafts.These approaches target masked diffusion, where revealed tokens and mask ratio expose block progress.Source: §5 · pp. 10–11
  • Uniform-state diffusion hides progress in the tokens, so OBA must carry sampler state (z, t, σ) across commitment; at 256-token blocks, its paired pass pays only when saved steps α exceed r−1.For DiffusionGemma, r = 1.26, so the break-even condition is α > 0.26.Source: §5 · p. 11
  • On DiffusionGemma, FlowBlock’s readiness threshold delays warm-up enough to yield a 4.5% throughput gain over Native (C=256), compared with OBA’s 16.3%.The comparison uses the same sampler state and isolates when the warm-up block starts.Source: §5 · p. 11

6. Conclusion

One Block Ahead denoises a warm-up block alongside the current block, then promotes it with its sampler state. On DiffusionGemma, it improves token and throughput measures, with a 1.37-point accuracy cost.

6. Conclusion

  • OBA denoises the next block alongside the current one under a block-causal mask, then promotes the warm-up block with its tokens and sampler state.The current block is denoised as if alone; promotion crosses the commitment boundary with the warm-up state.Source: §6 · p. 11
  • 22.25 tokens per forward pass versus 15.64, and 16.3% higher single-request throughput, with a 1.37-point accuracy cost on DiffusionGemma averaged over ten benchmarks.OBA’s α = 0.54 exceeds the paired-pass break-even cost r−1 = 0.26.Source: §6 · p. 11
  • Promoting tokens alone yields α = 0.25 and lowers throughput, while adding self-conditioning and then denoising time raises α to 0.41 and 0.54.The ablations attribute the gain to retaining sampler state across promotion.Source: §6 · p. 11
  • The break-even criterion is that saved passes per warm-up step, α, exceed the extra paired-pass cost, r−1.The results indicate that cross-block parallelism remains available even when a model already decodes one block efficiently.Source: §6 · p. 11

7. Limitations

The study evaluates OBA in a narrow model and serving setting, documents its implementation and use of AI coding tools, and provides a decoder state-machine description.

7. Limitations

  • The evaluation is limited to batch size one, one GPU type, and one model; the promotion mechanism is specific to uniform-state, block-autoregressive diffusion, unlike masked diffusion where progress is visible in masks.The authors identify DiffusionGemma as the only model they know of with both uniform-state and block-autoregressive properties.Source: §7 · p. 11

AI use statement

  • Generative AI coding agents implemented methods, baselines, kernels, the evaluation pipeline, and ran experiments; authors checked AI-assisted components, while tools also edited manuscript text.The authors state that AI tools did not propose research ideas, design methodology or experiments, interpret results, or write proofs.Source: AI use statement · pp. 11–12

Ethics statement

  • The work uses public benchmarks and a pretrained model, adds no capabilities beyond that model, and reports no human subjects, personal data, or significant ethical concerns.The authors state that the work adheres to the ICLR Code of Ethics.Source: Ethics statement · p. 12

Reproducibility statement

  • Reproduction details include the decoding algorithm, assumptions behind the break-even criterion, benchmark and hardware setup, sampler defaults, FlowBlock thresholds, and kernel adaptations.The authors note that throughput, unlike tokens per forward pass, depends on implementation.Source: Reproducibility statement · p. 12

A. Complete Decoder State Machine

  • The decoder runs paired current and warm-up blocks, commits only the current block, then promotes the warm-up tokens, self-conditioning signal, and time while resetting σ.It reserves at least one pass under the newly committed context and seeds a fresh warm-up block after promotion.Source: A. Complete Decoder State Machine · pp. 15–16

B. Kernel Implementation

The implementation compares configurations on a shared GB300 inference engine and optimizes precision, attention, expert layers, sampler operations, and paired-block dispatch. The measurement protocol reports per-pass denoising cost and relative layout cost across benchmarks.

B. Kernel Implementation

  • All five configurations use one inference engine optimized for a single GB300, with later builds validated against Native and OBA outputs to separate algorithm differences from kernel quality.Configurations share model export, sampling parameters, stopping rule, and batch size, while each receives a profile for its own input shape.Source: B. Kernel Implementation · p. 16

B.1. Precision and Quantization

  • NVFP4 4-bit weights and activations replace 16-bit activations in mixture-of-experts layers, while dense projections use FP8 and rotary tables use FP32.Amax reduction, scale finalization, and row quantization are fused into one pass.Source: §B.1 · p. 16

B.2. Attention

  • Fused attention kernels parallelize over cache H to better occupy the device at batch size one, and ring-buffered sliding-attention keys and values make two live blocks affordable at evaluated context lengths.The fallback ragged masked softmax is vectorized, with a scalar variant retained for exact validation.Source: §B.2 · p. 16

B.3. Mixture-of-Experts and Dense Layers

  • The mixture-of-experts and dense branches run concurrently on separate streams after their independent computations, with a bit-exact transformation that improves throughput in an A/B test.The implementation also fuses and parallelizes routing, gather, combine, and activation operations.Source: §B.3 · p. 17

B.4. Sampler

  • The self-conditioning output pass is specialized at two blocks per multiprocessor, doubling occupancy without spilling; one reduction computes the entropy-bounded acceptance rule for both halves.Source: §B.4 · p. 17

B.5. Kernels per Block Layout

  • OBA uses paired block-causal dispatch for two 256-query visibility views and processes shared producers and experts over 512 positions, while committing only the leading 256.The paired pass computes activation scales from the current block alone; grouped 512-row reductions can still change logits and sampler acceptances, so comparisons use matched prompts and seeds rather than identical trajectories.Source: §B.5 · pp. 17–18

B.6. Measurement Protocol

  • The protocol averages per-benchmark denoising time per pass and relative cost r across ten benchmarks; a paired pass costs less than a dense pass over the same 512 positions.The caption explains that Dense Mask’s lower per-pass time than Native (C=512) reflects shorter responses and fewer long-context passes.Source: §B.6 · p. 18, Table 3 · p. 18

C. Proofs

Under the stated mask and exact-arithmetic assumptions, the warm-up block does not affect the current block’s logits or, with matched random variates, its resulting sampler state and stopping decision.

C. Proofs

  • The current-block logits match a Native (C=256) pass and are independent of the warm-up state under the stated mask and operation assumptions.The proposition assumes exact arithmetic, with non-attention operations acting positionwise or reducing only over the current block.Source: C. Proofs · pp. 18–19
  • With identical sampler random variates, the paired pass produces the same current-block state and stopping decision as Native (C=256).Source: C. Proofs · p. 19
  • Current-block queries see the same keys as in the Native pass, while the remaining operations cannot introduce dependence on the warm-up block.The proof establishes this layer by layer, yielding identical current-block hidden states and logits.Source: C. Proofs · p. 19
  • The proposition is structural: shape-specific reductions and grouped mixture-of-experts calls may differ numerically, so experiments match prompts and seeds rather than trajectories.Source: C. Proofs · p. 19

D. Estimation of 𝛼

The appendix estimates the promotion factor from benchmark-level pass counters and relates it to measured throughput, while noting that the fit depends on simplifying assumptions.

D. Estimation of 𝛼

  • The fit estimates α separately per benchmark from aggregate counters across cases and three seeds, then reports an unweighted mean across benchmarks.It uses the assumptions of Section 3.3.Source: D. Estimation of 𝛼 · p. 19
  • Denoising passes are counted as F = Q − K; Native (C=256) supplies the reference mean passes per block because every block starts from noise.The block estimate counts a partially filled final response block as half a block on average.Source: D. Estimation of 𝛼 · p. 19
  • The recurrence fit models first-block overhead and the transient after promotion; averaging block passes directly understates α, while dropping first blocks overstates it.The first block pays n̄, while the second block is promoted with n̄ warm-up steps and needs fewer than n⋆ passes.Source: D. Estimation of 𝛼 · pp. 19–20
  • OBA’s benchmark-averaged ᾱ is 0.47, compared with the fitted α of 0.54.Source: D. Estimation of 𝛼 · p. 20
  • OBA clears the break-even α > r − 1 criterion, while Dense Mask does not; measured TPS ratios are 1.17 for OBA and 0.79 for Dense Mask.The TPF ratio for OBA is 1.43 versus 1 + ᾱ = 1.47 because paired decoding reduces denoising passes but not the cache-append pass.Source: D. Estimation of 𝛼 · p. 20
  • The estimate assumes constant α and block-independent n̄; its approximation and omitted term bias fits most on GSM8K and MMLU, averaging about three and five blocks per response.Source: D. Estimation of 𝛼 · p. 20

E. Evaluation Protocol

The evaluation uses a configured population of 4,736 cases, matched across configurations and seeds, and reports benchmark-level means with each benchmark weighted equally.

E. Evaluation Protocol

  • The configured evaluation population is 4,736 cases, with benchmark-specific primary graders and matched case identities, prompt hashes, and seeds across configurations.It comprises the listed benchmark subsets, not every example in the original releases; accuracy is single-generation accuracy averaged over three seeds.Source: E. Evaluation Protocol · p. 20
  • Across ten benchmarks and three seeds, OBA scores 82.70 against 84.07 for Native (C=256), a 1.37-point accuracy loss.Reported figures are unweighted means of per-benchmark values, so each benchmark counts equally rather than weighting by case count.Source: E. Evaluation Protocol · p. 21
  • Serving forwards comprise prefill, cache-append, and denoising passes, with each paired-block call counted as one pass; TPF divides generated tokens by denoising passes.The benchmark average is the unweighted mean of per-benchmark ratios.Source: E. Evaluation Protocol · p. 21

F. FlowBlock Adaptation

FlowBlock’s adapted readiness gate rarely runs more than one block at a time as its window widens; position-level measurements show high survival of promoted-block first-pass acceptances.

F. FlowBlock Adaptation

  • FlowBlock derives readiness from rejected positions newly accepted at the current step, admitting a successor above 0.6 on general and math tasks or 0.8 on coding tasks.At W=2 it differs from OBA in when it starts the next block and in switching pass width between 256 and 512 positions; promoted blocks retain tokens, self-conditioning, and denoising time.Source: F. FlowBlock Adaptation · p. 21
  • At every FlowBlock window width, more than 92% of denoising passes run one block, three blocks are live on at most 0.63%, and average TPF and TPS stay within 1% of W=2.Source: F. FlowBlock Adaptation · p. 21, Table 4 · p. 22

G. Position-Level Stability of Promoted Blocks

H. Per-Benchmark Results

This section organizes per-benchmark accuracy and wall-throughput results for the main comparison and ablations, alongside per-benchmark α and r diagnostics. For No Signal, negative α values are small and consistent with a warm-up that saves no denoising passes.

H. Per-Benchmark Results

Related papers

Loading 2610.19510v1…