XPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

Case study. A real GSM8K prompt decoded three ways under the same timing setup: autoregressive, the dFlash drafter alone, and XPress (ours). Each pane advances by the tokens it accepts per target-verification step, so XPress finishes first. Once the dFlash drafter finishes, the autoregressive pane is fast-forwarded (»») so you are not left watching it crawl.

TL;DR. Block-diffusion drafters like dFlash generate an entire block of draft tokens in a single forward pass, drastically reducing the overhead of multiple-token drafting in speculative decoding. The crucial final step of the single-pass discrete denoising process involves using the logit distribution at each position to sample conditionally independent tokens. The resulting draft is thus a set of per-position marginals, rather than a joint distribution: no draft token is guaranteed to depend on its predecessors. Such independently sampled marginals tend to produce sequences with tokens that are individually likely, but jointly improbable under the target model’s distribution, which verifies each token conditionally. This can cause early rejection and limits acceptance length. To address this, we propose XPress as a means to restore the missing causality in diffusion drafters. XPress is a lightweight causal refiner that reconciles the whole diffusion block at once through parallel refinement, restoring and propagating causal dependencies across the draft without a token-by-token loop. On Qwen3-8B, across seven math, code, and chat benchmarks, XPress raises acceptance length by ~30% on average (up to +56%) and its decoding throughput by ~1.3× on average (up to 1.7×) compared to the dFlash diffusion drafter.

1. Diffusion drafters and the problem with parallel prediction

Speculative decoding (SD) [Leviathan et al. 2023] accelerates autoregressive generation by using a lightweight draft model to propose future tokens, which the larger target model verifies in one parallel forward pass. A single multi-token verification pass costs about the same as a standard single-token target-model decoding step in low concurrency scenarios, so every draft token that matches target model outputs is another token produced at no additional target cost. The speedup achieved by SD is mainly governed by two factors: the acceptance length τ (the number of drafted tokens the target accepts per verification step) and the cost of the drafting process itself, T_draft. A larger acceptance length amortizes each target verification pass over more generated tokens, while a cheaper drafter reduces the overhead paid to produce them. This forms a natural trade-off: we want to maximize the expected acceptance length of the drafter, while avoiding a proportional increase to drafting overhead.

Autoregressive (AR) generation has long been the default approach to drafting models, and the EAGLE series [Li et al. 2024a][2024b][2025] is one of the most representative AR-drafting SD methods. EAGLE’s drafter is remarkably lightweight, as small as a single layer, yet it yields high-quality drafts. However, because drafting is autoregressive, generating n draft tokens involves n sequential forward passes of the draft model. As n becomes larger, the drafting overhead becomes increasingly pronounced, but acceptance length, which requires an unbroken chain of accepted verifications, does not. dFlash [Chen et al. 2026] resolves this via a block-diffusion model that proposes an entire block of draft tokens in a single forward pass. By turning n serial steps into a single parallel one, dFlash enables longer drafts (and downstream speedups for the target model) at near-constant overhead.

But the parallelism brought by the diffusion drafter carries an inherent limitation in accuracy. Unlike an AR drafter, where each position is conditioned on the preceding tokens, positions in a diffusion drafter are decoded from their marginal distributions. The token k is drawn without seeing what the token k−1 turned out to be, so the block is a set of individually plausible tokens with no guarantee of causality. This can violate natural linguistic dependencies across positions, even when every token is locally high-probability. For example, a drafter predicting each position independently can put a plural verb after a singular subject and produce “she are”, where each word is fine on its own but the verb contradicts the subject. At verification, these locally reasonable but jointly incoherent samples are rejected early by the left-to-right target model, limiting the achievable acceptance length.

One line of prior work tries to address this limitation by constructing a draft token tree. Tree-based drafting proposes a tree of candidate continuations and verifies the whole tree in a single target pass, so that the longest accepted path through the tree can be selected. Recent works, like PRESTO [Wang et al. 2026] and DDTree [Ringel & Romano 2026], have demonstrated that tree drafting can effectively enhance the achievable acceptance length of diffusion-drafter-based SD methods. Nevertheless, tree drafting has real limitations. Fundamentally, it hedges around the non-causality of the diffusion drafter by targeting recall rather than accuracy. On top of that, the required sparse, irregular tree attention is expensive and complex to serve. Moreover, because a candidate tree typically spans tens to hundreds of tokens, its gains fade quickly at large batch sizes, where the target pass is already compute-bound and those extra tokens are no longer free. The more direct fix is to see whether we can cheaply and effectively restore the lost causality for the diffusion drafter, which is the core research problem motivating XPress.

