From a0728ea5026c782b00f13aa91a67d990b1ae649a Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Thu, 23 Jul 2026 14:52:30 -0700 Subject: [PATCH] [spec decoding] fix inkling multi layer mtp draft extend cuda graph (#32254) --- .../sglang/srt/speculative/multi_layer_eagle_worker_v2.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index d451ac0ee..bc3f2abe2 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -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, )