[PD][Spec] Give a PD decode's first EAGLE draft step a proposal distribution

[salexspb]

PD stack

1. [PD] Wait for prefill completion before Mooncake early KV transfer — #41395.
2. [PD][Spec] Give a PD decode's first EAGLE draft step a proposal distribution — this PR.

Problem

With PD and EAGLE, whenever rejection sampling is enabled, the verifier is handed each draft step's proposal distribution. ROCm enables rejection sampling by default for EAGLE top-k 1 (#37134, _should_auto_enable_hip_rejection_sampling), which is why this crashes there out of the box. EagleWorkerV2.draft_forward starts that list with spec_info.draft_probs. PD prefill sends decode the first draft token, its probability and its hidden state, but not the distribution, so build_eagle_disagg_draft_input leaves draft_probs unset. The PD decode then fails on its first forward:

eagle_worker_v2.py, in draft_forward
    torch.stack(draft_probs_list, dim=1)
TypeError: expected Tensor as element 0 in argument 0, but got NoneType

Under a draft CUDA graph, replay skips the proposal copy, and the verify would read a stale buffer instead. This happens with PD + EAGLE (including native MTP) whenever rejection sampling is enabled. Aggregated serving is unaffected.

Fix

build_eagle_disagg_draft_input sets draft_probs to a one-hot FP32 proposal at each received draft token (pd_first_draft_proposal) whenever rejection sampling is on. Single-layer EAGLE gets one (batch, vocab) row. Multi-layer EAGLE gets the (batch, num_steps, vocab) chain, because its verify reads the received chain directly and otherwise fails with "draft_probs missing".

The verify stays exact. It accepts the token X with probability p(X) and otherwise resamples from the target with X removed, so the committed token is distributed as p, the verifier's target distribution. A greedy request's target is one-hot, so it accepts X exactly when X is the target argmax, with or without the true proposal, and its acceptance is unchanged. Sampled requests accept X with probability p(X) rather than Σ min(p, q), and only at their first decode step. Nothing is added to the PD wire.

#40953 instead transfers the full FP32 distribution through PD metadata, which needs a bootstrap wire change and costs vocab × 4 bytes per request. This PR is the minimal fix that restores correctness; the two could be combined.

Tests

test/registered/unit/speculative/test_pd_eagle_draft_proposal.py (CPU, 5 tests):

• the built draft input carries a one-hot FP32 proposal at each request's drafted token;
• the proposal's width matches the target vocab, as eagle_sample requires;
• there is no proposal without rejection sampling;
• multi-layer EAGLE gets a one-hot FP32 proposal at each received chain step;
• enumerating a single-token chain verify over a drafting q, the committed-token distribution equals p.

On this head: 5 passed. python -m compileall and import sglang are clean.

End-to-end validation (on a HiSparse fork build)

Before this change, a PD 4:4 decode (native MTP, EAGLE 3 steps, top-k 1, rejection sampling on) died at its first request with the torch.stack TypeError above. After it, in every configuration tested (aggregated TP8, PD 4:4 resident decode, PD 4:4 HiSparse decode), draft CUDA-graph capture completed, every service stayed up with no traceback, and 8/8 greedy chat prompts (including two needle prompts) were answered correctly, with mean accept length in the 2.6–2.65 range across configurations, matching the aggregated baseline.

These runs used single-layer EAGLE (native MTP). Multi-layer EAGLE PD is covered by the unit test only.

---

CI States

Latest PR Test (Base): ✅ Run #37887765090
Latest PR Test (Extra): ❌ Run #37887764918
Latest PR Test (AMD ROCm 10): ❌ Run #37887765201

[ShangmingCai]

LGTM

[ShangmingCai]

/tag-and-rerun-ci

[ShangmingCai]

/rerun-failed-ci