2. Causal refiner design in XPress

Rather than redesigning the diffusion drafter from scratch to restore causality, we instead ask if a small correction to existing outputs will suffice. What makes this plausible is a property already established for diffusion drafters: the correct token is frequently among the drafter’s top-k candidates at each position [Wang et al. 2026]. The drafter’s block-level marginals already narrow each position to a small candidate set, it just lacks the causal information needed to identify the right token within that set. The problem of correction therefore reduces to picking the right token from a narrow preexisting set, which should be feasible for a lightweight causal refiner. We thus formulate the correction process as a causal refiner with four properties:

1. Lightweight. The refiner has little room for new parameters or architectural complexity, as the drafter’s single parallel forward pass is already highly streamlined and performant, and this existing capability should be preserved.
2. Causal. Within that small resource budget, the refiner should inject real causal information, conditioning each token on its discretely sampled predecessors, rather than merely smoothing the drafter’s marginals locally.
3. Drafter-grounded. The refiner should make good use of what the diffusion drafter already computes, like its hidden states, which typically carry rich information about the block compared to the pure token id [S. L. Wang et al. 2026].
4. Low overhead. Refinement cost at inference should stay a small fraction of the total drafting time, adding as little latency as possible. In particular, we must avoid reintroducing a fully serial, left-to-right pass.

We propose XPress, a lightweight causal refiner instantiating all four properties above. It is lightweight, adding only ~80M parameters (161 MB in bf16) on top of the diffusion drafter; it injects causal information, conditioning each token on its predecessors rather than smoothing marginals locally; it is drafter-grounded, reusing the diffusion drafter’s hidden states rather than a bare token id; and it keeps overhead low, resolving the block in a few parallel iterations rather than a serial left-to-right pass.

Figure 1(a) shows the full pipeline of XPress: the diffusion drafter proposes the initial block in one pass, the target LM head reads out the base logits s_k, and the refiner adds a learned correction on top. Consistent with the aforementioned top-k observation, the refiner does not score the vocabulary from scratch. It adds a small per-position logit bias δ_k to the drafter’s own logits, which re-ranks the handful of candidates the drafter already favours. The proposed causal refiner is shown in Figure 1(b). Let V denote the vocabulary size, H the drafter’s hidden width, B the block length, and r=256 the low-rank dimension. All learned linear projections are named W with a descriptive subscript. The position-k logit bias is built in five steps:

The three inputs to the refiner are the drafter’s per-position hidden state h_k, the previous-token id x_k−1, and a block-global summary g = (1/B) Σ_j h_j, obtained by mean-pooling the drafter’s hidden states over the block. On the output side, s_k is the drafter’s own base logit vector for position k, δ_k is the learned correction, and ℓ_k is the corrected logit the block is re-decoded from. Among the learned maps, W_e ∈ ℝ^V×r is the token embedding, W_h, W_g ∈ ℝ^H×r the down-projections for h_k and g, W_in ∈ ℝ^3r×r the input fusion projection, and W_r ∈ ℝ^r×V the readout LM head. L is a per-channel lower-triangular mixer, so L_k,j is nonzero only for j≤k. Each design choice earns back one of the four properties:

• Lightweight. Everything except the two vocabulary matrices W_e, W_r lives in r-space. W_e and W_r are the embedding and prediction head required by any logit-bias model. On top of them the refiner adds only small r-space components, so it stays a correction rather than overwriting the base drafter outputs.
• Drafter-grounded. Feeding the per-position hidden h_k and the block-global summary g into the correction gives it strictly more signal than a bare token id [S. L. Wang et al. 2026], and these are obtained for free from the drafter.
• Causal. Positions exchange information in r-space, via lower-triangular mixing so that position k sees its whole prefix j≤k. Because the mix step (ii) happens after the prior sampled tokens are incorporated in the fuse step (i), this enables real causal conditioning. The mixer is also lighter than a conventional attention layer, since it is a single fixed triangular combination rather than a computed attention score, yet still expressive, encompassing the full set of learnable conv1d patterns.
• Low overhead. Mixing in r-space is lightweight, and a per-position r→2r→r residual MLP adds cheap nonlinear expressiveness. Steps (i) and (ii) stay linear, so the fuse-and-mix path folds into a single matrix at inference.

