[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, the LinearAttnBackend from #837 and the shared GatedDeltaNet layer from #841. 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 (--spec-algo), --speculative-num-steps (--spec-num-steps, k), --speculative-draft-model-path (--spec-draft-model).

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 see speculative decoding only through engine/spec_protocol.py: SpecRuntime is a Protocol whose one implementation, SpecRunner (spec/runner.py), is loaded by name, so the engine never imports freetoken.spec. SpecRunner is the target side and the order of a step. A draft method implements Proposer (spec/proposer.py): propose returns k drafts per request, and the other hooks (feature sink, post-verify, prefill, graph capture, release) default to no-ops. 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, which LinearAttnBackend sets for eager and captured verify batches. The backend now holds the LinearStatePool it is built with, and build_fla_metadata takes the pool as an argument, so attention/linear.py no longer reads the state pool from the global context (the hybrid-radix track metadata read its conv width from there too). In layers/linear_attn/gdn.py a verify goes through the decode branch: the verify conv (causal_conv1d_verify) in place of the decode conv, and the same fused recurrent kernel with disable_state_update and input_stash. The drain (Engine.commit_verify) 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. 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/layers/test_gdn.py: the GDN layer test now continues its prefill and decode with a verify of k+1 rows, a commit at a different accept count per request (1, 2 and k+1), and another decode step, each checked against the pure-torch reference over the request's whole sequence so far; a verify that wrote the live state or a wrong commit shows up in the next step.
• tests/spec/test_gdn_verify.py: the stash and replay kernels match decode bitwise for every accept count, 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.
• tests/models/qwen4_exp/test_ple.py passes its pool to build_fla_metadata instead of installing a global context for it.

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 fe28362, on main d043fb6. The results below are from fa0fba5; the two later commits change a comment and add the short flag aliases (tests/server passes, 640 tests).

Unit tests

FREETOKEN_QWEN36_MODEL=/path/to/Qwen3.6-27B 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_kda.py

FREETOKEN_QWEN36_MODEL gives the contract test its tokenizer. 1695 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 a2ce636 (#841 reproduced these runs token for token, so they stand for main d043fb6 too):

Checkpoint	main a2ce636	this branch	Token by token
Qwen/Qwen3.6-35B-A3B (GDN hybrid, MoE)	16384 tokens (hit the limit)	16384 tokens (hit the limit)	identical
Qwen3.8-Flash-Next-NVFP4 (qwen4_exp)	727 tokens, answer 70	727 tokens, answer 70	identical

AIME25 problems 1-3 submitted together on Qwen/Qwen3.6-35B-A3B (max_running_req=4, cache_type=hybrid_radix, greedy, 1024 tokens each with ignore_eos): all three identical to main a2ce636 token by token.

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.
• qwen4_exp targets also need verify handling for the PLE per-request state; without it a verify commits all k+1 rows into the PLE n-gram context. The MTP PR adds it.
• With FREETOKEN_MAMBA_SSM_DTYPE set to bfloat16 or float16 (fp32 is the default), the replay keeps the recurrent state in fp32 across the accepted rows and rounds once, while decode rounds after every token, so the two are no longer bitwise equal.