[Bugfix][KV Offloading] Handle queued request aborts without allocated KV blocks (#49146)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Jul 21, 2026 at 11:16 UTC
94ed0bf4e023255d8a0c98da962b6c01ee7dee9e
2 files changed
+49
-6
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+37
@@ -12,6 +12,9 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
12
)
13
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
14
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
15
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
16
+ OffloadingConnectorMetadata,
17
+)
18
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
19
OffloadingConnectorStats,
20
_ConnectorMetricName,
@@ -109,6 +112,40 @@ def test_last_block_offloaded_at_request_finish(
112
)
113
114
115
+@pytest.mark.parametrize("async_scheduling", [True, False])
116
+def test_abort_queued_request_does_not_build_store_job(
117
+ request_runner, async_scheduling: bool
118
+):
119
+ """Aborting a never-scheduled request must not store unallocated KV."""
120
+ block_size = 4
121
+ runner = request_runner(
122
+ block_size=block_size,
123
+ num_gpu_blocks=8,
124
+ async_scheduling=async_scheduling,
125
+ )
126
+
127
+ runner.new_request(token_ids=[0] * (block_size * 4))
128
+ runner.scheduler.schedule()
129
+
130
+ runner.new_request(token_ids=[1] * (block_size * 4))
131
+ queued_req_id = str(runner.req_id)
132
+ assert any(
133
+ request.request_id == queued_req_id for request in runner.scheduler.waiting
134
+ )
135
+
136
+ runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED)
137
+ req_status = runner.connector_scheduler._req_status[queued_req_id]
138
+ assert all(group_state.offload_keys for group_state in req_status.group_states)
139
+ assert all(not group_state.block_ids for group_state in req_status.group_states)
140
+
141
+ scheduler_output = runner.scheduler.schedule()
142
+
143
+ metadata = scheduler_output.kv_connector_metadata
144
+ assert isinstance(metadata, OffloadingConnectorMetadata)
145
+ assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values())
146
+ assert queued_req_id not in runner.connector_scheduler._req_status
147
+
148
+
149
def test_scheduler_reports_lookup_sync_delay(request_runner):
150
runner = request_runner(
151
block_size=4,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+12
-6
@@ -326,9 +326,12 @@ class RequestOffloadState:
326
group_state.block_ids.extend(new_blocks)
327
328
def storable_chunks(
329
- self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
329
+ self,
330
+ group_config: "GroupOffloadConfig",
331
+ group_state: RequestGroupState,
332
+ num_offloadable_tokens: int,
333
) -> int:
331
- """Number of leading offloaded chunks eligible for store.
334
+ """Number of allocated leading offloaded chunks eligible for store.
335
336
For eagle/MTP groups the volatile trailing chunk of the offloadable
337
range is excluded while decoding: the draft-layer KV of the last
@@ -345,7 +348,10 @@ class RequestOffloadState:
348
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
349
if group_config.is_eagle_group and is_decoding:
350
num_chunks = max(0, num_chunks - 1)
348
- return num_chunks
351
+ num_allocated_chunks = (
352
+ len(group_state.block_ids) // self.config.blocks_per_chunk
353
+ )
354
+ return min(num_chunks, num_allocated_chunks)
355
356
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
357
# max(): at the prefill->decode transition of a chunk-aligned prompt,
@@ -356,7 +362,7 @@ class RequestOffloadState:
362
):
363
group_state.next_stored_chunk_idx = max(
364
group_state.next_stored_chunk_idx,
359
- self.storable_chunks(group_config, num_offloadable_tokens),
365
+ self.storable_chunks(group_config, group_state, num_offloadable_tokens),
366
)
367
368
def update_num_hit_chunks(self, num_cached_tokens: int) -> None:
@@ -991,7 +997,7 @@ class OffloadingConnectorScheduler:
997
self.config.kv_group_configs, req_status.group_states
998
):
999
num_chunks = req_status.storable_chunks(
994
- group_config, num_offloadable_tokens
1000
+ group_config, group_state, num_offloadable_tokens
1001
)
1002
1003
start_chunk_idx = group_state.next_stored_chunk_idx
@@ -1068,7 +1074,7 @@ class OffloadingConnectorScheduler:
1074
group_config.sliding_window_size_in_chunks is not None
1075
)
1076
num_chunks = req_status.storable_chunks(
1071
- group_config, num_offloadable_tokens
1077
+ group_config, group_state, num_offloadable_tokens
1078
)
1079
start_chunk_idx = group_state.next_stored_chunk_idx
1080
block_ids = group_state.block_ids