For a Qwen3-8B target and a dFlash drafter, XPress adds 80.5M parameters, of which 96% are the two vocabulary maps that any logit-bias head needs. The causal refiner itself, meaning the hidden inputs, the mixer, and the MLP, is only 2.8M (3.4%), and none of it scales with the 152k-token vocabulary. Relative to the 1.05B-parameter dFlash drafter it attaches to (5 transformer layers plus the input projection; the drafter reuses the target’s embedding and LM head), the whole refiner is a 7.7% add-on, and the causal core alone is 0.26%.

3. Parallel refinement via Jacobi decoding

The causal mixer restricts the visibility of every position to its own prefix (j≤k), so the refiner naturally supports autoregressive generation. However, this left-to-right generation process pays an additional B−1-step loop cost over a block of B draft tokens. Instead of finalizing one position before moving to the next, XPress updates all positions at once and repeats this a few times, correcting any prior mistakes, an approach known as Jacobi decoding [Song et al. 2021]. The block converges in far fewer than B−1 iterations, which keeps the correction a small fraction of the draft step. Specifically, XPress seeds all positions from the diffusion drafter’s one-shot predictions, then updates them jointly:

1. Seed in parallel. Take the latest set of discrete token predictions.
2. Refine everything at once. One forward of the refiner re-corrects all positions, accounting for any changes in prior positions that surfaced in the last step.
3. Draw new tokens and repeat. Tokens converge to the ground truth AR output from left to right, in at most B−1 steps.

Formally, sequential greedy decoding computes y_k = argmax_v p_k(v | y_<k, h, g) in order, whereas Jacobi decoding solves the same equations from a seed Y⁽⁰⁾ by updating every position at once,

Each iteration updates every position from the block as it currently stands, and a position stops changing once the tokens before it have also halted. Because each position looks only leftward through the causal mixer, stability spreads rightward from the anchor token. The worst case is one position settling per iteration: given a stable substring of length n, the next prediction for position n+1 is guaranteed to also be stable in future iterations, as it depends only on the previous n tokens, plus the drafter features h and g which are precomputed constants. Thus B−1 iterations is guaranteed to reproduce the exact sequential decode. In practice, though, this worst case is very rare. The drafter’s seed is already a good guess, so on a typical iteration many positions settle at once, and K iterations lock in far more than K tokens. We find that K≈6 iterations is sufficient to yield acceptance length on par with a 15-step sequential decode.

4. Training XPress

The causal refiner is co-trained with the (co-adapted) drafter, using ground-truth token sequences and the predictions of the frozen target model. Both provide useful training signal, which we incorporate into two separate loss terms. The first is a teacher-forced cross-entropy against the ground truth token sequence. Conditioning each position on the ground-truth prefix y_<k, it maximizes the probability of the correct token x_k^*. This is the standard next-token training objective that instills language capability into the refiner. But it optimizes the data likelihood, whereas what sets the speedup of SD is the acceptance rate against the target model. Under speculative sampling, the probability that a token drawn from vocabulary distribution p is accepted against the target model distribution p^t can be expressed as Σ_x min(p(x), p^t(x)) = 1 − TV(p, p^t), with TV = ½‖p − p^t‖₁. So

and minimizing this total-variation distance to the target is exactly maximizing acceptance, so we make it the second loss term in our training objective. The two are complementary: the cross-entropy points the refiner at the right token, while the total-variation term shapes the distribution to the target the way acceptance is scored. The per-position loss is their weighted sum, with w_k = exp(−(k−1)/γ) emphasizing earlier positions, since an inference-time verification mismatch disqualifies not just that position but all following positions in the draft:

During a forward pass, the drafter produces a base distribution p^b, which the refiner then uses to produce the refined distribution p^r. A naive application of our loss to p^r yields a performant refiner, and a drafter co-adapted to its behavior. Yet this can be problematic at inference time: the Jacobi iteration begins from the drafter’s predictions, so if p^b drifts from the target p^t, the refiner is starting from a worse seed. We cannot differentiate through the token-sampling operation in the drafter that captures this dynamic, so we instead add an auxiliary loss on p^b, anchoring it to desired behavior.

An additional concern is the fact that minimizing the loss on p^r in a teacher-forced setting yields a refiner that is good at correcting gold prefixes, yet untested on the self-conditioned inputs it actually receives in the Jacobi decoding process. Therefore, we also introduce a consistency loss, in the spirit of consistency training for Jacobi decoding [Kou et al. 2024], to solve this misalignment between training and inference stages. Refiner forward passes are cheap by design, so during training we run a second forward pass whose token inputs come from the drafter’s argmax(p^b), and apply our two-term loss to that output p̂^r as well.

The full objective thus applies 𝑳 to three distributions:

