Wavefront-Parallel Decoding for Self-Correcting Diffusion Language Models

FlowBlock is a training-free parallel decoding framework for block-wise diffusion language models (dLLMs). Block-wise dLLMs such as LLaDA-2.x decode one block at a time, which makes inter-block decoding strictly serial. FlowBlock observes that the token-to-token (T2T) self-correction ability of LLaDA-2.1 series enable a downstream block start from an informative draft instead of a finalized predecessor — so block finality becomes a scheduling resource rather than a hard dependency.

FlowBlock is built on two mechanisms:

• Gated Wavefront Decoding (GWD) — admits blocks into a bounded wavefront only when a readiness gate (SPAWN_TH) is satisfied, jointly refines active blocks via T2T editing, and commits blocks in order under a windowed block-causal mask that preserves exact frozen-prefix KV-cache reuse.
• Heterogeneous Wavefront Packing (HWP) — gives each request its own wavefront while packing all active windows into dense, shape-stable batched forwards, so asynchronous progress stays GPU-efficient in batched serving.

✨ Highlights

• Up to 2.95× / 4.01× higher TPS than LLaDA-2.1 / LLaDA-2.0 under batched serving (1.57× / 3.17× at batch size 1), with up to 53.6% / 77.1% lower latency.
• Accuracy preserved or improved: +1.3 points on average across eight math / code benchmarks.
• No training: runs the unmodified LLaDA-2.1 series weights.

🛠️ Installation

Requirements: Linux, NVIDIA GPU(s) with ~80 GB memory for the 16B-A1B MoE checkpoints, CUDA 12.x driver.

git clone <this-repo-url> FlowBlock
cd FlowBlock

conda create -n flowblock python=3.11 -y
conda activate flowblock

# 1) Install PyTorch matching your CUDA version (validated: 2.8.0 + cu128)
pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128

# 2) Install the base dependencies (sglang, lm-eval, ...)
pip install -r requirements.txt

# 3) Install vLLM last — this is a separate step on purpose
pip install vllm==0.10.2

The dinfer package is used directly from source — the evaluation scripts set PYTHONPATH automatically, no pip install step is needed. Verify the setup:

PYTHONPATH="$PWD/dInfer/python" python -c \
  "from dinfer.decoding import FlowBlockJointDiffusionLLM, flowblock_packed; print('FlowBlock import OK')"

Models

Method	Checkpoint
flowblock, llada2.1	inclusionAI/LLaDA2.1-mini
llada2.0	inclusionAI/LLaDA2.0-mini

python dInfer/evaluations/download_hf_model.py \
  --repo_id inclusionAI/LLaDA2.1-mini --local_dir /path/to/LLaDA2.1-mini

Datasets

Benchmark data is fetched automatically on first run (GSM8K is downloaded to test_datasets/, the other tasks load from the Hugging Face Hub). For air-gapped machines, prefetch the code benchmarks once on a connected host:

TEST_DATASETS_ROOT=./test_datasets/code bash dInfer/evaluations/prefetch_code_datasets.sh

🚀 Running Evaluations

run_eval.sh is the single entry point. Pick a method, a task family (or a specific benchmark), and a model path:

# FlowBlock on math (defaults: GSM8K, W=2, SPAWN_TH=0.6)
bash run_eval.sh --method flowblock --task math --model-path /path/to/LLaDA2.1-mini

# FlowBlock on code (defaults: HumanEval, W=2, SPAWN_TH=0.8)
bash run_eval.sh --method flowblock --task code --model-path /path/to/LLaDA2.1-mini

# Serial baselines
bash run_eval.sh --method llada2.1 --benchmark math500 --model-path /path/to/LLaDA2.1-mini
bash run_eval.sh --method llada2.0 --task math --model-path /path/to/LLaDA2.0-mini

# FlowBlock and LLaDA-2.1 back-to-back on the same seed-fixed subset,
# 4 GPUs (data-parallel), batch size 8 per GPU
bash run_eval.sh --method flowblock,llada2.1 --task math \
  --model-path /path/to/LLaDA2.1-mini --gpus 0,1,2,3 --batch-size 8

Benchmarks

--task	--benchmark options	Default
math	gsm8k, math500, algebra, asdiv	gsm8k
code	humaneval, humaneval_plus, mbpp, mbpp_plus	humaneval

Key options and defaults

Option	Default	Meaning
--method	flowblock	flowblock | llada2.1 | llada2.0 (comma-separated to combine)
--window	2	FlowBlock wavefront window size W
--spawn-th	0.6 math / 0.8 code	Readiness gate threshold θ_spawn
--batch-size	4	Per-GPU batch size
--gen-length	2048	Max generation length
--block-length	32	Diffusion block length
--threshold	0.95 (2.0) / 0.5 math, 0.9 code (2.1 & FlowBlock)	Decode / commit threshold
--gpus	0	GPU ids; one data-parallel shard per GPU
--num-samples	full benchmark	Evaluate a seed-shuffled subset
--output-dir	outputs/<benchmark>_<timestamp>	Result directory

