[Bugfix][Spec] Fix multi-layer EAGLE DRAFT_EXTEND_V2 attn-TP logprob metadata capture (#28026)

Signed-off-by: Ting Sun <suntcrick@gmail.com>
2026-06-12 19:26:35 -07:00
parent d1a39b0c74
commit e9c3b262e4
1 changed files with 2 additions and 1 deletions
@@ -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