spec: defer verify() idle hidden_size to worker fixup (#25109)
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user