[Bug Fix] Resolve EAGLE cuda graph IMA under PD + DP + MTP with GLM-5.1 (#23037)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user