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

python/sglang/srt/model_executor/cpu_graph_runner.py

@@ -682,7 +682,7 @@ class CPUGraphRunner:
             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 @@ class CPUGraphRunner:
         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 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
         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,13 +3416,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # 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 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 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 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             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 @@ class ModelRunner(ModelRunnerKVCacheMixin):
 
             # 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 @@ class BaseCudaGraphRunner(ABC):
         """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 @@ class BaseCudaGraphRunner(ABC):
         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 @@ class BaseCudaGraphRunner(ABC):
     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 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             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 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             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 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         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 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             self.bs, stream_idx, variant_label
         )
 
-    def replay(
+    def execute(
         self,
         forward_batch: ForwardBatch,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
@@ -1070,7 +1070,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             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 @@ Backend selection comes from cuda_graph_config.prefill:
                       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 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             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 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     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 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             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 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         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