TD-Labs/Qwen3-4B-DCOCO

DCOCO drafter for Qwen/Qwen3-4B

This repository contains the DCOCO (DFlash Contextual Corrector) speculative-decoding drafter for Qwen/Qwen3-4B, from the paper Contextual Correction of Parallel Drafts for Speculative Decoding. It is a draft model: serve it together with Qwen/Qwen3-4B, which verifies every token, so the output distribution is exactly the target model's.

How it works

Diffusion-style drafters such as DFlash propose a whole block of tokens in a single pass, but each position is predicted without seeing which tokens were chosen before it. DCOCO keeps the top-16 candidates at every drafted position and scores how well each candidate fits each candidate at the next position, using candidate representations and a pair state over adjacent slots, all in one small batched pass. A prefix-first walk over those scores then picks the draft. This restores intra-block dependencies at a cost of a few percent of the drafting cycle.

Model details

Target (verifier) model	Qwen/Qwen3-4B
Architecture	DCOCODraftModel: 5-layer DFlash drafter (hidden size 2560, bidirectional in-block attention) plus the DCOCO corrector
Corrector	K = 16 candidates per slot, candidate width 128, pair-state width 128, 2-layer scoring MLP (hidden 256); 1,443,073 parameters
Draft block	15 drafted tokens per round (anchor position predicts); up to 16 tokens committed per verification
Target layers read	1, 9, 17, 25, 33
Training data	~800K prompts from the stem, math and code splits of Nemotron Post-Training Dataset v2 and CodeAlpaca-20k, with continuations regenerated by Qwen/Qwen3-4B (greedy, thinking off, up to 4,096 tokens)
Training	6 epochs of DFlash drafter training, then 6 epochs of joint training with the corrector; AdamW, lr 6e-4, cosine schedule
Format	BF16 safetensors (2.63 GB), vLLM-native config

Usage

DCOCO runs in vLLM, through the speculative-decoding dflash method; vLLM selects the corrector from the checkpoint's DCOCODraftModel architecture. This requires a vLLM build with DCOCO support, such as the DCOCO image (vllm/vllm-openai:v0.31.0 with the DCOCO overlay).

vllm serve Qwen/Qwen3-4B \
  --speculative-config '{"method": "dflash", "model": "TD-Labs/Qwen3-4B-DCOCO", "num_speculative_tokens": 15}'

Check that the server log shows Resolved architecture: DCOCODraftModel. With any other architecture, vLLM serves the drafter as plain DFlash without the corrector. Then query it like any OpenAI-compatible server:

curl localhost:8000/v1/chat/completions -H 'content-type: application/json' -d '{
  "model": "Qwen/Qwen3-4B",
  "messages": [{"role": "user", "content": "How many positive whole-number divisors does 196 have?"}],
  "chat_template_kwargs": {"enable_thinking": false}
}'

num_speculative_tokens can be lowered (down to 1) but not raised above 15.

Results

All results below are from the paper (Tables 14, 15, 16, 17 and 20). Setup: one NVIDIA B200 GPU, vLLM 0.27, concurrency 1 unless noted, thinking off with at most 2,048 generated tokens (4,096 with thinking on), 128 prompts per benchmark (all 30 of AIME-25, all 80 of MT-Bench, 880 of SPEED-Bench). Throughput (tok/s) is the median per-request rate (output tokens / request latency, including prefill); speedup is relative to the target model alone. τ is the mean accepted length per verification round, including the bonus token. T=0 is a single greedy run; T=1 is the mean over seeds 42, 43 and 44.

Per-benchmark results (thinking off)

Greedy decoding (T=0) in the first four columns, sampling (T=1) in the last four.

