diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index f208c41084e7..302d7a1f372d 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py @@ -206,13 +206,13 @@ def _post_process_after_profile(self, prof_context): # for NPU, profile data will be saved to disk for further analysis. pass - def replay( + def execute( self, forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: if forward_batch.needs_forward_metadata_init(): - self.replay_prepare(forward_batch, pp_proxy_tensors) + self.load_batch(forward_batch, pp_proxy_tensors) else: # In speculative decoding, these two fields are still needed. self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 5bd12c2760a4..5c265f3a9310 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -682,7 +682,7 @@ def _get_skip_cross_attention(self, forward_batch: ForwardBatch) -> bool: return True return bool(forward_batch.encoder_lens.max() == 0) - def can_run(self, forward_batch: ForwardBatch): + def can_run_graph(self, forward_batch: ForwardBatch): is_bs_supported = ( forward_batch.batch_size in self.graphs if self.disable_padding @@ -952,7 +952,7 @@ def prepare_replay( self.model_runner.attn_backend.init_forward_metadata(captured_forward_batch) return captured_forward_batch - def replay( + def execute( self, forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 33119e196ed4..513b66873580 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -799,7 +799,7 @@ def build_prefill_registry( carried from the batch (a read input) rather than written in-graph. Padding policies match the inline copy/zero in - ``PiecewiseCudaGraphRunner.replay_prepare``: ``input_ids`` / ``positions`` + ``PiecewiseCudaGraphRunner.load_batch``: ``input_ids`` / ``positions`` / ``out_cache_loc`` / ``mrope_positions`` / ``input_embeds`` reset their padded tail ``[raw_num_tokens:padded_num_tokens]`` to ``0`` (the padded tokens *are* processed by the graph, so they must be benign), then the head diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index d4bbe7455fef..686f1265a02b 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -516,7 +516,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # Attention planning state. True iff attention metadata for this batch has # already been planned outside ModelRunner.forward (multi-step draft - # pre-plan, plan-stream replay_prepare, hand-built spec batches), so the + # pre-plan, plan-stream load_batch, hand-built spec batches), so the # forward path must not plan again. Only such pre-planners may set this — # ModelRunner / graph runners never mark after their own planning. The # marker is only valid for the planning regime (backend set) it was set @@ -542,7 +542,7 @@ def mark_forward_metadata_ready(self, replan_equivalent: bool = False): Call right next to the out-of-forward planning action (e.g. ``draft_attn_backend.init_forward_metadata(fb)`` or - ``graph_runner.replay_prepare(fb)``). Records the batch shapes so + ``graph_runner.load_batch(fb)``). Records the batch shapes so staleness is detectable; pass ``replan_equivalent=True`` only when a forward-path re-plan is equivalent to the pre-plan (see field docs). diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 126ccfae5257..fe0c94d708ff 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -3416,13 +3416,13 @@ def forward_extend( # Check piecewies cuda graph can_run_graph = ( self.prefill_cuda_graph_runner is not None - and self.prefill_cuda_graph_runner.can_run(forward_batch) + and self.prefill_cuda_graph_runner.can_run_graph(forward_batch) ) if get_cp_strategy() is not None: can_run_graph = False if can_run_graph: # TODO: device_timer.wrap is too broad here — it also includes - # replay_prepare time. Move timing into the prefill cuda graph + # load_batch time. Move timing into the prefill cuda graph # runner to capture only the model.forward part. ctx = ( self.device_timer.wrap(metadata={"category": "extend"}) @@ -3430,7 +3430,7 @@ def forward_extend( else contextlib.nullcontext() ) with ctx: - ret = self.prefill_cuda_graph_runner.replay(forward_batch, **kwargs) + ret = self.prefill_cuda_graph_runner.execute(forward_batch, **kwargs) return (ret, can_run_graph) if not self.server_args.enable_pdmux: @@ -3704,7 +3704,7 @@ def _forward_raw( can_run_graph = bool( mode_check() and self.decode_cuda_graph_runner - and self.decode_cuda_graph_runner.can_run(forward_batch) + and self.decode_cuda_graph_runner.can_run_graph(forward_batch) ) if ( @@ -3717,7 +3717,7 @@ def _forward_raw( # Replay cuda graph if applicable if can_run_graph: - ret = self.decode_cuda_graph_runner.replay( + ret = self.decode_cuda_graph_runner.execute( forward_batch, pp_proxy_tensors=pp_proxy_tensors, ) diff --git a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py index 0cb801333f77..0f3eaea333a3 100644 --- a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py @@ -113,17 +113,17 @@ class BaseCudaGraphRunner(ABC): replay dispatch, and output slicing. Methods: - - can_run(forward_batch) — should forward_batch go through cuda + - can_run_graph(forward_batch) — should forward_batch go through cuda graph replay (vs eager fallback)? - capture_prepare(size, ...) — build the dummy ForwardBatch and - per-capture local state needed by capture_one_shape. - - capture() — outer capture loop; iterates over shapes and calls + per-shape local state needed by capture_one_shape. + - capture() — one-time setup; iterates over shapes and calls capture_one_shape for each. - capture_one_shape(size, ...) — drive one model forward at this shape into the backend's captured artifact. - - replay_prepare(forward_batch, ...) — pad to the nearest captured + - load_batch(forward_batch, ...) — pad to the nearest captured bucket, populate static input buffers, init attention metadata. - - replay(forward_batch, ...) — dispatch one batch through cuda + - execute(forward_batch, ...) — dispatch one batch through cuda graph replay. Notes: @@ -151,7 +151,7 @@ def _pad_to_bucket(raw_size: int, buckets: Sequence[int]) -> int: """Return the smallest buckets[i] >= raw_size. Caller's can_run must reject raw_size > max(buckets) before - reaching replay_prepare; this assertion makes the contract + reaching load_batch; this assertion makes the contract explicit (bisect_left returns len(buckets) when the value exceeds all buckets, which would otherwise IndexError below with no diagnostic). @@ -164,7 +164,7 @@ def _pad_to_bucket(raw_size: int, buckets: Sequence[int]) -> int: return buckets[index] @abstractmethod - def can_run(self, forward_batch: ForwardBatch) -> bool: ... + def can_run_graph(self, forward_batch: ForwardBatch) -> bool: ... @abstractmethod def capture_prepare(self, size: int, *args, **kwargs) -> Any: ... @@ -176,14 +176,14 @@ def capture(self) -> None: ... def capture_one_shape(self, size: int, *args, **kwargs) -> Any: ... @abstractmethod - def replay_prepare( + def load_batch( self, forward_batch: ForwardBatch, **kwargs, ) -> Any: ... @abstractmethod - def replay( + def execute( self, forward_batch: ForwardBatch, **kwargs, diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index fb75707e4b94..ccfe8135dde8 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -510,7 +510,7 @@ def _resolve_lora_variant(self, forward_batch: ForwardBatch): return "lora" return "nolora" - def can_run(self, forward_batch: ForwardBatch): + def can_run_graph(self, forward_batch: ForwardBatch): # Disable for token embedding overrides (dynamic per-request) if forward_batch.replace_embeds is not None: return False @@ -955,7 +955,7 @@ def recapture_if_needed(self, forward_batch: ForwardBatch): self.backend.cleanup() self.capture() - def replay_prepare( + def load_batch( self, forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, @@ -963,7 +963,7 @@ def replay_prepare( self.deepep_adapter.replay() if not forward_batch.needs_forward_metadata_init(): - # Pre-planned (plan-stream replay_prepare already ran). + # Pre-planned (plan-stream load_batch already ran). # In speculative decoding, these two fields are still needed. self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) @@ -1057,7 +1057,7 @@ def replay_prepare( self.bs, stream_idx, variant_label ) - def replay( + def execute( self, forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, @@ -1070,7 +1070,7 @@ def replay( else contextlib.nullcontext() ) with timer_ctx, self.backend.replay_session(): - self.replay_prepare(forward_batch, pp_proxy_tensors) + self.load_batch(forward_batch, pp_proxy_tensors) output = self.backend.replay(self._replay_graph_key, forward_batch) if isinstance(output, LogitsProcessorOutput): diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 1616c69cdaa3..3da471de5197 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -19,7 +19,7 @@ torch.compile's internal cache. Multi-batch supported. - "breakable" — BreakableCudaGraphBackend: segmented capture (no torch.compile). Captures with bs=1; rejects multi-req - prefill in can_run. + prefill in can_run_graph. - "full" — rejected at config validation; not supported for prefill. - "disabled" — handled at the model_runner level — runner not constructed. @@ -420,7 +420,7 @@ def _prepare_forward_metadata_for_replay( static_forward_batch=static_forward_batch, ) - def can_run(self, forward_batch: ForwardBatch) -> bool: + def can_run_graph(self, forward_batch: ForwardBatch) -> bool: if forward_batch.input_embeds is not None: return False if forward_batch.replace_embeds is not None: @@ -451,7 +451,7 @@ def can_run(self, forward_batch: ForwardBatch) -> bool: return False if num_tokens > self.max_num_tokens: return False - # No backend-level shape check here: replay_prepare bucket-pads + # No backend-level shape check here: load_batch bucket-pads # num_tokens up to the nearest captured shape, so eligibility is # bounded by num_tokens <= self.max_num_tokens (already # checked above), not by exact shape membership. @@ -648,7 +648,7 @@ def run_once(): post_warmup_hook=post_warmup_hook, ) - def replay_prepare(self, forward_batch: ForwardBatch, **kwargs) -> ForwardBatch: + def load_batch(self, forward_batch: ForwardBatch, **kwargs) -> ForwardBatch: """Pad, populate static buffers, and build the static_forward_batch the model code reads during replay. """ @@ -782,11 +782,11 @@ def _slot(name): self._static_num_tokens = static_num_tokens return static_forward_batch - def replay( + def execute( self, forward_batch: ForwardBatch, **kwargs ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]: with self.backend.replay_session(): - static_forward_batch = self.replay_prepare(forward_batch, **kwargs) + static_forward_batch = self.load_batch(forward_batch, **kwargs) static_num_tokens = len(static_forward_batch.input_ids) raw_num_tokens = self.raw_num_tokens diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index 4a240f438685..9936606d63e6 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -153,7 +153,9 @@ def prepare_for_draft_extend( # Supply CPU mirror (extend_seq_lens are all num_draft_tokens) so # backend max() reads from list without a per-iter D2H sync. forward_batch.extend_seq_lens_cpu = [num_draft_tokens] * bs - can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch) + can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph( + forward_batch + ) if not batch.forward_mode.is_idle() and not can_cuda_graph: draft_model_runner.attn_backend.init_forward_metadata(forward_batch) # Planned pre-pad; do NOT opt into post-pad re-plan. DSA's indexer @@ -260,7 +262,9 @@ def prepare_for_draft( draft_input.positions = batch.seq_lens.repeat_interleave(topk, dim=0) batch.capture_hidden_mode = capture_mode forward_batch = ForwardBatch.init_new(batch, draft_model_runner) - can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch) + can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph( + forward_batch + ) return forward_batch, can_cuda_graph diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 120b2b8c079c..1f5df7fb3603 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -73,12 +73,12 @@ def prepare_for_verify( can_run_cuda_graph = bool( target_worker.model_runner.decode_cuda_graph_runner - and target_worker.model_runner.decode_cuda_graph_runner.can_run( + and target_worker.model_runner.decode_cuda_graph_runner.can_run_graph( verify_forward_batch ) ) if can_run_cuda_graph: - target_worker.model_runner.decode_cuda_graph_runner.replay_prepare( + target_worker.model_runner.decode_cuda_graph_runner.load_batch( verify_forward_batch ) elif not batch.forward_mode.is_idle(): 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 737523f0c8a4..fbbafdce18c3 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -71,7 +71,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): loop (capture()), bucket-padding helper (_pad_to_bucket), and the backend-driven capture/replay scaffolding. EAGLE-specific bits — buffer dataclass, dummy ForwardBatch construction in - capture_one_shape, replay output unwrap, and can_run — are + capture_one_shape, replay output unwrap, and can_run_graph — are overridden. EAGLE does not call DecodeCudaGraphRunner.__init__ (that init @@ -253,9 +253,9 @@ def _make_graph_key(self, bs, stream_idx=None, variant_label=None): return ShapeKey(size=bs) # ----------------------------------------------------------------- - # can_run + # can_run_graph # ----------------------------------------------------------------- - def can_run(self, forward_batch: ForwardBatch): + def can_run_graph(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = ( max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs @@ -423,7 +423,7 @@ def _postprocess_output_to_raw_bs(self, out, raw_bs): # ----------------------------------------------------------------- # Replay # ----------------------------------------------------------------- - def replay(self, forward_batch: ForwardBatch): + def execute(self, forward_batch: ForwardBatch): assert forward_batch.out_cache_loc is not None self.deepep_adapter.replay() buffers = self.buffers 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 b0c6733c71dc..37e9d0a04e57 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 @@ -71,7 +71,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): Subclasses DecodeCudaGraphRunner to inherit the outer capture loop + backend scaffolding. Overrides capture_one_shape, - replay, can_run for EAGLE-specific draft-extend semantics. + replay, can_run_graph for EAGLE-specific draft-extend semantics. """ def __init__( @@ -255,7 +255,7 @@ def _cache_loc_dtype(self): def _make_graph_key(self, bs, stream_idx=None, variant_label=None): return ShapeKey(size=bs) - def can_run(self, forward_batch: ForwardBatch): + def can_run_graph(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = ( max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs @@ -427,7 +427,7 @@ def run_once(): ), ) - def replay(self, forward_batch: ForwardBatch): + def execute(self, forward_batch: ForwardBatch): assert forward_batch.out_cache_loc is not None self.deepep_adapter.replay() buffers = self.buffers diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 31330bf307be..c5a7c6963c5f 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -338,12 +338,12 @@ def eagle_prepare_for_verify( # Run attention backend plan and cuda graph preparation can_run_cuda_graph = bool( target_worker.model_runner.decode_cuda_graph_runner - and target_worker.model_runner.decode_cuda_graph_runner.can_run( + and target_worker.model_runner.decode_cuda_graph_runner.can_run_graph( verify_forward_batch ) ) if can_run_cuda_graph: - target_worker.model_runner.decode_cuda_graph_runner.replay_prepare( + target_worker.model_runner.decode_cuda_graph_runner.load_batch( verify_forward_batch ) verify_forward_batch.mark_forward_metadata_ready() diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 22ce04998329..232c1c858c56 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -444,7 +444,7 @@ def draft(self, batch: ScheduleBatch): # Run draft if can_cuda_graph: parent_list, top_scores_index, draft_tokens = ( - self.cuda_graph_runner.replay(forward_batch) + self.cuda_graph_runner.execute(forward_batch) ) else: if ( @@ -767,7 +767,7 @@ def _draft_extend_for_decode( # Run draft extend batch in the main compute stream can_cuda_graph = ( self.cuda_graph_runner_for_draft_extend - and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch) + and self.cuda_graph_runner_for_draft_extend.can_run_graph(forward_batch) ) canary_ctx = ( @@ -783,7 +783,7 @@ def _draft_extend_for_decode( ) with canary_ctx: if can_cuda_graph: - draft_logits_output = self.cuda_graph_runner_for_draft_extend.replay( + draft_logits_output = self.cuda_graph_runner_for_draft_extend.execute( forward_batch ) else: @@ -1379,7 +1379,7 @@ def verify(self, batch: ScheduleBatch): ).cpu() # Run target verify batch in the main compute stream (GPU compute). - # Metadata init is skipped iff cuda-graph already ran replay_prepare — + # Metadata init is skipped iff cuda-graph already ran load_batch — # eagle_prepare_for_verify marked the batch in exactly that case; the # non-cuda-graph path stays unmarked and gets forward_extend's init # (post-pad). diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 2ee4a53b019c..785dc6efdf83 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -192,7 +192,7 @@ def _make_graph_key(self, bs, stream_idx=None, variant_label=None): def _replay_graph(self, shape_key, forward_batch): return self.backend.replay(shape_key, forward_batch) - def can_run(self, forward_batch: ForwardBatch): + def can_run_graph(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = max(forward_batch.global_num_tokens_cpu) // ( self.topk * self.topk @@ -336,7 +336,7 @@ def _postprocess_output_to_raw_bs(self, out, raw_bs): parent_list, top_scores_index, draft_tokens = (t[:raw_bs] for t in out) return parent_list, top_scores_index, draft_tokens - def replay(self, forward_batch: ForwardBatch): + def execute(self, forward_batch: ForwardBatch): self.deepep_adapter.replay() buffers = self.buffers diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 6bb8fdbadb05..3be413951ad3 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -434,12 +434,13 @@ def draft(self, batch: ScheduleBatch): self._set_positions(forward_batch) self._expand_for_topk_draft(forward_batch) - can_run_cuda_graph = self.cuda_graph_runner and self.cuda_graph_runner.can_run( - forward_batch + can_run_cuda_graph = ( + self.cuda_graph_runner + and self.cuda_graph_runner.can_run_graph(forward_batch) ) if can_run_cuda_graph: - parent_list, top_scores_index, draft_tokens = self.cuda_graph_runner.replay( - forward_batch + parent_list, top_scores_index, draft_tokens = ( + self.cuda_graph_runner.execute(forward_batch) ) else: forward_batch.can_run_dp_cuda_graph = False diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index f08e9c1a189f..3a1e97d8ce19 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -289,7 +289,7 @@ def _replay_graph(self, shape_key, forward_batch): def _make_graph_key(self, bs, stream_idx=None, variant_label=None): return ShapeKey(size=bs) - def can_run(self, forward_batch: ForwardBatch): + def can_run_graph(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = ( max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs @@ -536,7 +536,7 @@ def init_replay_state( if forward_batch.extend_seq_lens_cpu is not None: self.extend_seq_lens_cpu[:raw_bs] = forward_batch.extend_seq_lens_cpu - def replay(self, forward_batch: ForwardBatch, init_state: bool = True): + def execute(self, forward_batch: ForwardBatch, init_state: bool = True): assert forward_batch.out_cache_loc is not None self.deepep_adapter.replay() buffers = self.buffers @@ -739,5 +739,5 @@ def get_runner(self, step): def get_last_runner(self): return self.runners[-1] if self.runners else None - def can_run(self, forward_batch): - return self.runners[0].can_run(forward_batch) + def can_run_graph(self, forward_batch): + return self.runners[0].can_run_graph(forward_batch) 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 92d99231a6ee..a0c884eb330d 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -541,7 +541,7 @@ def _draft_extend_for_decode( # Run draft extend batch in the main compute stream can_cuda_graph = ( self.cuda_graph_runner_for_draft_extend - and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch) + and self.cuda_graph_runner_for_draft_extend.can_run_graph(forward_batch) ) ret_topk_p_list = [] ret_topk_index_list = [] @@ -574,7 +574,7 @@ def _draft_extend_for_decode( # log_info_on_rank0(logger, f"step: {step}, forward_batch.input_ids: {forward_batch.input_ids}") if can_cuda_graph: draft_logits_output = ( - self.cuda_graph_runner_for_draft_extend.get_runner(step).replay( + self.cuda_graph_runner_for_draft_extend.get_runner(step).execute( forward_batch, init_state=(step == 0) ) ) @@ -840,7 +840,7 @@ def verify( ), ) # NOTE: metadata init is skipped here unconditionally, although - # eagle_prepare_for_verify only plans when cuda-graph replay_prepare ran. + # eagle_prepare_for_verify only plans when cuda-graph load_batch ran. # eagle_worker_v2 re-inits the non-graph path instead (post-pad); this # worker has not adopted that fix, so preserve its behavior verbatim. # On NPU with --disable-cuda-graph, non-graph verify needs metadata init diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index b12b3320306a..e9d14666e0a2 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -675,8 +675,8 @@ def run_eagle_draft_extend_cuda_graph_runner_case( ) adapter.prepare_replay_state(graph_fixture, case, draft_inputs, settings) - testcase.assertTrue(graph_runner.can_run(graph_batch)) - actual = graph_runner.replay(graph_batch) + testcase.assertTrue(graph_runner.can_run_graph(graph_batch)) + actual = graph_runner.execute(graph_batch) adapter.assert_outputs_close(actual, expected, settings) finally: _reset_cuda_graph_test_buffers() diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index af2dca92a0fa..bfbabf1d618d 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -544,8 +544,8 @@ def run_eagle_draft_cuda_graph_runner_case( ) adapter.prepare_replay_state(graph_fixture, case, draft_inputs, settings) - testcase.assertTrue(graph_runner.can_run(graph_batch)) - actual = graph_runner.replay(graph_batch) + testcase.assertTrue(graph_runner.can_run_graph(graph_batch)) + actual = graph_runner.execute(graph_batch) adapter.assert_outputs_close(actual, expected, settings) finally: _reset_cuda_graph_test_buffers() @@ -590,8 +590,8 @@ def run_frozen_kv_mtp_cuda_graph_runner_case( graph_runner = _capture_frozen_kv_mtp_graph_runner(graph_worker) adapter.prepare_replay_state(graph_fixture, case, draft_inputs, settings) - testcase.assertTrue(graph_runner.can_run(graph_batch)) - actual = graph_runner.replay(graph_batch) + testcase.assertTrue(graph_runner.can_run_graph(graph_batch)) + actual = graph_runner.execute(graph_batch) adapter.assert_outputs_close(actual, expected, settings) finally: _reset_cuda_graph_test_buffers()