[AMD][Spec] Prune unused EAGLE draft-extend LM-head rows

[chuyeh]

Motivation

During EAGLE draft extend, each request processes a fixed four-position window: one target-model token followed by three draft tokens. The draft model must process all four positions to update its KV cache, but only the hidden state at the last accepted position is used to predict the next draft token.

The existing path projects all four hidden-state rows through the large-vocabulary LM head and then discards three rows. This PR moves row selection before the LM head so only the required row is projected.

Modifications

• Publish the per-request selected draft-extend row through spec_info on eligible ROCm paths.
• Capture full hidden states for the next draft iteration while projecting only selected rows through the LM head.
• Size graph logits buffers for one row per request when pruning is active.
• Keep the established full-row path when a gathered buffer is required or on non-ROCm platforms.

The independent top-k1 softmax-skip optimization is intentionally handled by #35872.

Correctness

test_eagle_worker_v2_topk1_fastpath.py: 11 passed, 7 subtests passed

The tests cover hidden-state preservation, pruning eligibility, graph output sizing, and fallback behavior. A full InferenceX sweep completed every request at C4/C32/C64/C128.

Performance

• Hardware: 2× MI355X, TP2 per arm
• Image: rocm/sgl-dev:v0.5.18-rocm724-mi35x-20260825
• Model: amd/Qwen3.5-397B-A17B-MXFP4
• Workload: InferenceX 8192/1024 random lengths, range ratio 0.8
• AITER unified attention; EAGLE 3 steps / top-k 1 / 4 draft tokens; FP8 E4M3 KV

Conc.	Baseline tok/s	PR tok/s	Throughput	Baseline TTFT	PR TTFT	TTFT	Baseline TPOT	PR TPOT	TPOT	Median E2E
4	6,053.89	6,126.07	+1.19%	328.00 ms	327.55 ms	−0.14%	5.501 ms	5.379 ms	−2.22%	−1.42%
32	14,082.48	14,182.76	+0.71%	365.04 ms	361.19 ms	−1.06%	19.878 ms	19.718 ms	−0.81%	−1.74%
64	16,900.01	17,208.39	+1.82%	407.37 ms	495.03 ms	+21.52%	33.506 ms	33.005 ms	−1.50%	−2.40%
128	20,014.44	20,168.92	+0.77%	667.53 ms	658.09 ms	−1.41%	57.150 ms	56.865 ms	−0.50%	−0.70%

Median throughput improves by 0.98% and median TPOT by 1.15%. C32 was repeated after an inconsistent first throughput result; the repeat above aligns throughput with its TPOT and E2E changes. C64 TTFT is a prefill-scheduling outlier: decode TPOT, E2E latency, and throughput all improve at that point.

Checklist

• Keep the PR scoped to LM-head row pruning.
• Preserve gathered-buffer and non-ROCm fallbacks.
• Add focused unit coverage.
• Run end-to-end serving benchmarks.

---

CI States

Latest PR Test (Base): ❌ Run #33704304480
Latest PR Test (Extra): ❌ Run #33704304000
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33704304319

[chuyeh]

/rerun-failed-ci

[chuyeh]

/rerun-failed-ci

[chuyeh]

/rerun-failed-ci

[chuyeh]

/rerun-failed-ci

[chuyeh]

/rerun-failed-ci

base-b-test-1-gpu-small (2) failed on test_openai_completion_rust.py::test_logprobs_have_zero_kl_against_python_frontend (Rust vs Python token_logprobs exact-equality). Unrelated to this PR's EAGLE draft-extend change; same partition passed on latest scheduled main. Rerunning the failed jobs.

[chuyeh]

Closing as superseded by #35546, which now owns selected-row LM-head pruning for the draft-extend CUDA-graph path.

After merging current main, I reduced this PR to only its remaining ROCm eager-path delta and reran Qwen3.5 with CUDA graphs disabled:

Concurrency	Δ total throughput	Δ median TPOT	Δ median E2E
4	+0.46%	+0.98%	−2.17%
8	−0.11%	−0.29%	+1.17%

Both arms completed all requests with identical input/output token counts. The signs reverse between C4 and C8, so the residual eager-only effect is indistinguishable from run variance. The earlier graph-enabled performance numbers primarily measured functionality that is now upstream through #35546.

Validation after the rescope: 10 focused worker tests passed, 44 inherited selected-row infrastructure tests passed, and pre-commit passed.