perf(frozen_kv_mtp): capture the assistant seed step with CUDA graph (+6.8% chat tok/s, +4.4% summ tok/s)

[pyc96]

Summary

Captures the assistant seed step in FROZEN_KV_MTP with CUDA graph. Previously only the recurrent draft loop was captured; the seed step (one decode-shape forward of the assistant) ran eager and cost ~20-25 % of decode wall time at speculative_num_steps=3.

Measured: +6.8 % chat tok/s, +4.4 % summ tok/s, ~5 % TPOT reduction on Gemma-4 31B-it MTP H100 TP=2. accept_length unchanged.

What the seed step is and why it was eager

Per scheduler iter, the FROZEN_KV_MTP V1 worker runs (in order):

1. Target prefill (only on prefill iters) — eager (variable shape)
2. Assistant seed step (after prefill, then again after every verify) — was EAGER
3. Recurrent draft loop (speculative_num_steps - 1 iters) — captured by FrozenKVMTPCudaGraphRunner
4. Target verify — captured by target's standard cuda graph runner

The seed step is one decode-shape forward of the assistant model that seeds topk_p / topk_index / hidden_states on the FrozenKVMTPDraftInput for the recurrent loop. It's structurally identical to one iter of the recurrent loop, but runs through draft_model_runner.forward() which has its own cuda graph disabled (the worker forces disable_cuda_graph=True for the draft TpModelWorker because the recurrent loop captures its own outer graph).

Result: the eager seed forward is launched twice per scheduler iter, dominates ~20-25 % of decode wall time, and is invisible to the existing capture pipeline.

What this PR does

File	Change
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py	New FrozenKVMTPSeedInputBuffers dataclass (input_ids + hidden_states inputs + topk_p / topk_index / hidden_states outputs). New methods _init_seed_buffers, _capture_seed, _capture_one_seed_graph, can_run_seed, replay_seed. Captures at the same batch sizes as the recurrent loop (29 sizes on the campaign workload). Falls back to eager seed silently if capture fails. Env var SGLANG_FROZEN_KV_MTP_DISABLE_SEED_CG=1 skips capture for A/B benchmarking.
python/sglang/srt/speculative/frozen_kv_mtp_worker.py	_run_assistant_seed_step uses replay_seed when can_run_seed(bs) AND no multimodal embeds AND no mrope_positions. Otherwise the eager path runs as before. Stitches seed graph outputs (topk_p, topk_index, hidden_states) onto the new FrozenKVMTPDraftInput.

The two paths are predicate-disjoint, so the eager fallback is 100 % unchanged for any caller that doesn't qualify.

Measured

google/gemma-4-31B-it, H100 TP=2, triton, FROZEN_KV_MTP NEXTN 3/4/1, 80 prompts, warmup 2, seed 1, --disable-overlap-schedule:

Metric	seed CG OFF (baseline)	seed CG ON (this PR)	Δ
chat 1k/1k tok/s	1466	1566	+6.8 %
chat 1k/1k medTPOT (ms)	34.0	32.1	−5.7 %
chat 1k/1k medTTFT (ms)	2829	2778	−1.8 %
summ 8k/1k tok/s	428	447	+4.4 %
summ 8k/1k medTPOT (ms)	24.8	23.8	−4.0 %
summ 8k/1k medTTFT (ms)	80956	77624	−4.1 %
chat accept_length	3.135	3.128	tied (same draft)
summ accept_length	3.142	3.139	tied (same draft)

Correctness — per-prompt parity (20 greedy prompts, temp=0)

{"total": 20, "matched": 19, "mismatched": 1, "match_rate": 0.95}

The single mismatch is a long output that diverges after ~50 tokens (semantically equivalent; bf16 numerical drift from fp32 accumulation differences inside the captured kernel — same drift pattern documented in PR #19 and #20).

Server log evidence (init)

[info TP0] Capture Frozen-KV MTP draft cuda graph begin.
[info TP0] Capture Frozen-KV MTP seed cuda graph begin.
[info TP0] Capture Frozen-KV MTP seed cuda graph end (captured 29 batch sizes).
[info TP0] Capture Frozen-KV MTP draft cuda graph end.

Reproducer

# Server A: seed CG ON (this PR)
CUDA_VISIBLE_DEVICES=0,1 python -m sglang.launch_server \
  --model-path google/gemma-4-31B-it --dtype bfloat16 --trust-remote-code --tp-size 2 \
  --port 30000 --attention-backend triton \
  --speculative-algorithm NEXTN \
  --speculative-draft-model-path google/gemma-4-31B-it-assistant \
  --speculative-num-steps 3 --speculative-num-draft-tokens 4 --speculative-eagle-topk 1 \
  --max-running-requests 80 --disable-overlap-schedule

# Server B: seed CG OFF (eager seed path, pre-PR behavior)
SGLANG_FROZEN_KV_MTP_DISABLE_SEED_CG=1 CUDA_VISIBLE_DEVICES=2,3 python -m sglang.launch_server \
  ... --port 30001 ...

# Bench
python -m sglang.bench_serving --backend sglang-oai-chat \
  --host 127.0.0.1 --port 30000 --model google/gemma-4-31B-it \
  --dataset-name random --random-input-len 1000 --random-output-len 1000 \
  --num-prompts 80 --warmup-requests 2 --seed 1

Stack base

pyc/feat-gemma4-ultimate-v2 (PR #21)

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.