[Bug Fix] Resolve EAGLE cuda graph IMA under PD + DP + MTP with GLM-5.1 (#23037)

This commit is contained in:
Yuxuan Zhang
2026-05-01 13:53:52 -07:00
committed by GitHub
parent b58fa60a1f
commit 79bc2505a5
@@ -26,6 +26,7 @@ from sglang.srt.model_executor.forward_batch_info import (
)
from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.srt.speculative.spec_utils import maybe_detect_nan, maybe_detect_oob
from sglang.srt.utils import (
require_attn_tp_gather,
require_gathered_buffer,
@@ -391,6 +392,10 @@ class EAGLEDraftCudaGraphRunner:
buffers.seq_lens.fill_(self.seq_len_fill_value)
buffers.out_cache_loc.zero_()
buffers.positions.zero_()
buffers.topk_p.zero_()
buffers.topk_index.zero_()
buffers.hidden_states.zero_()
buffers.req_pool_indices.zero_()
num_tokens = bs * self.num_tokens_per_bs
@@ -400,6 +405,17 @@ class EAGLEDraftCudaGraphRunner:
forward_batch.out_cache_loc
)
buffers.positions[:raw_num_token].copy_(forward_batch.positions)
maybe_detect_nan(
forward_batch.spec_info.topk_p,
"EagleDraftCudaGraphRunner.replay: topk_p",
)
maybe_detect_oob(
forward_batch.spec_info.topk_index,
0,
self.model_runner.model_config.vocab_size,
"EagleDraftCudaGraphRunner.replay: topk_index vs vocab_size="
f"{self.model_runner.model_config.vocab_size}",
)
buffers.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p)
buffers.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index)
buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)