refactor(runner): rename runner replay/load/can_run for the shared surface (#28384)
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+2
-2
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user