[Spec] fix EAGLE v2 verify metadata init order on non-cuda-graph path (#26244)

This commit is contained in:
Liangsheng Yin
2026-05-24 18:49:33 -07:00
committed by GitHub
parent 64e2b54a8f
commit 850887dc63
2 changed files with 7 additions and 7 deletions
@@ -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