Benchmark	No-spec tok/s	DCOCO tok/s	Speedup	τ	No-spec tok/s	DCOCO tok/s	Speedup	τ
MATH-500	366.2	2165.5	5.91×	7.967	365.2	1644.7	4.50×	6.041
AIME-25	364.4	2085.9	5.72×	7.368	364.0	1194.1	3.28×	4.623
GSM8K	369.6	1638.4	4.43×	6.082	368.6	1525.7	4.14×	5.610
HumanEval	370.6	1475.8	3.98×	5.351	370.1	1372.1	3.71×	4.949
MBPP	370.2	1435.3	3.88×	5.297	369.4	1260.7	3.41×	4.603
LiveCodeBench	369.2	1754.1	4.75×	7.127	369.2	1301.1	3.52×	5.110
MT-Bench	369.0	661.3	1.79×	2.946	369.0	639.1	1.73×	2.669
Arena-Hard-v2	369.6	692.3	1.87×	2.716	369.9	631.4	1.71×	2.393
Alpaca	368.1	658.7	1.79×	2.515	367.2	631.3	1.72×	2.357
SPEED-Bench	369.1	759.0	2.06×	2.753	368.8	668.3	1.81×	2.425
Average	368.6	1332.6	3.62×	5.012	368.1	1086.8	2.95×	4.078

Comparison with other parallel drafters (average over the ten benchmarks)

Thinking off:

Method	tok/s (T=0)	Speedup (T=0)	τ (T=0)	tok/s (T=1)	Speedup (T=1)	τ (T=1)
DFlash	1096.9	2.98×	4.162	927.6	2.52×	3.476
Domino	1202.7	3.27×	4.861	995.7	2.71×	3.977
DSpark	980.1	2.66×	3.889	838.4	2.28×	3.281
DCOCO	1332.6	3.62×	5.012	1086.8	2.95×	4.078

Thinking on (extended reasoning):

Method	tok/s (T=0)	Speedup (T=0)	τ (T=0)	tok/s (T=1)	Speedup (T=1)	τ (T=1)
DFlash	753.5	2.08×	2.778	673.5	1.87×	2.460
Domino	785.2	2.17×	3.089	696.7	1.93×	2.702
DSpark	670.6	1.85×	2.565	610.2	1.69×	2.313
DCOCO	862.2	2.38×	3.129	765.3	2.12×	2.741

Throughput under concurrent requests (thinking off)

Aggregate output tokens/s, averaged over the ten benchmarks.

Concurrent requests	DFlash	Domino	DSpark	DCOCO	DCOCO vs. best other
1 (T=0)	1096.9	1202.7	980.1	1332.6	+10.8%
4 (T=0)	4097.1	4032.1	3517.5	4660.3	+13.7%
8 (T=0)	6942.6	6560.9	5953.9	7781.7	+12.1%
16 (T=0)	10940.3	9633.9	9448.6	12215.5	+11.7%
32 (T=0)	14181.1	11966.5	12746.7	15602.7	+10.0%
1 (T=1)	927.6	995.7	838.4	1086.8	+9.1%
4 (T=1)	3345.6	3226.4	2929.6	3721.2	+11.2%
8 (T=1)	5697.0	5306.7	5015.9	6191.9	+8.7%
16 (T=1)	9092.8	7830.4	7965.6	9834.0	+8.2%
32 (T=1)	11676.9	9628.3	10525.8	12476.2	+6.8%

Limitations

• The drafter is trained for Qwen/Qwen3-4B and must be served with exactly that target model.
• It needs the target's full vocabulary (no reduced draft vocabulary).
• It was trained on thinking-off generations; the speedup is lower with extended thinking (see above).
• The speedup depends on the task: highest on math and code, lower on open-ended chat.

Citation

@article{eisenstadt2026dcoco,
  title  = {Contextual Correction of Parallel Drafts for Speculative Decoding},
  author = {Eisenstadt, Roy and Lidar, Harel and Gonen, Tal and Luxembourg, Omer and
            Zimerman, Itamar and Wolf, Lior and David, Eli and Naiman, Ilan},
  year   = {2026}
}