The first loss optimizes the refiner under teacher forcing, the second drafter-anchor loss preserves the quality of the diffusion drafter itself, and the final consistency loss allows the refiner to better operate under inference conditions.

5. Relation to concurrent works

Two concurrent works, Domino [Huang et al. 2026] and DSpark [Cheng et al. 2026], share our goal of restoring causality to a diffusion drafter, each by attaching a correction head that makes a position’s logits depend on the tokens that came before. Both restore causal structure, but both pay for it in ways XPress does not: they decode the correction serially, and feed the refiner narrow inputs.

Domino uses a GRU to walk the diffusion block left to right, carrying a recurrent hidden state and emitting a per-position logit correction. Because position k’s correction depends on the realized prefix, causal structure is restored. But the GRU is a recurrence: it must step through the block one position at a time. DSpark aims to shed the overhead and complexity of Domino’s left-to-right GRU by attaching an even lighter Markov head. For each position it adds a bias that depends only on the identity of the previous token,

Here W₁ ∈ ℝ^V×r embeds the previous token id and W₂ ∈ ℝ^r×V reads the correction back out. This is exactly the degenerate corner of XPress’s design space (§2): keep only the token-id input, drop the causal mixer and the MLP, and the refiner collapses to the same bigram bias δ_k = W_e[x_k−1] W_r.

