[Spec] fix EAGLE v2 verify metadata init order on non-cuda-graph path (#26244)
This commit is contained in:
@@ -306,11 +306,9 @@ class EagleVerifyInputV2Mixin:
|
||||
)
|
||||
if can_run_cuda_graph:
|
||||
target_worker.model_runner.graph_runner.replay_prepare(verify_forward_batch)
|
||||
else:
|
||||
if not batch.forward_mode.is_idle():
|
||||
target_worker.model_runner.attn_backend.init_forward_metadata(
|
||||
verify_forward_batch
|
||||
)
|
||||
# Non-cuda-graph: defer init to forward_extend, which runs after
|
||||
# `_forward_raw -> prepare_mlp_sync_batch` pads the batch. Initing
|
||||
# here would use pre-pad shapes and trip DSv4 indexer shape match.
|
||||
|
||||
return verify_forward_batch, can_run_cuda_graph
|
||||
|
||||
|
||||
@@ -1039,12 +1039,14 @@ class EAGLEWorkerV2(BaseSpecWorker):
|
||||
verify_input.retrieve_next_token.shape
|
||||
).cpu()
|
||||
|
||||
# Run target verify batch in the main compute stream (GPU compute)
|
||||
# Run target verify batch in the main compute stream (GPU compute).
|
||||
# Only skip metadata init when cuda-graph already ran replay_prepare;
|
||||
# the non-cuda-graph path needs forward_extend's init (post-pad).
|
||||
forward_batch_output = self.target_worker.forward_batch_generation(
|
||||
batch=None,
|
||||
forward_batch=verify_forward_batch,
|
||||
is_verify=True,
|
||||
skip_attn_backend_init=True,
|
||||
skip_attn_backend_init=can_run_cuda_graph,
|
||||
)
|
||||
logits_output = forward_batch_output.logits_output
|
||||
|
||||
|
||||
Reference in New Issue
Block a user