[BugFix] Fix async scheduling + chunked prefill + preemption (#28787)

Signed-off-by: Nick Hill <nhill@redhat.com>
This commit is contained in:
Nick Hill 2025-11-16 14:46:46 -08:00 committed by GitHub
parent 03ee48111d
commit 80b6080ddc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 8 additions and 9 deletions

View File

@ -65,9 +65,8 @@ def test_without_spec_decoding(
(True, "mp", True, None, False),
(True, "uni", True, None, False),
(False, "mp", True, None, True),
# Async scheduling + preemption + chunked prefill needs to be fixed (WIP)
# (True, "mp", True, None, True),
# (True, "uni", True, None, True),
(True, "mp", True, None, True),
(True, "uni", True, None, True),
]
run_tests(
@ -103,9 +102,8 @@ def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):
(False, "mp", True, spec_config_short, True),
(True, "uni", True, spec_config, False),
(True, "uni", True, spec_config_short, False),
# Async scheduling + preemption + chunked prefill needs to be fixed (WIP)
# (True, "mp", True, spec_config, True),
# (True, "uni", True, spec_config_short, True),
(True, "mp", True, spec_config, True),
(True, "uni", True, spec_config_short, True),
]
run_tests(

View File

@ -778,9 +778,7 @@ class Scheduler(SchedulerInterface):
assert not scheduled_in_prev_step
resumed_req_ids.add(req_id)
if not scheduled_in_prev_step:
all_token_ids[req_id] = req.all_token_ids[
: req.num_computed_tokens + num_tokens
]
all_token_ids[req_id] = req.all_token_ids.copy()
new_block_ids.append(
req_to_new_blocks[req_id].get_block_ids(allow_none=True)
)

View File

@ -97,6 +97,9 @@ class ConstantList(Generic[T], Sequence):
def __repr__(self):
return f"ConstantList({self._x})"
def copy(self) -> list[T]:
return self._x.copy()
class CpuGpuBuffer:
"""Buffer to easily copy tensors between CPU and GPU."""