[jason-fxz]

Refs #828

Second PR of the speculative decoding plan in #828: the target side of a verify step, driven by the scheduler, on top of the multi-row seams from #833. No draft method is registered yet (the MTP PR adds the first one), so nothing changes unless a proposer is plugged in; the tests plug in a stub. The full stack is in the draft PoC #832.

Summary

With --speculative-algorithm set, every decode step becomes a verify step: each running request gets k draft tokens from a Proposer, the target scores the k+1 rows in one forward, and the scheduler commits the accepted drafts plus one target token.

New flags: --speculative-algorithm, --speculative-num-steps (k), --speculative-draft-model-path.

Design

scheduler        decode batch -> verify batch: k+1 rows per request, slots allocated for all of them
Proposer         k drafts per request, written into token_pool after the pending token
target forward   one uniform extend over the k+1 rows (the verify graph family when captured)
rejection        sample every row as decode would; accept the prefix that matches the drafts
drain            commit n = accepted + 1 tokens; roll back the unused slots; replay the accepted rows into the GDN state

• The engine and the scheduler import only engine/spec_protocol.py. spec/runner.py (SpecRunner) is the target side and the order of a step; spec/proposer.py is the draft-side interface. Method code goes under spec/<method>/ in later PRs.
• Drafts are argmax, so sampling each verify row with the normal decode sampler and accepting the matching prefix is exact rejection sampling (spec/rejection.py). Each row samples with its request's temperature, top-k, top-p and greedy flag. Presence and frequency penalties are applied per row, with the drafts before a row counted as already generated.
• GatedDeltaNet state is only read during a verify. Each row's recurrence inputs and conv window go into a LinearVerifyScratch owned by the state pool and handed to the GDN layers through FLAMetadata.verify_scratch. The drain replays the accepted rows into the live state (gdn_state_replay, a Triton kernel whose result matches stepwise decode bitwise) and restores the conv window, before a finished request gives its slot back.
• When decode steps have more than one row, the target captures the verify graph family in place of the decode one. With speculative decoding on, the scheduler uses the normal loop (no overlap scheduling).
• Memory: a proposer declares fixed bytes and draft KV bytes per target KV token; both are taken out of the KV pool budget, counted by the rebuild check and reported in /v1/cache/status (geometry.unit_bytes.kv_per_token). The verify scratch is counted with the GDN state pool.

Rejected at startup: TP > 1; targets with sliding-window or DSV4 window-tier KV; decode attention backends that cannot capture multi-row graphs; proposers that sample their drafts. With --attention-backend X,fi, where X captures multi-row graphs on device, the backend becomes X and the change is logged.

Tests

• tests/spec/test_verify_contract.py: a tiny hybrid (GDN + attention) model with dummy weights, served end to end through the scheduler with a stub proposer whose drafts come from a plain greedy run and are wrong at a different place each step (so every accept count occurs). The speculative run must reproduce the greedy run token for token, except at a near tie, which the test checks is one. fa and triton, with and without CUDA graphs.
• tests/spec/test_gdn_verify.py: a GDN verify gives the same outputs as stepwise decode, leaves the live state untouched, and after the commit matches decode after n tokens for every n; the stash and replay kernels match decode bitwise, also for a launch with more layer x row x head programs than one CUDA grid axis holds.
• tests/spec/test_rejection.py: greedy prefix match; each verify row samples with its own request's parameters (top-k, greedy); each verify row counts the drafts before it as committed for penalties.
• tests/scheduler/test_spec_drain.py: what a verify drain commits, which slots it gives back, and that the target state is committed before a finished request frees its GDN slot.
• tests/kvcache/test_cache_unit_bytes.py: kv_bytes_per_token counts the draft KV.

Test environment

• GPU: 1x NVIDIA H100 80GB HBM3, driver 580.95.05
• CPU: Intel Xeon Platinum 8480+ (26 vCPUs), 221 GB RAM, Ubuntu 22.04.5
• CUDA 13.0, torch 2.14.1+cu130, Python 3.12.12
• This branch at 4bc0185, on main e328cf6. The end-to-end runs below used c20f4014; the only later code change restores the GDN conv window with one indexed copy instead of one copy per request (unit tests rerun on 4bc0185).

Unit tests

FREETOKEN_QWEN36_MODEL=/path/to/Qwen3.6-27B PYTHONPATH=python pytest tests/spec tests/scheduler tests/engine tests/kvcache tests/attention tests/moe/test_offload.py tests/server tests/tokenizer tests/models

FREETOKEN_QWEN36_MODEL gives the contract test its tokenizer. 1686 passed, 202 skipped, 1 failed. The failure is tests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensors[apodex/Apodex-1.1-mini-NVFP4[index-only]] (a local checkpoint with quant_method: modelopt_mixed); it fails the same way on main.

End-to-end with speculative decoding off

No draft method is registered here, so serving without the new flags must match main. AIME25 problem 1 through freetoken.llm.LLM(model_path, dtype=torch.bfloat16, max_running_req=4) with the chat template, greedy (temperature=0, max_tokens=16384, natural end), compared token by token with the same run on main 9221fd7 (the commit before #833 merged):

Checkpoint	main 9221fd7	this branch	Token by token
Qwen/Qwen3.6-35B-A3B (GDN hybrid, MoE)	16384 tokens (hit the limit)	16384 tokens (hit the limit)	identical
Qwen/Qwen3-30B-A3B (full attention, MoE)	4709 tokens, answer 70	4709 tokens, answer 70	identical

Smoke test on this branch, passed:

FREETOKEN_TEST_MODEL=/path/to/Qwen3.6-35B-A3B FREETOKEN_AIME25_JSONL=/path/to/aime25.jsonl PYTHONPATH=python pytest tests/e2e/test_aime.py

Known gaps

• No draft method is registered here; speed and MTP on/off results come with the MTP PR.
• The contract test uses a dense hybrid model; MoE targets are not covered by a test in this PR.
• The verify step's sampling memory (logits for k+1 rows per request) is not reserved in the KV budget.
• A linear-attention layer without a verify path (the glm5_next KDA layer) is kept out only because its attention backend cannot capture verify graphs; there is no direct check.
• With speculative decoding, the tool-call anchor GDN snapshot is taken only when a step starts exactly at the anchor, so it is usually missed.
• Proposer.slot_states (per-request draft state on the GDN slots) is declared but not yet wired into the pool or the budget; no proposer uses it yet.