From e9c3b262e4eada54bb4972c1f3dc9bad3e80d1b3 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sat, 13 Jun 2026 10:26:35 +0800 Subject: [PATCH] [Bugfix][Spec] Fix multi-layer EAGLE DRAFT_EXTEND_V2 attn-TP logprob metadata capture (#28026) Signed-off-by: Ting Sun --- .../multi_layer_eagle_draft_extend_cuda_graph_runner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 7a5d97fcb..d41a3357e 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -336,7 +336,8 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): global_num_tokens_for_logprob_cpu = [num_tokens] * self.dp_size elif self.require_attn_tp_gather: global_num_tokens_cpu = [num_tokens] - global_num_tokens_for_logprob_cpu = [bs] + # DRAFT_EXTEND_V2 produces logits for all tokens, not bs (see mlp branch above) + global_num_tokens_for_logprob_cpu = [num_tokens] else: global_num_tokens_cpu = None