[spec decoding] fix inkling multi layer mtp draft extend cuda graph (#32254)

2026-07-23 14:52:30 -07:00
parent ed26a111b2
commit a0728ea502
1 changed files with 6 additions and 2 deletions
@@ -676,8 +676,12 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
        # Batch 2: Draft extend
        draft_extend_input = EagleDraftExtendInput(
            hidden_states=batch_result.logits_output.hidden_states,
-            # Actual width: the multi-layer chain fills num_steps + 1 rows/req.
-            num_tokens_per_req=self.speculative_num_steps + 1,
+            # Actual width: the multi-layer chain fills num_steps + 1 rows/req,
+            # plus the boundary-KV front rows when the widened window is active
+            # (must match the capture width in the draft-extend graph runner).
+            num_tokens_per_req=self.speculative_num_steps
+            + 1
+            + self.draft_extend_num_front_tokens,
            num_tokens_for_logprob_per_req=1,
            num_front_tokens=self.draft_extend_num_front_tokens,
        )