From 6fe3e6f17050f387fae3321f17ef1ed5e6b7f7c0 Mon Sep 17 00:00:00 2001 From: Lzy17 Date: Thu, 24 Sep 2026 10:42:14 -0500 Subject: [PATCH] [AMD][DI][CI] Seed the draft proposal distribution on the PD decode path GLM-5.2 EP16 with MTP dies about a minute into serving: File "speculative/eagle_worker_v2.py", line 932, in draft_forward torch.stack(draft_probs_list, dim=1) TypeError: expected Tensor as element 0 in argument 0, but got NoneType The prefill worker sends topk_p and topk_index but no proposal distribution, so build_eagle_disagg_draft_input leaves draft_probs None, and the eager draft loop seeds its list with that field. ROCm auto-enables rejection sampling for EAGLE topk=1, so the seed is required there. The MTP tests that pass all capture CUDA graphs, and that runner allocates a zeroed draft_probs buffer. This recipe sets --disable-cuda-graph and takes the eager path, which had no equivalent. Zeros, matching the graph path: the sampler rejects q == 0 and resamples that position from the target. --- .../srt/speculative/eagle_disaggregation.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/python/sglang/srt/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py index 37335b21cd2d..41fe81e88c82 100644 --- a/python/sglang/srt/speculative/eagle_disaggregation.py +++ b/python/sglang/srt/speculative/eagle_disaggregation.py @@ -92,6 +92,21 @@ def build_eagle_disagg_draft_input( if torch.any(torch.all(dsa_topk_indices < 0, dim=1)).item(): dsa_topk_indices = None + # The prefill worker sends topk_p / topk_index but no proposal + # distribution, and the eager draft loop seeds its list with this field + # (eagle_worker_v2.draft_forward). Zeros stand in for the first draft + # token's q: the sampler rejects q == 0 and resamples that position from + # the target, which is what the graph path's zeroed buffer already does. + draft_probs = ( + torch.zeros( + (topk_index.shape[0], batch.model_config.vocab_size), + device=batch.device, + dtype=torch.float32, + ) + if spec.speculative_use_rejection_sampling + else None + ) + requires_dsa_seed_for_cuda_graph = _requires_dsa_seed_for_cuda_graph( batch.model_config.hf_config, spec.speculative_eagle_topk, @@ -100,6 +115,7 @@ def build_eagle_disagg_draft_input( spec_info = EagleDraftInput( topk_p=topk_p, topk_index=topk_index, + draft_probs=draft_probs, hidden_states=hidden_states, bonus_tokens=last_tokens_tensor, dsa_topk_indices=dsa_topk_indices,