[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:
|
if can_run_cuda_graph:
|
||||||
target_worker.model_runner.graph_runner.replay_prepare(verify_forward_batch)
|
target_worker.model_runner.graph_runner.replay_prepare(verify_forward_batch)
|
||||||
else:
|
# Non-cuda-graph: defer init to forward_extend, which runs after
|
||||||
if not batch.forward_mode.is_idle():
|
# `_forward_raw -> prepare_mlp_sync_batch` pads the batch. Initing
|
||||||
target_worker.model_runner.attn_backend.init_forward_metadata(
|
# here would use pre-pad shapes and trip DSv4 indexer shape match.
|
||||||
verify_forward_batch
|
|
||||||
)
|
|
||||||
|
|
||||||
return verify_forward_batch, can_run_cuda_graph
|
return verify_forward_batch, can_run_cuda_graph
|
||||||
|
|
||||||
|
|||||||
@@ -1039,12 +1039,14 @@ class EAGLEWorkerV2(BaseSpecWorker):
|
|||||||
verify_input.retrieve_next_token.shape
|
verify_input.retrieve_next_token.shape
|
||||||
).cpu()
|
).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(
|
forward_batch_output = self.target_worker.forward_batch_generation(
|
||||||
batch=None,
|
batch=None,
|
||||||
forward_batch=verify_forward_batch,
|
forward_batch=verify_forward_batch,
|
||||||
is_verify=True,
|
is_verify=True,
|
||||||
skip_attn_backend_init=True,
|
skip_attn_backend_init=can_run_cuda_graph,
|
||||||
)
|
)
|
||||||
logits_output = forward_batch_output.logits_output
|
logits_output = forward_batch_output.logits_output
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user