diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 0c45cba5a7d..29d206b426b 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -12,6 +12,9 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import ( ) from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( + OffloadingConnectorMetadata, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, _ConnectorMetricName, @@ -109,6 +112,40 @@ def test_last_block_offloaded_at_request_finish( ) +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_abort_queued_request_does_not_build_store_job( + request_runner, async_scheduling: bool +): + """Aborting a never-scheduled request must not store unallocated KV.""" + block_size = 4 + runner = request_runner( + block_size=block_size, + num_gpu_blocks=8, + async_scheduling=async_scheduling, + ) + + runner.new_request(token_ids=[0] * (block_size * 4)) + runner.scheduler.schedule() + + runner.new_request(token_ids=[1] * (block_size * 4)) + queued_req_id = str(runner.req_id) + assert any( + request.request_id == queued_req_id for request in runner.scheduler.waiting + ) + + runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED) + req_status = runner.connector_scheduler._req_status[queued_req_id] + assert all(group_state.offload_keys for group_state in req_status.group_states) + assert all(not group_state.block_ids for group_state in req_status.group_states) + + scheduler_output = runner.scheduler.schedule() + + metadata = scheduler_output.kv_connector_metadata + assert isinstance(metadata, OffloadingConnectorMetadata) + assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values()) + assert queued_req_id not in runner.connector_scheduler._req_status + + def test_scheduler_reports_lookup_sync_delay(request_runner): runner = request_runner( block_size=4, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index ce119e5bbe2..2e1875ae73c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -326,9 +326,12 @@ class RequestOffloadState: group_state.block_ids.extend(new_blocks) def storable_chunks( - self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int + self, + group_config: "GroupOffloadConfig", + group_state: RequestGroupState, + num_offloadable_tokens: int, ) -> int: - """Number of leading offloaded chunks eligible for store. + """Number of allocated leading offloaded chunks eligible for store. For eagle/MTP groups the volatile trailing chunk of the offloadable range is excluded while decoding: the draft-layer KV of the last @@ -345,7 +348,10 @@ class RequestOffloadState: is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens if group_config.is_eagle_group and is_decoding: num_chunks = max(0, num_chunks - 1) - return num_chunks + num_allocated_chunks = ( + len(group_state.block_ids) // self.config.blocks_per_chunk + ) + return min(num_chunks, num_allocated_chunks) def advance_stored_idx(self, num_offloadable_tokens: int) -> None: # max(): at the prefill->decode transition of a chunk-aligned prompt, @@ -356,7 +362,7 @@ class RequestOffloadState: ): group_state.next_stored_chunk_idx = max( group_state.next_stored_chunk_idx, - self.storable_chunks(group_config, num_offloadable_tokens), + self.storable_chunks(group_config, group_state, num_offloadable_tokens), ) def update_num_hit_chunks(self, num_cached_tokens: int) -> None: @@ -991,7 +997,7 @@ class OffloadingConnectorScheduler: self.config.kv_group_configs, req_status.group_states ): num_chunks = req_status.storable_chunks( - group_config, num_offloadable_tokens + group_config, group_state, num_offloadable_tokens ) start_chunk_idx = group_state.next_stored_chunk_idx @@ -1068,7 +1074,7 @@ class OffloadingConnectorScheduler: group_config.sliding_window_size_in_chunks is not None ) num_chunks = req_status.storable_chunks( - group_config, num_offloadable_tokens + group_config, group_state, num_offloadable_tokens ) start_chunk_idx = group_state.next_stored_chunk_idx block_ids = group_state.block_ids