Skip to content

[AMD][DI][CI] Seed the draft proposal distribution on the PD decode path - #41145

Open
Lzy17 wants to merge 2 commits into
sgl-project:mainfrom
Lzy17:ci/pd-eagle-draft-probs-seed
Open

Lzy17 wants to merge 2 commits into
sgl-project:mainfrom
Lzy17:ci/pd-eagle-draft-probs-seed

Conversation

@Lzy17

@Lzy17 Lzy17 commented Sep 24, 2026 •

Copy link
Copy Markdown
Collaborator

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

Cause

build_eagle_disagg_draft_input is the only constructor of EagleDraftInput
on the PD decode path. The prefill worker sends topk_p and topk_index but
no proposal distribution, so it leaves draft_probs unset, and the eager draft
loop seeds its list with that field:

# eagle_worker_v2.py:719
if get_spec().speculative_use_rejection_sampling:
    draft_probs_list: List[torch.Tensor] = [spec_info.draft_probs]

The list is stacked at the end of the loop, so a None seed is a guaranteed
TypeError as soon as the first draft step completes.

The graph path never hit this because EAGLEDraftCudaGraphRunner allocates a
zeroed draft_probs buffer under the same condition
(eagle_draft_cuda_graph_runner.py:197). The eager path had no equivalent.
Every MTP test that passes today captures graphs; this recipe sets
--disable-cuda-graph.

Scope: this is not AMD-only

The tag is on the PR because that is where the nightly caught it, but the
defect is in common code and reachable on any backend:

  • speculative_use_rejection_sampling is a documented user-facing flag
    (fields/spec.py:116), default False. Any CUDA user who sets it explicitly
    and runs PD + EAGLE without CUDA graphs hits the identical crash. ROCm only
    differs in that _should_auto_enable_hip_rejection_sampling turns it on by
    default, so AMD reaches the path without opting in — which is why the nightly
    found it first.
  • [PD] Keep EAGLE DP graph and token metadata consistent #32196 (merged) makes the eager path easier to reach for everyone: when
    topk == 1 and a DSA seed is required but dsa_topk_indices is None, the
    batch is now marked cuda_graph_compatible=False and falls back to eager.
    That is exactly the path with the missing seed.

Fix

Allocate zeros under the same condition the graph runner already uses, so the
two paths agree. Zeros are the correct seed rather than a placeholder: the
sampler rejects q == 0 and resamples that position from the target
distribution, which is the behaviour the zeroed graph buffer has always
produced.

Rebased onto #32196; the two changes are independent additions to the same
function.

Verification

MI355X nightly, ROCm 10, GLM-5.2 EP16 + MTP, PD 1P1D with
--disable-cuda-graph: crashes on main, serves to completion with this
patch, GSM8K 0.949.


CI States

Latest PR Test (Base): ✅ Run #36804310116
Latest PR Test (Extra): ❌ Run #36804309955
Latest PR Test (AMD ROCm 10): ❌ Run #36804310105

@Lzy17

Lzy17 commented Sep 24, 2026

Copy link
Copy Markdown
Collaborator Author

Confirmed on MI355X, ROCm 10 (v0.5.20-rocm10-mi35x-20260923), GLM-5.2 FP4
2P1D EP16 + MTP:

before after
GSM8K crashed before scoring 0.949
concurrency sweep 0/7 7/7
draft_probs_list TypeError yes none

One note on why zeros are the right value rather than a workaround: the graph
path's buffer is also permanently zeros here. replay only copies into it when
spec_info.draft_probs is not None, which on the PD path it never is, and
nothing else writes it. So this makes the eager path match the graph path
exactly.

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.
@Lzy17
Lzy17 force-pushed the ci/pd-eagle-draft-probs-seed branch from 317b4fe to 6fe3e6f Compare September 30, 2026 07:16
@HaiShaw HaiShaw added the run-ci CI: run the baseline test suite on this PR label Sep 30, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants