From 3e367f9bcdf903ed580d328f9071acb3cb9ef2d7 Mon Sep 17 00:00:00 2001 From: Tarushii Goel Date: Mon, 20 Apr 2026 16:29:16 -0700 Subject: [PATCH] [sgl] fix incorrect behavior in cuda graph draft extend (#22832) --- .../multi_layer_eagle_worker_v2.py | 26 ++++++++++++------- 1 file changed, 17 insertions(+), 9 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 1b216d46f..8940aae97 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -538,20 +538,28 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): # Update req_to_hidden_states_pool for KV Cache reversion if ( - self.cuda_graph_runner_for_draft_extend is not None - and forward_batch.extend_seq_lens is not None + forward_batch.extend_seq_lens is not None + and self.cuda_graph_runner_for_draft_extend is not None ): - last_cuda_graph_runner = ( - self.cuda_graph_runner_for_draft_extend.get_last_runner() - ) + if can_cuda_graph: + last_runner = self.cuda_graph_runner_for_draft_extend.get_last_runner() + hidden_states = last_runner.buffers.hidden_states + req_pool_indices = last_runner.buffers.req_pool_indices + extend_seq_lens = last_runner.buffers.extend_seq_lens + extend_start_loc = last_runner.buffers.extend_start_loc + else: + hidden_states = draft_logits_output.logits_output.hidden_states + req_pool_indices = forward_batch.req_pool_indices + extend_seq_lens = forward_batch.extend_seq_lens + extend_start_loc = forward_batch.extend_start_loc assign_hidden_states_pool_triton( - last_cuda_graph_runner.buffers.hidden_states, - last_cuda_graph_runner.buffers.req_pool_indices, + hidden_states, + req_pool_indices, self.req_to_hidden_states_pool, self.speculative_num_steps - 1, forward_batch.batch_size, - last_cuda_graph_runner.buffers.extend_seq_lens, - last_cuda_graph_runner.buffers.extend_start_loc, + extend_seq_lens, + extend_start_loc, ) # Reorganize the spec info for the next batch