[EAGLE] Prune draft-extend logits to selected rows

[YAMY1234]

Motivation

During EAGLE draft extension under CUDA graph execution, the model computes full-vocabulary logits for every draft row even though the speculative worker only consumes selected per-request rows. The unnecessary LM-head rows increase graph time and temporary memory usage on the decode critical path.

Modifications

• Select the rows consumed by the speculative verifier before applying the LM head during draft-extend CUDA graph execution.
• Keep the full six-row draft forward unchanged so hidden-state and speculative-state evolution remain identical.
• Preserve the eager path and existing output contract while sizing the graph logits buffer for the selected rows.

Accuracy Tests

Natural-acceptance metric	Baseline	Candidate
Completed requests	112 / 112	112 / 112
Output length	860	860
Mean acceptance length	3.684250	3.652805
Acceptance rate	0.536375	0.530488
CUDA graph enabled	Yes	Yes
Queue / retractions / errors	0 / 0 / 0	0 / 0 / 0

Speed Tests and Profiling

• Setup: Qwen3.5-397B-A17B-NVFP4 on one node with 4 GB300 GPUs, TP4/DP4/EP4 with DP attention, FP8 KV cache, TRTLLM-MHA, FlashInfer CuTeDSL MoE plus FlashInfer A2A for target and draft, NEXTN with 5 speculative steps and 6 draft tokens, online-NVFP4 draft weights, ReplaySSM circular history, Mamba extra_buffer, and CUDA graph tiers through batch 32.
• Workload: one frozen 88,786-token input repeated as a 112-request cohort, 112 warmup requests with OSL 64, 112 measurement requests with exact OSL 860, temperature 0, fixed AL 4.8, and no request replenishment during the measured pure-generation interval.
• Measurement gate: all four ranks must simultaneously report achieved batch 28, CUDA graph enabled, queue depth 0, and no completion/drain row inside the selected interval.

Serving metric	Baseline	Candidate	Delta
Stable rank intervals	17 / 17 / 17 / 17	17 / 17 / 17 / 16	—
Mean context tokens/rank	103,367.5	102,725.7	-0.6209%
Target-step host time	29.688688 ms	28.996720 ms	-0.691968 ms / -2.3307%
Fixed-AL4.8 output tokens/s/GPU	4,526.977	4,635.007	+2.3864%

• To remove the small context difference, an independent exact-context match paired 40 windows at the same mean context of 102,342.4 tokens/rank and measured 4,542.499 to 4,653.987 output tokens/s/GPU (+2.4543%).
• The final zero-resident-context drain row is excluded from both arms.

Nsight metric for draft-extend graph 10658	Baseline	Candidate	Delta
Graph time/GPU-step	1.113471 ms	0.974980 ms	-0.138491 ms / -12.44%
Selected-row projection	0.371574 ms	0.296726 ms	-0.074848 ms
Large output copy	0.098430 ms	0.017771 ms	-0.080659 ms
Graph capture memory	0.15 GB	0.07 GB	-0.08 GB / -53.3%

• The Nsight capture contains 40 actual-B28 steps on every rank with CUDA graph enabled and queue depth 0.
• Draft-decode graph 10637 still performs four full-vocabulary projections per target step by design; this optimization only owns draft-extend graph 10658.
• Output tokens/s/GPU is an engine-side pure-generation metric derived from the fixed-B28 target-step window, not a service-level AIPerf total-throughput metric.
• As a cross-engine diagnostic only, applying the measured within-SGL gain to the fixed-work authority gives 4,608.074 output tokens/s/GPU versus TRT-LLM at 4,791.788, or 96.166% with a 3.834% remaining gap; this does not replace a direct final integration comparison.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #33596531149
Latest PR Test (Extra): ❌ Run #33596530896
Latest PR Test (AMD ROCm 7.2): ❌ Run #33596531105

[YAMY1234]

/tag-and-rerun-ci