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

This commit is contained in:
Cheng Wan
2026-06-19 01:45:04 -07:00
committed by GitHub
parent 31c0a98066
commit 1c6331cbd6
20 changed files with 73 additions and 68 deletions
@@ -206,13 +206,13 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
# 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)
@@ -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,
@@ -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
@@ -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).
@@ -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,
)
@@ -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,
@@ -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):
@@ -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
@@ -153,7 +153,9 @@ class EagleDraftWorkerBase(ABC):
# 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 @@ class EagleDraftWorkerBase(ABC):
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
+2 -2
View File
@@ -73,12 +73,12 @@ class DFlashVerifyInput(SpecInput):
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():
@@ -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 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
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 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
# -----------------------------------------------------------------
# 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
@@ -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 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
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 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
),
)
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
+2 -2
View File
@@ -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()
@@ -444,7 +444,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
# 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 @@ class EagleDraftWorker(EagleDraftWorkerBase):
# 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 @@ class EagleDraftWorker(EagleDraftWorkerBase):
)
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 @@ class EAGLEWorkerV2(BaseSpecWorker):
).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).
@@ -192,7 +192,7 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
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 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
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
@@ -434,12 +434,13 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
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
@@ -289,7 +289,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
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 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
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 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
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)
@@ -541,7 +541,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
# 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 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
# 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 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
),
)
# 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
@@ -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()
@@ -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()