The two heads differ in mechanism but share two weaknesses. The first is that both are serial. Domino’s GRU passes its state from one position to the next, and DSpark’s bias for token k is indexed by the sampled id of token k−1. So either way a block of B tokens takes B−1 steps that must run in order. (In SGLang’s DSpark implementation the head is a plain Python for-loop [sglang #30261]; CUDA-graph capture unrolls it and replays the kernels with negligible launch overhead, but the B−1 steps still execute strictly one after another because each waits on the last.) The second weakness, most acute in DSpark, is narrow inputs. DSpark’s Markov head conditions on a single previous token id, limiting it to repairing local two-token clashes, blind to both the rest of the block and the drafter’s hidden states, precomputed representations that encode far more than any single token id. Domino’s GRU state carries more than one token, but it still never sees the whole block at once. XPress, on the other hand, mixes over positions explicitly, and can even see into the future (in a limited fashion) despite its causal structure, by accessing the entire block’s worth of drafter hidden states.

6. Experiments

Setup. We use Qwen3-8B as the target model and dFlash as the base drafter. The block size for the diffusion drafter is 16, the sampling temperature is 0, and the maximum number of generated tokens is 2048. All heads compared share the same training recipe and the same harness at inference (CUDA graph and torch.compile), measured on H200. The results are for a single-sequence batch size.

Acceptance length. XPress raises acceptance length over the Markov head on all seven benchmarks, by +3.0% to +7.0% (mean +5.0%), with the widest margins on code (LiveCodeBench +7.0%, HumanEval +5.0%). Read against the bare dFlash diffusion drafter, both heads do most of the work, but XPress captures more of it, +29% over the drafter’s own τ on average versus +23% for the Markov head.

How many iterations are needed? More Jacobi iterations lock in more of the prefix, but that is not the same as acceptance rising monotonically with K. When the accepted prefix already reaches past the settled region, one more iteration can overwrite a not-yet-converged tail token that happened to match the target, so τ can dip slightly. In practice τ rises quickly, crosses the Markov baseline within a few iterations, and plateaus by K≈7; pushing to K=16 never beats the plateau, so a small K captures essentially all of the gain.

Drafting-time latency. As shown in Figure 3, XPress’s cost grows with the number of Jacobi iterations: each Jacobi iteration adds about 75 µs, so the refiner runs from 150 µs at K=1 to 681 µs at K=8. The Markov head is a fixed 601 µs, its 15-step serial decode regardless of K, so the two cross near K=7. As shown in Figure 4, XPress never needs to run that far. It matches or beats the Markov head’s per-step acceptance on every benchmark by K=4, and at K=4 the refiner costs just 379 µs, a 1.6× speedup over the Markov head. In other words, at the first point where XPress is already more accurate, it is also markedly faster. Pushing on to the accuracy plateau at K=6 still leaves it cheaper (530 µs, a 1.13× speedup).

Figure 5. Composition of the draft-side latency per block for the Markov head (gsm8k, block 16, H200/sdpa). The drafter forward and the base lm_head are shared across heads; the refiner is the CUDA-graphed 15-step serial decode, about a quarter of the draft cost.

Figure 5 breaks the draft step into its three parts: the drafter’s forward pass, the base lm_head readout, and the causal refiner. The first two are shared by every head; the only difference is the refiner. But that refiner cost is non-trivial: even run as a compiled unrolled loop, the Markov head’s 15-step serial decode is about 24% of the draft-side cost. XPress acts to minimize this slice, replacing the 15-step serial decode with a handful of Jacobi iterations.

End-to-end throughput. As shown in Figure 6, XPress reaches up to 8.2× over autoregressive decoding (GSM8K) and averages 6.2×, versus 5.9× for the Markov head and 4.9× for the bare drafter. Against the Markov head, though, the margin is only +5%, well short of the drafting-time speedup above. The reason is that each decode step is dominated by the target’s verification pass. The drafting-time speedup would carry through more fully in settings where the draft is a larger share of each step, for example larger blocks or a cheaper, quantized target, where we would expect the latency advantage to surface end to end.

7. Conclusion

XPress is a lightweight causal refiner for block-diffusion drafters. It reads the drafter’s own hidden states and reconciles the whole block through a few parallel Jacobi iterations, restoring the token-to-token dependencies that a parallel drafter drops, without falling back to a serial, left-to-right decode. Across seven math, code, and chat benchmarks on Qwen3-8B, XPress raises the dFlash drafter’s acceptance length by about 30% and its end-to-end throughput by about 1.3×. Moreover, by running four Jacobi iterations, XPress already matches or beats the Markov head’s acceptance length while running about 1.6× faster.

There are several parts that could be further optimized. The iteration count K is currently fixed for a whole run, but most blocks converge with a small number of iterations, so an adaptive rule that stops a block once it stops changing, or a schedule that spends iterations only where the draft is still unsettled, could cut the average K with no loss in acceptance. More broadly, the philosophy introduced by XPress is not specific to speculative decoding. Reconciling a block of mutually dependent predictions in a few parallel iterations may help wherever a model emits many interdependent outputs at once.

8. Citation

If you find XPress useful, please cite:

@article{xpress2026,
  title   = {XPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding},
  author  = {Wang, Zheng and Wertheimer, Davis and Lim, Yu Chin Fabian and
             Srivatsa, Mudhakar and Ganti, Raghu K. and Zhang, Minjia and Wang, Naigang},
  year    = {2026},
}

9. References

1. J. Chen, Y. Liang, Z. Liu. DFlash: Block Diffusion for Flash Speculative Decoding. 2026. arXiv:2602.06036
2. J. Huang, Y. Zhang, Q. Zhang, H. Lin, H. Xu, L. Zhang. Domino: Decoupling Causal Modeling from Autoregressive Drafting in Speculative Decoding. 2026. arXiv:2605.29707
3. X. Cheng, X. Yu, C. Shao, J. Li, Y. Xiong, et al. (DeepSeek-AI). DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation. 2026. arXiv:2607.05147
4. Y. Leviathan, M. Kalman, Y. Matias. Fast Inference from Transformers via Speculative Decoding. ICML 2023. arXiv:2211.17192
5. Y. Li, F. Wei, C. Zhang, H. Zhang. EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. ICML 2024. arXiv:2401.15077
6. Y. Li, F. Wei, C. Zhang, H. Zhang. EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees. EMNLP 2024. arXiv:2406.16858
7. Y. Li, F. Wei, C. Zhang, H. Zhang. EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test. 2025. arXiv:2503.01840
8. Z. Wang, Z. Ye, Q. Cheng, Y. Fu, Z. Wang, F. Zhu, H. Zhao, J. Kautz, P. Molchanov, H. Shi, M. Zhang. PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding. 2026. arXiv:2607.22634
9. L. Ringel, Y. Romano. Accelerating Speculative Decoding with Block Diffusion Draft Trees. (DDTree.) 2026. arXiv:2604.12989
10. Y. Song, C. Meng, R. Liao, S. Ermon. Accelerating Feedforward Computation via Parallel Nonlinear Equation Solving. ICML 2021. arXiv:2002.03629
11. A. Santilli et al. Accelerating Transformer Inference for Translation via Parallel Decoding. ACL 2023. arXiv:2305.10427
12. Y. Fu, P. Bailis, I. Stoica, H. Zhang. Break the Sequential Dependency of LLM Inference Using Lookahead Decoding. ICML 2024. arXiv:2402.02057
13. S. Kou, L. Hu, Z. He, Z. Deng, H. Zhang. CLLMs: Consistency Large Language Models. ICML 2024. arXiv:2403.00835
14. S. L. Wang, P. Isola, B. Cheung. The Truth Lies Somewhere in the Middle (of the Generated Tokens). 2026. arXiv:2605.09969