refactor(runner): rename runner replay/load/can_run for the shared surface

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)

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,

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

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).

python/sglang/srt/model_executor/model_runner.py

@@ -3416,21 +3416,21 @@ 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"})
                 if self.device_timer
                 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,
                 )

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,

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,15 +955,15 @@ 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,
     ):
         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):

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
 

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
 
 

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():

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

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

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()