XPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

XPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding August 2, 2026 12 minutes 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, Tdraft. 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: 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. 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. 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]. 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) The full pipeline. The block-diffusion drafter produces hidden states, the target LM head reads out the base logits, and the refiner adds a learned logits bias to form the final logits. (b) Inside the refiner. The three inputs, the token id We[xk−1], the global hidden state g, and the per-position hidden state hk, are down-projected into r-space, fused, mixed causally across block positions, passed through the r-space MLP, and read back out to vocabulary by the shared low-rank head. Hover over a module to see what it is. 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 sk, 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: (i) fuse: ak = Win[ hkWh ‖ g Wg ‖ We[xk−1] ] (ii) mix: ck = ak + Σj≤k Lk,j ⊙ aj (iii) MLP: zk = ck + MLP(ck) (iv) readout: δk = zk Wr ∈ ℝV (v) correct: ℓk = sk + δk The three inputs to the refiner are the drafter’s per-position hidden state hk, the previous-token id xk−1, and a block-global summary g = (1/B) Σj hj, obtained by mean-pooling the drafter’s hidden states over the block. On the output side, sk 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, We ∈ ℝV×r is the token embedding, Wh, Wg ∈ ℝH×r the down-projections for hk and g, Win ∈ ℝ3r×r the input fusion projection, and Wr ∈ ℝr×V the readout LM head. L is a per-channel lower-triangular mixer, so Lk,j is nonzero only for j≤k. Each design choice earns back one of the four properties: Lightweight. Everything except the two vocabulary matrices We, Wr lives in r-space. We and Wr 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 hk 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 t [truncated for AI cost control]