polish code for prefill restrictions (#2992)

This commit is contained in:
Zero Rains
2025-07-23 20:12:01 +08:00
committed by GitHub
parent abd238fc12
commit 5d1788c7b5
2 changed files with 5 additions and 13 deletions

View File

@@ -278,7 +278,7 @@ class PaddleDisWorkerProc():
if self.local_rank % mp_num_per_node == 0:
if self.task_queue.num_tasks() > 0:
# VL only support 1 batch to prefill
if not self.fd_config.model_config.enable_mm or self.worker.prefill_finished():
if not self.fd_config.model_config.enable_mm or not self.worker.prefill_finished():
if self.nnode > 1:
self.task_queue.read_finish_flag.set(1)
else:
@@ -332,8 +332,7 @@ class PaddleDisWorkerProc():
# Execute model to generate token. The generated token will be written to the buffer.
# These generated tokens can be obtained through get_output op.
self.worker.execute_model(req_dicts)
if not self.fd_config.model_config.enable_mm:
self.exist_prefill_task_signal.value[0] = self.worker.prefill_finished()
self.exist_prefill_task_signal.value[0] = self.worker.prefill_finished()
def determine_num_available_blocks(self) -> None:
"""Profiles the peak memory usage of the model to determine how many