[jason-fxz]

Refs #828

First PR of the speculative decoding plan in #828: the engine and scheduler seams for decode steps that run several query rows per request. Nothing in this PR runs more than one row: every new knob defaults to one row per request, so serving behaves as on main. The full stack that uses these seams is in the draft PoC #832.

Summary

A speculative verify step runs k+1 rows per request through the target and then commits up to k+1 tokens. Today the engine and the scheduler assume one row and one token per request per decode step in several places. This PR turns those assumptions into explicit parameters, all fixed at 1 for now:

• EngineConfig.decode_rows_per_req (always 1 here) drives every per-request decode allowance: the scheduler's admission and in-flight token headroom, the longest request (rows - 1 short of the page-table width, so a widened step still fits), the CPU MoE executor's token capacity, and with GPU decode the MoE slot pools (the decode-batch clamp, the startup split and the rebuild check).
• Batch.rows_per_req (1 for decode) and a verify phase; a verify batch must set rows_per_req (asserted). The scheduler takes one GDN slot and one page-table row per request, every rows_per_req rows.
• Req.widen_pending(k) widens a request by k provisional rows; Req.commit_pending(n) and CacheManager.rollback_tail(...) keep the committed rows and free the rest.
• The scheduler's per-token commit moves into _commit_tokens(req, tokens), which appends tokens in order and stops at the first one that finishes the request (EOS, stop string or length, same rules as before). DetokenizeMsg.next_token: int becomes next_tokens: List[int]; the detokenizer, LLM and the tokenizer server read the list.

Nothing in this PR calls the multi-row pieces yet. The verify PR is the first caller of widen_pending / commit_pending / rollback_tail, the verify phase, Batch.features and FeatureSink, GraphRunner with more than one row per request, BackendInfo.uniform_extend_graph, report_batch(generated=) and the count _commit_tokens returns; the MTP PR is the first caller of Context.override with arguments and of GraphRunner with several outputs. The tests below cover multi-row attention graphs against the eager path and the pages a rollback keeps; the other pieces are exercised by their callers in those PRs (the verify PR's contract test runs multi-row CUDA graphs on fa and triton).

Design

• GraphRunner captures any forward callable for rows_per_req rows per request instead of the model's forward for one row. Its output buffers are sized from the first (largest) capture, so it is no longer tied to [bs, vocab] logits. Engine graphs pass forward=self.model.forward, out_dtype=float32, i.e. what they captured before. The other new keyword arguments (phase, ctx_override, batch_hook, pool) are unused by the engine here; they let a draft capture its own graphs on its own backend, in the target's graph memory pool.
• init_capture_graph(max_seq_len, bs_list, rows_per_req=1) on every attention backend. fa, triton and flashinfer capture rows_per_req > 1 (flashinfer through its prefill wrapper); dsa (and the GLM-5 Next backend built on it), dsv4, dsv41, m3, qsa and trtllm raise NotImplementedError for rows_per_req != 1. BackendInfo.uniform_extend_graph records which is which.
• Context.override(kv_cache=, attn_backend=, hide_target_state=) lets a draft forward run on its own KV pool and backend with the target's GDN pool hidden; forward_batch(batch, **override) scopes it with the batch.
• Engine.forward_logits(batch) is split out of forward_batch, and ParallelLMHead.logits(x) out of the LM head forward. The LM head hands its input rows to batch.features (a FeatureSink, None unless a draft asks for them).

Tests

New tests:

• tests/attention/test_capture_rows_per_req.py: attention graphs captured with 1 and 4 rows per request replay what the eager path computes on fa, triton and flashinfer.
• tests/scheduler/test_rollback_tail.py: after a step widened by 3 rows commits 1..4 tokens, the request keeps exactly the pages of its committed tokens (page sizes 1 and 4) and the next step allocates no page twice.

Existing tests change only for the renamed fields and the new signatures (next_token -> next_tokens, GraphRunner(forward=...), init_capture_graph(..., rows_per_req=1) on a fake backend, _commit_tokens on a scheduler stub, decode_rows_per_req on a stub engine config); tests/engine/test_cache_budget.py gains a case with two rows per request for the GPU-decode MoE pool clamp.

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 c21b961 (main 9221fd7 merged in), A/B against main 9221fd7. After it, 3b8f873 changes two comments only, 6d52946 scales the GPU-decode MoE pool limits by decode_rows_per_req (no change at 1), 1cb80f7 removes the unread Context.moe_offload_cache field (no behavior change), and the last commit trims the new tests. Unit tests rerun on the last commit.

Unit tests

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

1672 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 A/B against main

AIME25 problem 1 through freetoken.llm.LLM(model_path, dtype=torch.bfloat16, max_running_req=4) with the chat template, on main 9221fd7 and on this branch.

Greedy, one request at a time (temperature=0, max_tokens=16384, natural end), compared token by token:

Checkpoint	main	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

Decode speed: the checkpoint's recommended sampling from generation_config.json (Qwen3.6-35B-A3B: T=1.0, top_p=0.95, top_k=20; Qwen3-30B-A3B: T=0.6, top_p=0.95, top_k=20), 1024 decode tokens with ignore_eos, a 32-token warmup, median of 3 runs:

Checkpoint	main tok/s	this branch tok/s
Qwen/Qwen3.6-35B-A3B	159.0	159.0
Qwen/Qwen3-30B-A3B	168.0	168.0

Batch: problems 1-8 of AIME25 submitted together (max_running_req=4, so requests 5-8 queue until slots free up), Qwen/Qwen3.6-35B-A3B, recommended sampling, max_tokens=4096. All 8 requests succeed on both sides and run to 4096 tokens; on both sides the output reads coherently, and the final 300 characters of every request are the same on main and on this branch. Wall time 62.5 s on both.

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