From 5a4394a34272b10a1f53bc6c635c8ded1fa01070 Mon Sep 17 00:00:00 2001 From: yinghui <32845984+cicirori@users.noreply.github.com> Date: Sat, 22 Nov 2025 18:55:16 +0100 Subject: [PATCH] align code style eagle draft&draft_extend cuda graph runner (#13533) --- .../eagle_draft_cuda_graph_runner.py | 29 +++++++++---------- .../eagle_draft_extend_cuda_graph_runner.py | 11 +++---- 2 files changed, 19 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 2038ba4d9..cd9d171fe 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -33,10 +33,6 @@ from sglang.srt.utils import ( if TYPE_CHECKING: from sglang.srt.speculative.eagle_worker import EAGLEWorker -import logging - -logger = logging.getLogger(__name__) - class EAGLEDraftCudaGraphRunner: def __init__(self, eagle_worker: EAGLEWorker): @@ -51,29 +47,28 @@ class EAGLEDraftCudaGraphRunner: self.output_buffers = {} self.enable_torch_compile = model_runner.server_args.enable_torch_compile self.disable_padding = model_runner.server_args.disable_cuda_graph_padding - self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args) self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args) self.require_mlp_sync = require_mlp_sync(model_runner.server_args) self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args) - self.dp_size = self.model_runner.dp_size self.tp_size = self.model_runner.tp_size - self.topk = model_runner.server_args.speculative_eagle_topk + self.dp_size = self.model_runner.dp_size self.speculative_num_steps = model_runner.server_args.speculative_num_steps + self.topk = model_runner.server_args.speculative_eagle_topk self.enable_profile_cuda_graph = ( model_runner.server_args.enable_profile_cuda_graph ) self.enable_pdmux = False self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() - server_args = model_runner.server_args # Batch sizes to capture self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner) - self.num_tokens_per_bs = server_args.speculative_eagle_topk # Attention backend + self.num_tokens_per_bs = self.topk self.max_bs = max(self.capture_bs) self.max_num_token = self.max_bs * self.num_tokens_per_bs + self.model_runner.draft_attn_backend.init_cuda_graph_state( self.max_bs, self.max_num_token ) @@ -92,10 +87,6 @@ class EAGLEDraftCudaGraphRunner: with torch.device(model_runner.device): self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64) self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32) - self.seq_lens = torch.full( - (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32 - ) - self.extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32) self.out_cache_loc = torch.zeros( (self.max_num_token * self.speculative_num_steps,), dtype=torch.int64 ) @@ -103,6 +94,10 @@ class EAGLEDraftCudaGraphRunner: self.mrope_positions = torch.zeros( (3, self.max_num_token), dtype=torch.int64 ) + self.seq_lens = torch.full( + (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32 + ) + self.extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32) self.topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32) self.topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64) self.hidden_states = torch.zeros( @@ -194,9 +189,9 @@ class EAGLEDraftCudaGraphRunner: out_cache_loc = self.out_cache_loc[: num_tokens * self.speculative_num_steps] positions = self.positions[:num_tokens] mrope_positions = self.mrope_positions[:, :num_tokens] + hidden_states = self.hidden_states[:num_seqs] topk_p = self.topk_p[:num_seqs] topk_index = self.topk_index[:num_seqs] - hidden_states = self.hidden_states[:num_seqs] if self.require_mlp_tp_gather: self.global_num_tokens_gpu.copy_( @@ -264,6 +259,7 @@ class EAGLEDraftCudaGraphRunner: positions=positions, mrope_positions=mrope_positions, global_num_tokens_gpu=global_num_tokens, + global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob, dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(), global_dp_buffer_len=global_dp_buffer_len, spec_algorithm=self.model_runner.spec_algorithm, @@ -271,7 +267,6 @@ class EAGLEDraftCudaGraphRunner: capture_hidden_mode=( spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL ), - global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob, ) # Attention backend @@ -303,6 +298,7 @@ class EAGLEDraftCudaGraphRunner: self.deepep_adapter.capture(is_extend_in_batch=False) self._capture_init(run_once) + out = self._capture_graph( graph, get_global_graph_memory_pool(), stream, run_once ) @@ -333,6 +329,7 @@ class EAGLEDraftCudaGraphRunner: index = bisect.bisect_left(self.capture_bs, max_batch_size) else: index = bisect.bisect_left(self.capture_bs, raw_bs) + bs = self.capture_bs[index] if bs != raw_bs: self.seq_lens.fill_(self.seq_len_fill_value) @@ -342,7 +339,6 @@ class EAGLEDraftCudaGraphRunner: num_tokens = bs * self.num_tokens_per_bs # Common inputs - self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens) self.out_cache_loc[: raw_num_token * self.speculative_num_steps].copy_( forward_batch.out_cache_loc @@ -351,6 +347,7 @@ class EAGLEDraftCudaGraphRunner: self.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p) self.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index) self.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states) + self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) # TODO(ch-wan): support num_token_non_padded if self.require_gathered_buffer: diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index ca638cec8..0cfe2ef8f 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -57,16 +57,17 @@ class EAGLEDraftExtendCudaGraphRunner: self.require_mlp_sync = require_mlp_sync(model_runner.server_args) self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args) self.tp_size = self.model_runner.tp_size - self.dp_size = model_runner.server_args.dp_size + self.dp_size = self.model_runner.dp_size self.speculative_num_steps = model_runner.server_args.speculative_num_steps self.topk = model_runner.server_args.speculative_eagle_topk self.enable_profile_cuda_graph = ( model_runner.server_args.enable_profile_cuda_graph ) self.enable_pdmux = False + self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() + self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner) self.padded_static_len = -1 - self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() # Attention backend self.num_tokens_per_bs = self.speculative_num_steps + 1 @@ -232,11 +233,11 @@ class EAGLEDraftExtendCudaGraphRunner: seq_lens_cpu = self.seq_lens_cpu[:bs] extend_seq_lens = self.extend_seq_lens[:bs] extend_seq_lens_cpu = self.extend_seq_lens_cpu[:bs] - accept_length = self.accept_length[:bs] out_cache_loc = self.out_cache_loc[:num_tokens] positions = self.positions[:num_tokens] mrope_positions = self.mrope_positions[:, :num_tokens] hidden_states = self.hidden_states[:num_tokens] + accept_length = self.accept_length[:bs] next_token_logits_buffer = self.next_token_logits_buffer[ : bs if self.forward_mode == ForwardMode.DRAFT_EXTEND else num_tokens ] @@ -293,6 +294,8 @@ class EAGLEDraftExtendCudaGraphRunner: seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu, next_token_logits_buffer=next_token_logits_buffer, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, req_to_token_pool=self.model_runner.req_to_token_pool, token_to_kv_pool=self.model_runner.token_to_kv_pool, out_cache_loc=out_cache_loc, @@ -308,8 +311,6 @@ class EAGLEDraftExtendCudaGraphRunner: spec_info=spec_info, capture_hidden_mode=CaptureHiddenMode.LAST, attn_backend=self.eagle_worker.draft_extend_attn_backend, - extend_seq_lens=extend_seq_lens, - extend_seq_lens_cpu=extend_seq_lens_cpu, padded_static_len=self.padded_static_len, )