Run bash run_eval.sh --help for the full list. Results land in:

<output-dir>/<method>/summary.json    # accuracy, TPS, latency, TPF per method
<output-dir>/global_summary.log       # side-by-side comparison table

For multi-configuration sweeps (block lengths, batch sizes, thresholds), edit RUN_COMBOS in dInfer/evaluations/run_flowblock_joint_batch_compare_loop.sh and run it directly; aggregate_global_results.py merges everything into one table afterwards.

⚠️ Code benchmarks execute model-generated Python (HF_ALLOW_CODE_EVAL=1). Run them only in an isolated environment (container / VM without secrets).

📊 Results Preview

Under the paper settings (block length 32), FlowBlock consistently dominates the serial baselines in the throughput–latency plane while matching or improving accuracy.

Single-request decoding (batch size 1)

Main results at batch size 1 on a single 80 GB GPU (block length 32, generation length 2048; FlowBlock uses W=2 with per-task spawn thresholds). Best value per benchmark in bold.

Benchmark	LLaDA-2.0				LLaDA-2.1				FlowBlock
	Acc	TPF	TPS	Lat	Acc	TPF	TPS	Lat	Acc	TPF	TPS	Lat
GSM8K	92.49	2.04	82.6	4.17	92.49	4.46	176.6	2.05	92.65	6.73	254.2	1.34
MATH500	73.40	2.59	104.6	11.25	77.00	5.34	211.7	5.40	76.20	8.86	332.1	3.44
Minerva-Algebra	91.24	3.00	120.9	6.50	93.60	6.29	247.6	2.87	93.51	10.14	378.9	1.81
ASDIV	92.36	2.03	81.9	2.94	92.89	4.43	174.5	1.50	92.80	6.36	239.7	1.01
HumanEval	84.76	4.37	159.6	2.16	82.93	4.93	180.1	1.79	86.59	8.41	279.4	1.15
HumanEval+	79.27	4.37	162.5	2.13	77.44	4.93	182.8	1.76	82.32	8.41	282.4	1.14
MBPP	79.16	2.75	105.1	2.24	82.20	3.26	124.2	2.09	83.61	4.43	162.1	1.53
MBPP+	84.39	2.78	108.7	2.18	87.04	3.30	125.8	1.96	88.36	4.47	161.3	1.39
Average	84.63	2.99	115.7	4.20	85.70	4.62	177.9	2.43	87.00	7.23	261.2	1.60

_{Acc = accuracy (%) · TPF = tokens per forward pass · TPS = tokens per second · Lat = per-request latency (s).}

On average FlowBlock reaches 87.00 accuracy — 1.3 points above LLaDA-2.1 (85.70) — while delivering 1.45× / 2.35× the throughput of LLaDA-2.1 / LLaDA-2.0 (up to 1.57× / 3.17×).

Batched serving

Math benchmarks — batched serving (TPS ↑ / latency ↓ vs. batch size):

Code benchmarks — batched serving:

📁 Repository Layout

FlowBlock
├── run_eval.sh                        # ← user-facing evaluation entry point
├── requirements.txt
├── assets/                            # logo & result figures
└── dInfer/
    ├── python/dinfer/
    │   ├── decoding/                  # FlowBlock + LLaDA-2.1 decoders
    │   │   ├── generate_uniform.py    #   GWD: gated wavefront, joint M2T/T2T refine, in-order commit
    │   │   ├── flowblock_packed.py    #   HWP: per-request wavefronts packed into dense forwards
    │   │   ├── parallel_strategy.py   #   threshold / joint-threshold decode policies
    │   │   ├── diffusion_runner.py    #   SGLang model runner (compile / CUDA-graph)
    │   │   └── utils.py               #   KV cache, block iterators, token arrays
    │   ├── decoding_llada_origin/     # original LLaDA-2.0 decoding (serial baseline)
    │   └── model/                     # LLaDA-2 MoE modeling (HF + SGLang)
    └── evaluations/
        ├── run_flowblock_joint_batch_compare.sh       # comparison harness (used by run_eval.sh)
        ├── run_flowblock_joint_batch_compare_loop.sh  # multi-config sweep driver
        ├── eval_dinfer_sglang.py                      # lm-eval entry with dInfer backend
        ├── prepare_task_subset.py                     # seed-fixed subset & GPU sharding
        ├── aggregate_dp8_results.py / aggregate_global_results.py
        ├── val_gsm8k.py / val_math.py / val_asdiv.py / val_code.py
        └── tasks/                                     # lm-eval task configs

🙏 Acknowledgements

This project is built upon DMax and the dInfer inference framework, and uses SGLang, lm-evaluation-harness, and the LLaDA-2.x model family. Many thanks to the authors of these projects.

📄 License & Citation

Released under the Apache 2.0 License.

The FlowBlock paper is currently under review; a citation entry will be added once it is public.