[Auto Sync] Update scheduler.py (20251017) (#11738)

This commit is contained in:
Yineng Zhang
2025-10-17 12:36:07 -07:00
committed by GitHub
parent 8fcc69e7c4
commit b79f75fd53

View File

@@ -2166,6 +2166,12 @@ class Scheduler(
batch.prepare_for_decode() batch.prepare_for_decode()
return batch return batch
# placeholder for override
def update_cache_from_scheduler(
self, schedule_batch: ScheduleBatch, batch_result: GenerationBatchResult
):
pass
def run_batch( def run_batch(
self, batch: ScheduleBatch self, batch: ScheduleBatch
) -> Union[GenerationBatchResult, EmbeddingBatchResult]: ) -> Union[GenerationBatchResult, EmbeddingBatchResult]:
@@ -2242,6 +2248,7 @@ class Scheduler(
batch_or_worker_batch batch_or_worker_batch
) )
future_indices_or_next_token_ids = batch_result.next_token_ids future_indices_or_next_token_ids = batch_result.next_token_ids
self.update_cache_from_scheduler(batch, batch_result)
# NOTE: future_indices_or_next_token_ids is used in ScheduleBatch, # NOTE: future_indices_or_next_token_ids is used in ScheduleBatch,
# which can probably be replaced by future_indices later [TODO(lsyin)]. # which can probably be replaced by future_indices later [TODO(lsyin)].