Skip to content

Commit 52a2bb0

Browse files
committed
[Bugfix][Mamba] Fix finished_requests_ids of Mamba-like models
Signed-off-by: Sixue(Cecil) Wang <cecilwang@preferred.jp>
1 parent dc1b4a6 commit 52a2bb0

File tree

3 files changed

+12
-0
lines changed

3 files changed

+12
-0
lines changed

vllm/core/scheduler.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -646,6 +646,12 @@ def get_and_reset_finished_requests_ids(self) -> List[str]:
646646
self._finished_requests_ids = list()
647647
return finished_requests_ids
648648

649+
def get_async_stopped_request_ids(self):
650+
if self._async_stopped:
651+
return [seq_group.request_id for seq_group in self._async_stopped]
652+
else:
653+
return []
654+
649655
def _schedule_running(
650656
self,
651657
budget: SchedulingBudget,

vllm/engine/async_llm_engine.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -309,6 +309,9 @@ async def step_async(
309309
finished_requests_ids = self.scheduler[
310310
virtual_engine].get_and_reset_finished_requests_ids()
311311

312+
finished_requests_ids += self.scheduler[
313+
virtual_engine].get_async_stopped_request_ids()
314+
312315
# Maybe switch from async mode to sync mode
313316
if not allow_async_output_proc and len(ctx.output_queue) > 0:
314317
self._process_model_outputs(ctx=ctx)

vllm/engine/llm_engine.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1385,6 +1385,9 @@ def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:
13851385
if finished_request_id in self.seq_id_to_seq_group:
13861386
del self.seq_id_to_seq_group[finished_request_id]
13871387

1388+
finished_requests_ids += self.scheduler[
1389+
virtual_engine].get_async_stopped_request_ids()
1390+
13881391
# Maybe switch from async mode to sync mode
13891392
if not allow_async_output_proc and len(ctx.output_queue) > 0:
13901393
self._process_model_outputs(ctx=ctx)

0 commit comments

Comments
 (0)