A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 20:30 UTC
[BugFix][kv_offload] Fix offloading decodes with async scheduling (#33881)
Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Or Ozeri committed
Mar 11, 2026 at 22:43 UTC
7ee5d5093b369d5c55199bc4613c9afdecabe0b7
3 files changed
+65
-20
tests/v1/kv_connector/unit/test_offloading_connector.py
+50
-15
@@ -148,17 +148,23 @@ class TransferSummary:
148
149
class RequestRunner:
150
def __init__(
151
- self, offloaded_block_size: int, gpu_block_size: int, num_gpu_blocks: int
151
+ self,
152
+ offloaded_block_size: int,
153
+ gpu_block_size: int,
154
+ num_gpu_blocks: int,
155
+ async_scheduling: bool = True,
156
):
157
self.offloaded_block_size: int = offloaded_block_size
158
self.gpu_block_size: int = gpu_block_size
159
self.num_gpu_blocks: int = num_gpu_blocks
160
+ self.async_scheduling: bool = async_scheduling
161
162
self.req_id: int = -1
163
164
vllm_config = create_vllm_config(
165
block_size=gpu_block_size, max_num_batched_tokens=1000
166
)
167
+ vllm_config.scheduler_config.async_scheduling = async_scheduling
168
vllm_config.kv_transfer_config = KVTransferConfig(
169
kv_connector="OffloadingConnector",
170
kv_role="kv_both",
@@ -313,6 +319,8 @@ class RequestRunner:
319
320
tokens_iter = iter(decoded_tokens)
321
token_id = next(tokens_iter, None)
322
+ prev_scheduler_output = None
323
+ prev_model_runner_output = None
324
while True:
325
assert self.scheduler.requests
326
@@ -354,7 +362,16 @@ class RequestRunner:
362
if self.scheduler.running:
363
token_id = next(tokens_iter, None)
364
357
- self.scheduler.update_from_output(scheduler_output, model_runner_output)
365
+ if self.async_scheduling:
366
+ # in async scheduling we update the output of the previous step
367
+ if prev_model_runner_output is not None:
368
+ self.scheduler.update_from_output(
369
+ prev_scheduler_output, prev_model_runner_output
370
+ )
371
+ prev_scheduler_output = scheduler_output
372
+ prev_model_runner_output = model_runner_output
373
+ else:
374
+ self.scheduler.update_from_output(scheduler_output, model_runner_output)
375
376
if (
377
prev_token_id == EOS_TOKEN_ID
@@ -365,6 +382,11 @@ class RequestRunner:
382
continue
383
384
if token_id is None:
385
+ if self.async_scheduling:
386
+ # sample last token
387
+ self.scheduler.update_from_output(
388
+ prev_scheduler_output, prev_model_runner_output
389
+ )
390
break
391
392
self._parse_transfers()
@@ -445,11 +467,14 @@ class RequestRunner:
467
def request_runner():
468
runners = []
469
448
- def runner_factory(offloaded_block_size, gpu_block_size, num_gpu_blocks):
470
+ def runner_factory(
471
+ offloaded_block_size, gpu_block_size, num_gpu_blocks, async_scheduling
472
+ ):
473
runner = RequestRunner(
474
offloaded_block_size=offloaded_block_size,
475
gpu_block_size=gpu_block_size,
476
num_gpu_blocks=num_gpu_blocks,
477
+ async_scheduling=async_scheduling,
478
)
479
runners.append(runner)
480
return runner
@@ -466,7 +491,8 @@ def generate_store_output(block_hashes: Iterable[BlockHash]):
491
)
492
493
469
-def test_offloading_connector(request_runner):
494
+@pytest.mark.parametrize("async_scheduling", [True, False])
495
+def test_offloading_connector(request_runner, async_scheduling: bool):
496
offloaded_block_size = 12
497
gpu_block_size = 4
498
num_gpu_blocks = 100
@@ -476,6 +502,7 @@ def test_offloading_connector(request_runner):
502
offloaded_block_size=offloaded_block_size,
503
gpu_block_size=gpu_block_size,
504
num_gpu_blocks=num_gpu_blocks,
505
+ async_scheduling=async_scheduling,
506
)
507
508
# 3 blocks, store just the middle block (skip first and last)
@@ -498,26 +525,28 @@ def test_offloading_connector(request_runner):
525
runner.run(decoded_tokens=[0])
526
runner.manager.prepare_store.assert_called()
527
501
- # 1 more block, now set block_hashes_to_store = []
528
+ # 1 more block (+ token for async scheduling)
529
+ # now set block_hashes_to_store = []
530
runner.manager.prepare_store.side_effect = (
531
lambda block_hashes: generate_store_output([])
532
)
505
- runner.run(decoded_tokens=[0] * offloaded_block_size)
533
+ runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
534
507
- # 1 more block, now check touch was called with all 6 blocks
535
+ # 1 more block (+ token for kicking off offloading)
536
+ # now check touch was called with all 6 blocks
537
runner.manager.prepare_store.side_effect = (
538
lambda block_hashes: generate_store_output(block_hashes)
539
)
511
- runner.run(decoded_tokens=[0] * offloaded_block_size)
540
+ runner.run(
541
+ decoded_tokens=[0] * (offloaded_block_size + 1),
542
+ expected_stored_gpu_block_indexes=(15, 16, 17),
543
+ )
544
runner.manager.touch.assert_called()
545
block_hashes1 = list(runner.manager.touch.call_args.args[0])
546
assert len(block_hashes1) == 6
547
548
# terminate request
517
- runner.run(
518
- decoded_tokens=[EOS_TOKEN_ID],
519
- expected_stored_gpu_block_indexes=(15, 16, 17),
520
- )
549
+ runner.run(decoded_tokens=[EOS_TOKEN_ID])
550
551
# create a new request differing only on the last token
552
runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
@@ -608,7 +637,8 @@ def test_offloading_connector(request_runner):
637
assert event.medium == "B"
638
639
611
-def test_request_preemption(request_runner):
640
+@pytest.mark.parametrize("async_scheduling", [True, False])
641
+def test_request_preemption(request_runner, async_scheduling: bool):
642
offloaded_block_size = 12
643
gpu_block_size = 4
644
num_gpu_blocks = 100
@@ -617,6 +647,7 @@ def test_request_preemption(request_runner):
647
offloaded_block_size=offloaded_block_size,
648
gpu_block_size=gpu_block_size,
649
num_gpu_blocks=num_gpu_blocks,
650
+ async_scheduling=async_scheduling,
651
)
652
653
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
@@ -674,7 +705,8 @@ def test_request_preemption(request_runner):
705
)
706
707
677
-def test_concurrent_lookups_of_the_same_prefix(request_runner):
708
+@pytest.mark.parametrize("async_scheduling", [True, False])
709
+def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
710
offloaded_block_size = 12
711
gpu_block_size = 4
712
num_gpu_blocks = 100
@@ -683,6 +715,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner):
715
offloaded_block_size=offloaded_block_size,
716
gpu_block_size=gpu_block_size,
717
num_gpu_blocks=num_gpu_blocks,
718
+ async_scheduling=async_scheduling,
719
)
720
721
# store 1 blocks
@@ -732,7 +765,8 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner):
765
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
766
767
735
-def test_abort_loading_requests(request_runner):
768
+@pytest.mark.parametrize("async_scheduling", [True, False])
769
+def test_abort_loading_requests(request_runner, async_scheduling: bool):
770
offloaded_block_size = 12
771
gpu_block_size = 4
772
num_gpu_blocks = 100
@@ -741,6 +775,7 @@ def test_abort_loading_requests(request_runner):
775
offloaded_block_size=offloaded_block_size,
776
gpu_block_size=gpu_block_size,
777
num_gpu_blocks=num_gpu_blocks,
778
+ async_scheduling=async_scheduling,
779
)
780
781
# store 1 blocks
tests/v1/kv_connector/unit/utils.py
+7
-2
@@ -31,6 +31,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( #
31
from vllm.utils.hashing import sha256
32
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
33
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
34
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
35
from vllm.v1.core.sched.scheduler import Scheduler, SchedulerOutput
36
from vllm.v1.kv_cache_interface import (
37
FullAttentionSpec,
@@ -143,7 +144,7 @@ def create_scheduler(
144
vllm_config: VllmConfig,
145
num_blocks: int = 10000,
146
kv_cache_config: KVCacheConfig | None = None,
146
-) -> Scheduler:
147
+) -> Scheduler | AsyncScheduler:
148
"""Initialize Scheduler For Testing."""
149
block_size = vllm_config.cache_config.block_size
150
if kv_cache_config is None:
@@ -163,7 +164,11 @@ def create_scheduler(
164
],
165
)
166
vllm_config.cache_config.num_gpu_blocks = num_blocks
166
- return Scheduler(
167
+
168
+ scheduler_cls = (
169
+ AsyncScheduler if vllm_config.scheduler_config.async_scheduling else Scheduler
170
+ )
171
+ return scheduler_cls(
172
vllm_config=vllm_config,
173
kv_cache_config=kv_cache_config,
174
log_stats=True,
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+8
-3
@@ -416,7 +416,9 @@ class OffloadingConnectorScheduler:
416
417
req = self._requests[req_id]
418
new_tokens = scheduler_output.num_scheduled_tokens[req_id]
419
- total_tokens = req.num_computed_tokens + new_tokens
419
+ expected_tokens = req.num_computed_tokens + new_tokens
420
+ # with async scheduling, some tokens may be missing
421
+ total_tokens = min(expected_tokens, req.num_tokens)
422
num_blocks = total_tokens // self.offloaded_block_size
423
start_block_idx = self._next_stored_block_idx.get(req_id, 0)
424
num_new_blocks = num_blocks - start_block_idx
@@ -424,8 +426,8 @@ class OffloadingConnectorScheduler:
426
if num_new_blocks <= 0:
427
continue
428
427
- # NOTE: In async scheduling, placeholders may temporarily make
428
- # len(req.block_hashes) < num_blocks * self.block_size_factor.
429
+ num_gpu_blocks = num_blocks * self.block_size_factor
430
+ assert len(req.block_hashes) >= num_gpu_blocks
431
432
new_block_hashes = self._get_block_hashes(
433
req, start_idx=start_block_idx, end_idx=num_blocks
@@ -529,6 +531,9 @@ class OffloadingConnectorScheduler:
531
req_id = request.request_id
532
self._requests.pop(req_id, None)
533
self._request_block_ids.pop(req_id, None)
534
+
535
+ # TODO(orozery): possibly kickoff offload for last block
536
+ # which may have been deferred due to async scheduling
537
self._next_stored_block_idx.pop(req_id, None)
538
539
request_being_stored = req_id in self._reqs_being_stored