---

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. Measured on a single H200 at batch 1; the acceptance lengths shown are for this one prompt, not the dataset averages reported below.

---

🧠 Abstract

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 end-to-end decoding throughput by ~1.3× on average (up to 1.7×) compared to the original dFlash diffusion drafter.

End-to-end decoding speedup over autoregressive baseline (Qwen3-8B, block 16, batch 1). xPress improves on the dFlash drafter at both temperatures on every benchmark, and the gains survive lossless sampling.

---

🔍 How it works

XPress adds a small causal refiner on top of the drafter (left). The refiner turns the drafter's logits into a corrected block by conditioning each position on the previous block token, and the block is resolved with K parallel Jacobi passes instead of a left-to-right loop.

Everything but the two vocabulary projections lives in an r-dimensional bottleneck (r = 256, right), which makes the head roughly 16× lighter than working at the model's hidden size:

(a) the diffusion drafter with the causal refiner in the loop. The refiner reads the drafter's hidden states and adds a logit bias on top of the target LM head's base logits. (b) inside the refiner: previous-token embedding, per-position and block-global hidden states are fused to r, mixed causally across the block, passed through an r-space SwiGLU MLP, and read out to the vocabulary as a bias.

---

📋 What's in this repo

The exact HF-based harness that produced the paper's acceptance-length and throughput numbers (Qwen3-8B target, dFlash-b16 drafter, block 16), plus a self-contained vLLM serving integration:

---

🚀 Getting Started

🖥️ Environment Setup

conda create -n xpress python=3.11 -y
conda init bash && exec $SHELL      # skip if `conda activate` already works in your shell
conda activate xpress
pip install -r requirements.txt

Notes

• All reported acceptance lengths and throughputs were measured on a single NVIDIA H200. Verified with torch==2.11.0+cu129. Any CUDA 12.x-capable node works. If your driver is older than the wheel's CUDA, prepend the compat layer, e.g. export LD_LIBRARY_PATH=/usr/local/cuda-13.0/compat:$LD_LIBRARY_PATH.
• flash-attn is OPTIONAL: the harness auto-falls back to torch SDPA (published numbers were measured with SDPA).
• First run compiles Triton/inductor kernels (~5 min); later runs reuse the cache (./inductor_cache, override with INDUCTOR_CACHE=...).

📦 Checkpoints

Checkpoints are passed by Hugging Face repo id and download automatically:

To use local files instead, point XPRESS_CKPT / MK_CKPT (or --xpress-refiner-path / --markov-refiner-path) at a .pt path.

---

⚖️ Reproduce the paper numbers

One script for both temperatures — ./bench_fair_fast.sh <dataset> [temperature]:

# T=0 acceptance + paired throughput (published protocol; interleaved rounds,
# warmup discarded, paired ratios):
./bench_fair_fast.sh gsm8k
# T=1 lossless sampling (frozen-Gumbel + honest-q verification):
./bench_fair_fast.sh gsm8k 1
# All benchmarks: gsm8k math500 humaneval mbpp aime25 livecodebench mt-bench

Bare-drafter baseline (the drafter column below): pass --drafter-only to benchmark_compile_all.py (draft = the drafter's own K=0 seed, no refiner; works with or without --use-graph — the seed sampling itself is CUDA-graphed).

Single-method / custom runs go through benchmark_compile_all.py directly (see --help); the .sh wrapper only sets the published protocol.

📊 Results

Both temperatures, all seven benchmarks: speedup over the autoregressive baseline (Sp.) and the acceptance length (τ) that produces it, against the plain dFlash drafter and the Markov-head baseline.

Each entry reports the throughput-optimal K for that setting: bench_fair_fast.sh sweeps K_LIST=4,5,6,7, and the numbers below are the K with the best tok/s together with the acceptance length measured at that same K.

Temperature = 0

Temperature = 1 (lossless sampling; τ averaged over 5 seeds)

Sp. = speedup over the autoregressive baseline; τ = acceptance length (mean accepted tokens per verification step); × Gain = xPress's throughput over the plain dFlash drafter. Qwen3-8B with a dFlash block-16 drafter, single-sequence decoding, at most 2048 generated tokens, on a single H200.

Tolerances: acceptance ±0.1 (bf16 varies slightly across GPU models/driver stacks). Absolute tok/s depends on your node. The paired ratio printed by the interleaved script is the number to compare.

---

⚡ Serving with vLLM

vllm_xpress/ runs xPress as a first-class vLLM V1 speculative-decoding method. Our implementation lives there as readable source (vllm_xpress/src/vllm/...: the refiner head, the draft model, the speculator and its Triton kernels); the only edits to existing upstream files are ~48 lines of method registration in registration.patch. setup_vllm.sh clones upstream vLLM at a pinned commit, installs our files, and self-checks the result:

conda create -n vllm-xpress python=3.11 -y && conda activate vllm-xpress
cd vllm_xpress && ./setup_vllm.sh
python bench_vllm_accept.py --dataset gsm8k --max-samples 128           # xPress
python bench_vllm_accept.py --method dspark --dataset gsm8k --max-samples 128   # Markov baseline

See vllm_xpress/README.md for the serving checkpoint, the K/batch sweeps, and the runtime log lines that confirm the integration is active.

---

🙏 Attribution

The Markov-head baseline runs DeepSeek's own implementation, vendored verbatim under refiners/deepspec/. See refiners/deepspec/PROVENANCE.md. Everything under refiners/xpress* is ours.

---

📚 Citation

If you find our work useful or relevant to your project and research, please kindly cite:

@misc{wang2026xpress,
      title={xPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding},
      author={Zheng Wang and Davis Wertheimer and Yu Chin Fabian Lim and Mudhakar Srivatsa and Raghu K. Ganti and Minjia Zhang and Naigang Wang},
      year={2026},
      eprint={2608.02438},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2608.02438},
}