[jason-fxz]

Refs #828

Adds DFlash 2 as a second speculative decoding method next to MTP (#844). A separate draft checkpoint reads the outputs of a few target layers and drafts up to block_size - 1 tokens in one forward (7 for the released drafts). The qwen3_5_moe and qwen4_exp families (HF and GGUF) accept any DFlash 2 draft that fits the target; tested with the three released drafts on their own targets: incoai/Qwen3.8-27B-DFlash2, incoai/Qwen3.6-35B-A3B-DFlash2 and incoai/Qwen3.8-Flash-Next-DFlash2.

ft serve --model /path/to/Qwen3.6-35B-A3B --speculative-algorithm dflash --speculative-draft-model-path /path/to/Qwen3.6-35B-A3B-DFlash2

Claims

• DFlash k=7 is 1.49x / 1.32x / 1.16x as fast as the faster MTP setting on Qwen3.8-27B / Qwen3.6-35B-A3B / Qwen3.8-Flash-Next-NVFP4 (one request, AIME25 problem 1, recommended sampling).
• Speculative decoding off and MTP produce the same greedy tokens as main.
• The verify keeps the target's output distribution for sampled drafts; greedy DFlash follows greedy decoding up to a bf16 near tie (checked on 27B and 35B-A3B).
• On a prefix-cache hit the draft's context is restored bitwise.

Validation

H100 80GB, driver 580.95.05, Intel Xeon Platinum 8480+ (26 vCPUs), CUDA 13.0, torch 2.14.1+cu130; this branch at 0f48b9d, main at 9cacd08. The last commit, a033e53, only drops redundant tests; the trimmed test files pass on it (62 passed, 4 skipped). FreeToken's CI does not run the test suite, so everything below ran locally.

Accuracy

Check	Setup	Result
Off and MTP vs main	greedy AIME25 problem 1 (up to 16384 tokens); off on all three targets, MTP k=3 on 35B-A3B	token for token identical
DFlash vs off	greedy, off pinned to fa (the backend speculative decoding switches to)	27B k=3 identical; 27B k=7 parts after 402 tokens and 35B-A3B k=3 / k=7 after 27, both bf16 near ties (top-2 logits within one ulp, checked on 4e09a28b, same outputs); Flash-Next parts after 154, not checked
AIME25 problem 1 answer	sampled, 5 runs per setting (table below)	DFlash 5/5 on every target and k; the only misses (off and MTP k=3 on 35B-A3B) hit the 16384-token limit
Sampled verify distribution	test_rejection.py, Monte Carlo over 6M rows	committed-token marginals within 5e-3 of the target distribution
Prefix-cache restore	35B-A3B; hits on a track snapshot, a finish snapshot, a finish after 2400 decoded tokens; with and without graphs	restored window bitwise equal to the snapshot; an eager shadow of every window write finds no rejected, stale or unwritten row
Smoke	test_aime.py with DFlash on Flash-Next; test_cache_rebuild.py with DFlash on 35B-A3B	both pass (5/5; every rebuild repeats the pre-rebuild output)

Speed

One request, AIME25 problem 1, recommended sampling (temperature=1.0, top_p=0.95, top_k=20), natural end. Each cell is the median of 5 runs of output tokens / generate() wall time, with the median tokens committed per verify step in parentheses.

Target	off	MTP (faster of k=3, k=7)	DFlash k=7	vs MTP	vs off
Qwen/Qwen3.8-27B (bf16)	50.1	144.4 (k=7, 4.51)	215.0 (5.47)	1.49x	4.29x
Qwen/Qwen3.6-35B-A3B (bf16)	193.1	359.6 (k=3, 3.25)	475.7 (5.39)	1.32x	2.46x
nvidia/Qwen3.8-Flash-Next-NVFP4	92.9	149.7 (k=3, 3.16)	174.2 (4.93)	1.16x	1.87x
Flash-Next GGUF IQ3_XXS	103.6	not run	173.5 (4.51)	-	1.67x

On the GGUF, DFlash k=3 is faster: 179.2 (3.23), 1.73x off.

All settings, main vs this branch, and memory

Each run is freetoken.llm.LLM(model, dtype=torch.bfloat16, max_running_req=4) with the chat template, one prompt, SamplingParams(max_tokens=16384, temperature=1.0, top_p=0.95, top_k=20), and speculative_algorithm / speculative_num_steps / speculative_draft_model_path per row. The GGUF is ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF (IQ3_XXS). Off and MTP ran on main 9cacd08 and on this branch; each pair produced the same token count in every run. KV tokens are what the engine sized at startup with default memory settings, with MoE cache slots in parentheses (35B-A3B and the GGUF cache every expert, Flash-Next NVFP4 22483 of 24576 when off).

Target	Setting	main	this branch	tokens/step	correct	KV tokens
Qwen3.8-27B	off	50.2	50.1	-	5/5	253435
	MTP k=3	123.9	124.0	3.10	5/5	225733
	MTP k=7	144.4	144.4	4.51	5/5	224687
	DFlash k=3	-	140.7	3.40	5/5	182950
	DFlash k=7	-	215.0	5.47	5/5	181742
Qwen3.6-35B-A3B	off	193.1	193.1	-	3/5	224198 (10240)
	MTP k=3	361.3	359.6	3.25	4/5	127668 (10240)
	MTP k=7	348.7	348.8	4.52	5/5	126088 (10240)
	DFlash k=3	-	378.5	3.52	5/5	144146 (10240)
	DFlash k=7	-	475.7	5.39	5/5	142110 (10240)
Qwen3.8-Flash-Next-NVFP4	off	93.0	92.9	-	5/5	8256 (22483)
	MTP k=3	150.1	149.7	3.16	5/5	8256 (21491)
	MTP k=7	146.0	146.1	4.42	5/5	8256 (21471)
	DFlash k=3	-	165.0	3.34	5/5	8192 (21793)
	DFlash k=7	-	174.2	4.93	5/5	8192 (21771)
Flash-Next GGUF IQ3_XXS	off	-	103.6	-	5/5	1019072 (24576)
	DFlash k=3	-	179.2	3.23	5/5	941696 (24576)
	DFlash k=7	-	173.5	4.51	5/5	939264 (24576)

The runs that miss the answer all hit the 16384-token limit. Flash-Next NVFP4's KV is at its 8192-token reserve floor with DFlash, so its draft takes expert cache slots instead.

The draft stays resident: 1924M / 526M / 603M parameters (27B / 35B-A3B / Flash-Next drafts), plus its K/V ring and a window of 2112 x hidden bf16 on each GDN slot (25 slots here: 516 / 206 / 258 MiB).

Reproduce

# unit tests (1745 passed, 211 skipped, 1 failed: test_quant_config on a local modelopt_mixed checkpoint, which fails the same way on main)
FREETOKEN_QWEN36_MODEL=/path/to/Qwen3.6-35B-A3B PYTHONPATH=python pytest tests/spec tests/layers tests/scheduler tests/engine tests/kvcache tests/attention tests/moe/test_offload.py tests/server tests/tokenizer tests/models tests/kernels/test_dflash.py

# smoke tests
FREETOKEN_TEST_MODEL=/path/to/Qwen3.8-Flash-Next-NVFP4 FREETOKEN_AIME25_JSONL=/path/to/aime25.jsonl FREETOKEN_AIME_SAMPLES=5 FREETOKEN_TEST_SPEC_ALGORITHM=dflash FREETOKEN_TEST_SPEC_DRAFT_PATH=/path/to/Qwen3.8-Flash-Next-DFlash2 PYTHONPATH=python pytest -s tests/e2e/test_aime.py
FREETOKEN_REBUILD_TEST_MODEL=/path/to/Qwen3.6-35B-A3B FREETOKEN_REBUILD_EXTRA_ARGS="--speculative-algorithm dflash --speculative-draft-model-path /path/to/Qwen3.6-35B-A3B-DFlash2" PYTHONPATH=python pytest tests/e2e/test_cache_rebuild.py

# one speed run: AIME25 problem 1 through the chat template, then tok/s = len(token_ids) / wall time
from freetoken.core import SamplingParams
from freetoken.llm import LLM

llm = LLM(model_path=MODEL, dtype=torch.bfloat16, max_running_req=4,
          speculative_algorithm="dflash", speculative_draft_model_path=DRAFT, speculative_num_steps=7)
ids = llm.tokenizer.encode(llm.tokenizer.apply_chat_template([{"role": "user", "content": problem}], add_generation_prompt=True, tokenize=False), add_special_tokens=False)
out = llm.generate([ids], SamplingParams(max_tokens=16384, temperature=1.0, top_p=0.95, top_k=20))

Tests added or changed

File	Checks against	Needs
tests/kernels/test_dflash.py	ring attention vs a gathered PyTorch reference; the selector never draws a q = 0 candidate	CUDA
tests/models/dflash/test_draft.py	draft forward (ring K/V, grouped conv, norm + rope) and selector walk vs an fp32 reference of the DFlash 2 math; a ring rebuilt from the window drafts the same block bitwise	CUDA
tests/spec/test_dflash_window.py	a verify's rejected rows and a prefill's rows past a tracked boundary stay out of every window read	-
tests/spec/test_rejection.py	sparse-q rejection vs the target distribution (Monte Carlo, with trusted, NaN, wrong-total and greedy rows)	CUDA
tests/spec/test_verify_contract.py	three batched greedy requests with a random-weight DFlash draft vs plain decoding (equal except at fp16 near ties), qwen3_5 and qwen4_exp, with and without graphs	CUDA, FREETOKEN_QWEN36_MODEL (tokenizer)
tests/kvcache/test_linear_state_pool_alloc.py, tests/models/qwen4_exp/test_ple.py	the window rides a slot copy; its bytes; track_src	-
tests/e2e/test_cache_rebuild.py, tests/e2e/test_aime.py	a rebuilt engine repeats the pre-rebuild greedy output (except with --moe-strategy hybrid); FREETOKEN_TEST_SPEC_* knobs	weights

Details

The released DFlash 2 drafts use sliding-window attention (window 2048).

• Runtime: a KV ring per request. The draft's K/V lives in a ring of 2112 positions: the window plus 64 rows of room for a verify's rejected rows and for prefill rows past a GDN snapshot boundary. The prefill and every verify turn the outputs of the target layers the draft reads into the draft's hidden-sized context rows (its fc + norm) and write their K/V into the ring. One draft forward over [last token, mask x 7] reads it; the target LM head gives 16 candidates per slot, and a selector picks one path (sampled from its distribution q, or argmax when greedy).
• Prefix cache: the context rows as a slot state. The same context rows, before the per-layer K/V projection, are also kept as a dflash_window slot state on the target's GDN state slot. They are snapshotted and restored with the GDN state, and on a prefix hit the ring is rebuilt from the last 2047 of them.
• Sampled verify: accept with min(1, p/q), on a reject resample from max(0, p - q), so the output follows the target distribution. All-greedy batches keep the prefix match of feat(spec): verify draft tokens on the target side #836.
• Shared changes: FeatureSink.layer_output gains a residual argument and is called after each layer the draft reads (qwen3_5_moe, qwen4_exp); Proposer.propose gets the pending tokens and returns a SpecProposal(tokens, cand, q); a proposer can declare slot states (EngineConfig.spec_slot_states); Sampler.probs and FLAMetadata.track_src are new.
• Refused at startup, from the draft's config.json: no draft path; a draft that is not DFlash2DraftModel (DFlash v1 included), or is quantized, causal, not all sliding-window, or uses a rope, partial or non-NeoX rotary, or output scaling this forward does not apply; selector sizes or head_dim that are not powers of two, a conv_group_size that does not divide the hidden size, or empty, repeated or negative target_layer_ids; k outside 1..block_size - 1 (the default is block_size - 1); a draft whose layers, hidden size, vocabulary or mask token do not fit the target; a target without GDN layers unless --cache-type naive.

Known gaps

• Not run end to end: concurrent requests (every run above serves one request; the contract test batches three on tiny models), Qwen3.5-family GGUF targets, other quantized targets, prompts beyond about 4k tokens, multimodal prompts, --moe-strategy hybrid / cpu.
• Prefill temporaries are not budgeted and use the memory_ratio headroom: the tapped features of the rows the draft keeps (at most 2112 per request; 256 MiB for an 8192-row chunk of four 2048-token prompts on 35B-A3B), and a prefix-hit restore that projects every restoring request's 2047 rows in one call, so it grows with max_running_req.
• One sparse-q verify takes 646 us against 266 us for MTP's argmax-draft verify (microbenchmark: CUDA graph replay, 4 requests, k = 7, top-k + top-p).
• Speculative decoding still needs --tp 1 and runs without overlap scheduling, as in feat(spec): verify draft tokens on the target side #836.