spec: defer verify() idle hidden_size to worker fixup (#25109)

This commit is contained in:
Liangsheng Yin
2026-05-12 22:13:52 -07:00
committed by GitHub
parent 938198e91c
commit adae4042a7
3 changed files with 18 additions and 7 deletions
+9 -4
View File
@@ -252,10 +252,13 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
accepted token logits.
"""
if batch.forward_mode.is_idle():
# hidden_size=None: worker fixup in forward_draft_extend_after_decode
# rebuilds via EagleDraftExtendInput.hidden_size_for(worker)
# (single source incl. EAGLE-3 aux widening).
draft_extend_input = EagleDraftExtendInput.create_idle_input(
device=batch.device,
hidden_size=batch.model_config.spec_hidden_size,
dtype=batch.model_config.dtype,
hidden_size=None,
dtype=None,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
return EagleVerifyOutput.create_idle(
@@ -645,10 +648,12 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
req_pool_indices=batch.req_pool_indices[unfinished_index_device],
)
else:
# hidden_size=None: worker fixup rebuilds via
# EagleDraftExtendInput.hidden_size_for(worker) (single source).
draft_extend_input = EagleDraftExtendInput.create_idle_input(
device=batch.device,
hidden_size=batch.model_config.spec_hidden_size,
dtype=batch.model_config.dtype,
hidden_size=None,
dtype=None,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
@@ -1161,8 +1161,11 @@ class EAGLEWorker(TpModelWorker):
if self.speculative_algorithm.is_standalone()
else CaptureHiddenMode.LAST
)
if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0:
# All reqs finished this verify; swap to an idle ExtendInput.
if draft_extend_input.input_ids.shape[0] == 0:
# Single source for hidden_size via hidden_size_for(self) (incl.
# EAGLE-3 aux widening). Two stub origins from verify(): fully-idle
# batch (DP attn rank w/o reqs) and active batch with all reqs
# finished. prepare_for_idle() is idempotent on already-idle.
batch = batch.copy()
batch.prepare_for_idle()
draft_extend_input = EagleDraftExtendInput.create_idle_input(
@@ -706,7 +706,10 @@ class MultiLayerEagleWorker(TpModelWorker):
if self.speculative_algorithm.is_standalone()
else CaptureHiddenMode.LAST
)
if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0:
if draft_extend_input.input_ids.shape[0] == 0:
# Single source for hidden_size via hidden_size_for(self) (incl.
# EAGLE-3 aux widening). Two stub origins from verify(): fully-idle
# batch and active batch with all reqs finished.
batch = batch.copy()
batch.prepare_for_idle()
draft_extend_input = EagleDraftExtendInput.create_idle_input(