[Core] Fix ValueError on KV load failure with a hybrid KV cache (#50388)
Signed-off-by: Kebe <mail@kebe7jun.com>
Kebe committed
Sep 13, 2026 at 22:03 UTC
c351fd3c649569d8a585eb22cfaf2ee728e5fc22
3 files changed
+380
-4
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+340
-2
@@ -2,18 +2,28 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable
5
-from unittest.mock import Mock
5
+from typing import Literal
6
+from unittest.mock import Mock, patch
7
8
import pytest
9
+import torch
10
11
+from vllm.v1.core.sched.output import SchedulerOutput
12
from vllm.v1.core.sched.scheduler import Scheduler
10
-from vllm.v1.request import Request, RequestStatus
13
+from vllm.v1.kv_cache_interface import (
14
+ KVCacheConfig,
15
+ KVCacheGroupSpec,
16
+ MLAAttentionSpec,
17
+ SlidingWindowMLASpec,
18
+)
19
+from vllm.v1.request import FinishReason, Request, RequestStatus
20
21
from .utils import (
22
create_model_runner_output,
23
create_request,
24
create_scheduler,
25
create_vllm_config,
26
+ make_kv_cache_config,
27
)
28
29
@@ -34,6 +44,18 @@ def scheduler():
44
return create_scheduler(vllm_config)
45
46
47
+@pytest.fixture
48
+def hybrid_scheduler():
49
+ """Scheduler for a hybrid model: full attention plus a Mamba group."""
50
+ vllm_config = create_vllm_config(kv_load_failure_policy="recompute")
51
+ kv_cache_config = make_kv_cache_config(
52
+ block_size=vllm_config.cache_config.block_size,
53
+ mamba_enabled=True,
54
+ num_blocks=10000,
55
+ )
56
+ return create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
57
+
58
+
59
@pytest.mark.parametrize(
60
"num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs",
61
[
@@ -337,3 +359,319 @@ def test_async_progressive_load_failure(
359
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
360
assert scheduler.failed_recving_kv_req_ids == {request.request_id}
361
assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
362
+
363
+
364
+def _schedule_hybrid_async_load(
365
+ scheduler: Scheduler, num_prompt_blocks: int, num_external_computed_blocks: int
366
+) -> Request:
367
+ """Put one request into an async external load on a hybrid model."""
368
+ request = create_request(num_tokens=num_prompt_blocks * scheduler.block_size)
369
+ scheduler.add_request(request=request)
370
+
371
+ scheduler.connector = Mock()
372
+ scheduler.connector.get_num_new_matched_tokens.side_effect = (
373
+ _make_get_num_new_matched_tokens(
374
+ {request.request_id: num_external_computed_blocks * scheduler.block_size},
375
+ async_load=True,
376
+ )
377
+ )
378
+ scheduler.connector.take_events.return_value = ()
379
+ return request
380
+
381
+
382
+def _report_invalid_blocks(
383
+ scheduler: Scheduler, scheduler_output, invalid_block_ids: set[int]
384
+) -> None:
385
+ scheduler.update_from_output(
386
+ scheduler_output,
387
+ create_model_runner_output(
388
+ reqs=[],
389
+ finished_recving=set(),
390
+ invalid_block_ids=invalid_block_ids,
391
+ use_eos=True,
392
+ ),
393
+ )
394
+
395
+
396
+@pytest.mark.parametrize("failed_group_idx", [0, 1])
397
+def test_hybrid_load_failure_recomputes_whole_request(
398
+ hybrid_scheduler: Scheduler, failed_group_idx: int
399
+):
400
+ """A hybrid model must recover, not crash.
401
+
402
+ Its groups sit on different block grids, so there is no single truncation
403
+ point; the request restarts from scratch instead.
404
+ """
405
+ num_prompt_blocks, num_external_computed_blocks = 10, 9
406
+ request = _schedule_hybrid_async_load(
407
+ hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks
408
+ )
409
+ scheduler_output = hybrid_scheduler.schedule()
410
+
411
+ block_ids_per_group = hybrid_scheduler.kv_cache_manager.get_block_ids(
412
+ request.request_id
413
+ )
414
+ assert len(block_ids_per_group) == 2, "expected a hybrid (multi-group) layout"
415
+ null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id
416
+ # The Mamba group backs only its aligned snapshot with a real block, so
417
+ # pick a block that actually belongs to this request.
418
+ failed_block = next(
419
+ block_id
420
+ for block_id in block_ids_per_group[failed_group_idx]
421
+ if block_id != null_block_id
422
+ )
423
+
424
+ _report_invalid_blocks(hybrid_scheduler, scheduler_output, {failed_block})
425
+
426
+ assert request.num_computed_tokens == 0
427
+ assert hybrid_scheduler.failed_recving_kv_req_ids == {request.request_id}
428
+ assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
429
+
430
+
431
+def test_hybrid_load_failure_ignores_null_block(hybrid_scheduler: Scheduler):
432
+ """The null block is shared by every request and holds no KV, so reporting
433
+ it must not restart anyone."""
434
+ num_prompt_blocks, num_external_computed_blocks = 10, 9
435
+ request = _schedule_hybrid_async_load(
436
+ hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks
437
+ )
438
+ scheduler_output = hybrid_scheduler.schedule()
439
+
440
+ null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id
441
+ _, mamba_block_ids = hybrid_scheduler.kv_cache_manager.get_block_ids(
442
+ request.request_id
443
+ )
444
+ # The Mamba group parks its unaligned positions on the null block, so a
445
+ # request that never failed would otherwise be restarted by this report.
446
+ assert mamba_block_ids[0] == null_block_id
447
+
448
+ num_computed_tokens = request.num_computed_tokens
449
+ _report_invalid_blocks(hybrid_scheduler, scheduler_output, {null_block_id})
450
+
451
+ assert request.num_computed_tokens == num_computed_tokens
452
+ assert not hybrid_scheduler.failed_recving_kv_req_ids
453
+
454
+
455
+def _make_deepseek_v4_five_group_config(num_blocks: int) -> KVCacheConfig:
456
+ """Build the scheduler-side projection of V4-Flash's packed KV layout."""
457
+
458
+ def mla_spec() -> MLAAttentionSpec:
459
+ return MLAAttentionSpec(
460
+ block_size=256,
461
+ num_kv_heads=1,
462
+ head_size=512,
463
+ dtype=torch.bfloat16,
464
+ tokens_per_state=4,
465
+ model_version="deepseek_v4",
466
+ )
467
+
468
+ def swa_spec(
469
+ *, block_size: int, sliding_window: int, head_size: int
470
+ ) -> SlidingWindowMLASpec:
471
+ return SlidingWindowMLASpec(
472
+ block_size=block_size,
473
+ num_kv_heads=1,
474
+ head_size=head_size,
475
+ dtype=torch.float32,
476
+ sliding_window=sliding_window,
477
+ model_version="deepseek_v4",
478
+ )
479
+
480
+ # DeepSeek-V4 packs these into five physical groups. The two SWA groups
481
+ # have identical scheduling semantics but are split to align layer tuples.
482
+ groups = [
483
+ KVCacheGroupSpec(["full_mla_and_indexer"], mla_spec()),
484
+ KVCacheGroupSpec(
485
+ ["c4_state"],
486
+ swa_spec(block_size=4, sliding_window=8, head_size=2048),
487
+ ),
488
+ KVCacheGroupSpec(
489
+ ["c128_state"],
490
+ swa_spec(block_size=8, sliding_window=128, head_size=2048),
491
+ ),
492
+ KVCacheGroupSpec(
493
+ ["swa.0"],
494
+ swa_spec(block_size=64, sliding_window=512, head_size=512),
495
+ ),
496
+ KVCacheGroupSpec(
497
+ ["swa.1"],
498
+ swa_spec(block_size=64, sliding_window=512, head_size=512),
499
+ ),
500
+ ]
501
+ return KVCacheConfig(
502
+ num_blocks=num_blocks,
503
+ kv_cache_tensors=[],
504
+ kv_cache_groups=groups,
505
+ )
506
+
507
+
508
+def _create_deepseek_v4_scheduler(
509
+ kv_load_failure_policy: Literal["recompute", "fail"] = "recompute",
510
+) -> Scheduler:
511
+ num_blocks = 10000
512
+ vllm_config = create_vllm_config(
513
+ block_size=256,
514
+ max_num_batched_tokens=4096,
515
+ max_model_len=8192,
516
+ # max_model_len only reaches SchedulerConfig; opt-125m still caps the
517
+ # model at 2048, which is below the 9-block requests used here.
518
+ hf_overrides={"max_position_embeddings": 8192},
519
+ kv_load_failure_policy=kv_load_failure_policy,
520
+ )
521
+ return create_scheduler(
522
+ vllm_config,
523
+ num_blocks=num_blocks,
524
+ kv_cache_config=_make_deepseek_v4_five_group_config(num_blocks),
525
+ hash_block_size=4,
526
+ )
527
+
528
+
529
+@pytest.fixture
530
+def deepseek_v4_scheduler() -> Scheduler:
531
+ return _create_deepseek_v4_scheduler()
532
+
533
+
534
+def _schedule_deepseek_v4_async_load(
535
+ scheduler: Scheduler,
536
+) -> tuple[Request, SchedulerOutput]:
537
+ num_external_tokens = 8 * scheduler.block_size
538
+ request = create_request(num_tokens=9 * scheduler.block_size, block_size=4)
539
+ scheduler.add_request(request)
540
+
541
+ scheduler.connector = Mock()
542
+ scheduler.connector.get_num_new_matched_tokens.side_effect = (
543
+ _make_get_num_new_matched_tokens(
544
+ {request.request_id: num_external_tokens}, async_load=True
545
+ )
546
+ )
547
+ scheduler.connector.take_events.return_value = ()
548
+ return request, scheduler.schedule()
549
+
550
+
551
+@pytest.mark.parametrize("failed_group_idx", range(5))
552
+def test_deepseek_v4_load_failure_recomputes_whole_request(
553
+ deepseek_v4_scheduler: Scheduler, failed_group_idx: int
554
+):
555
+ """A failed load in any V4-Flash KV group must not kill EngineCore."""
556
+ request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler)
557
+
558
+ block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids(
559
+ request.request_id
560
+ )
561
+ assert len(block_ids_per_group) == 5
562
+ null_block_id = (
563
+ deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id
564
+ )
565
+ failed_block_id = next(
566
+ block_id
567
+ for block_id in block_ids_per_group[failed_group_idx]
568
+ if block_id != null_block_id
569
+ )
570
+
571
+ deepseek_v4_scheduler.update_from_output(
572
+ scheduler_output,
573
+ create_model_runner_output(
574
+ reqs=[],
575
+ finished_recving=set(),
576
+ invalid_block_ids={failed_block_id},
577
+ use_eos=True,
578
+ ),
579
+ )
580
+
581
+ assert request.num_computed_tokens == 0
582
+ assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
583
+ assert deepseek_v4_scheduler.failed_recving_kv_req_ids == {request.request_id}
584
+
585
+
586
+def test_deepseek_v4_load_failure_ignores_shared_null_block(
587
+ deepseek_v4_scheduler: Scheduler,
588
+):
589
+ """The shared null block contains no request KV and cannot fail a load."""
590
+ request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler)
591
+ null_block_id = (
592
+ deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id
593
+ )
594
+ block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids(
595
+ request.request_id
596
+ )
597
+ assert any(null_block_id in group for group in block_ids_per_group[1:])
598
+ num_computed_tokens = request.num_computed_tokens
599
+
600
+ deepseek_v4_scheduler.update_from_output(
601
+ scheduler_output,
602
+ create_model_runner_output(
603
+ reqs=[],
604
+ finished_recving=set(),
605
+ invalid_block_ids={null_block_id},
606
+ use_eos=True,
607
+ ),
608
+ )
609
+
610
+ assert request.num_computed_tokens == num_computed_tokens
611
+ assert not deepseek_v4_scheduler.failed_recving_kv_req_ids
612
+
613
+
614
+def test_deepseek_v4_sync_load_failure_fails_and_evicts_real_blocks():
615
+ """Fail policy must terminate cleanly and never evict the null block."""
616
+ scheduler = _create_deepseek_v4_scheduler(kv_load_failure_policy="fail")
617
+ num_external_tokens = 8 * scheduler.block_size
618
+ request = create_request(num_tokens=9 * scheduler.block_size, block_size=4)
619
+ scheduler.add_request(request)
620
+
621
+ scheduler.connector = Mock()
622
+ scheduler.connector.get_num_new_matched_tokens.side_effect = (
623
+ _make_get_num_new_matched_tokens(
624
+ {request.request_id: num_external_tokens}, async_load=False
625
+ )
626
+ )
627
+ scheduler.connector.request_finished.return_value = (False, None)
628
+ scheduler.connector.take_events.return_value = ()
629
+ scheduler_output = scheduler.schedule()
630
+
631
+ block_ids_per_group = scheduler.kv_cache_manager.get_block_ids(request.request_id)
632
+ assert len(block_ids_per_group) == 5
633
+ null_block_id = scheduler.kv_cache_manager.block_pool.null_block.block_id
634
+ real_block_ids = {
635
+ block_id
636
+ for group_block_ids in block_ids_per_group
637
+ for block_id in group_block_ids
638
+ if block_id != null_block_id
639
+ }
640
+ failed_block_id = next(
641
+ block_id for block_id in block_ids_per_group[2] if block_id != null_block_id
642
+ )
643
+
644
+ evicted_block_ids: list[set[int]] = []
645
+ original_evict_blocks = scheduler.kv_cache_manager.evict_blocks
646
+
647
+ def evict_blocks_spy(block_ids: set[int]) -> None:
648
+ evicted_block_ids.append(set(block_ids))
649
+ original_evict_blocks(block_ids)
650
+
651
+ with (
652
+ patch.object(
653
+ scheduler.kv_cache_manager,
654
+ "evict_blocks",
655
+ side_effect=evict_blocks_spy,
656
+ ),
657
+ patch.object(scheduler, "_connector_finished", return_value=(False, None)),
658
+ ):
659
+ outputs = scheduler.update_from_output(
660
+ scheduler_output,
661
+ create_model_runner_output(
662
+ [request],
663
+ invalid_block_ids={failed_block_id},
664
+ use_eos=True,
665
+ ),
666
+ )
667
+
668
+ assert request.status == RequestStatus.FINISHED_ERROR
669
+ assert request.get_finished_reason() == FinishReason.ERROR
670
+ assert request.request_id not in scheduler.requests
671
+ assert evicted_block_ids == [real_block_ids]
672
+ assert null_block_id not in evicted_block_ids[0]
673
+
674
+ engine_outputs = next(iter(outputs.values()))
675
+ assert len(engine_outputs.outputs) == 1
676
+ assert engine_outputs.outputs[0].request_id == request.request_id
677
+ assert engine_outputs.outputs[0].finish_reason == FinishReason.ERROR
tests/v1/kv_connector/unit/utils.py
+2
@@ -155,6 +155,7 @@ def create_scheduler(
155
vllm_config: VllmConfig,
156
num_blocks: int = 10000,
157
kv_cache_config: KVCacheConfig | None = None,
158
+ hash_block_size: int | None = None,
159
) -> Scheduler | AsyncScheduler:
160
"""Initialize Scheduler For Testing."""
161
block_size = vllm_config.cache_config.block_size
@@ -185,6 +186,7 @@ def create_scheduler(
186
log_stats=True,
187
structured_output_manager=StructuredOutputManager(vllm_config),
188
block_size=block_size,
189
+ hash_block_size=hash_block_size,
190
)
191
192
vllm/v1/core/sched/scheduler.py
+38
-2
@@ -3058,6 +3058,11 @@ class Scheduler(SchedulerInterface):
3058
For observability, it also accumulates the total number of tokens that
3059
will need to be recomputed across all affected requests.
3060
3061
+ A hybrid model spreads one token span over KV cache groups that disagree
3062
+ on block size, so a block index maps to a different token position in
3063
+ each group and there is no single longest valid prefix. Those requests
3064
+ recompute in full instead.
3065
+
3066
Args:
3067
requests: The set of requests to scan for invalid blocks.
3068
invalid_block_ids: IDs of invalid blocks.
@@ -3081,18 +3086,49 @@ class Scheduler(SchedulerInterface):
3086
# these requests must be rescheduled, but only the first will recompute
3087
# it. This set tracks blocks already marked for recomputation.
3088
marked_invalid_block_ids: set[int] = set()
3089
+ # Positions a group holds no KV for -- skipped sliding-window blocks,
3090
+ # and every Mamba block but the aligned snapshot -- all point at this
3091
+ # one block, shared by every request. It carries no request's data, so
3092
+ # it can neither fail for a request nor be evicted on its behalf.
3093
+ null_block_id = self.kv_cache_manager.block_pool.null_block.block_id
3094
for request in requests:
3095
is_affected = False
3096
marked_invalid_block = False
3097
req_id = request.request_id
3088
- # TODO (davidb): add support for hybrid memory allocator
3089
- (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id)
3098
+ # One block-id list per KV cache group.
3099
+ req_block_ids_per_group = self.kv_cache_manager.get_block_ids(req_id)
3100
# We iterate only over blocks that may contain externally computed
3101
# tokens
3102
req_num_computed_tokens = (
3103
request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0)
3104
)
3105
3106
+ if len(req_block_ids_per_group) > 1:
3107
+ # No single truncation point exists across groups, so the whole
3108
+ # request is recomputed. Blocks are still marked, so a request
3109
+ # sharing one of them is rescheduled too.
3110
+ request_invalid_block_ids = {
3111
+ block_id
3112
+ for req_block_ids in req_block_ids_per_group
3113
+ for block_id in req_block_ids
3114
+ if block_id != null_block_id and block_id in invalid_block_ids
3115
+ }
3116
+ if request_invalid_block_ids:
3117
+ marked_invalid_block_ids |= request_invalid_block_ids
3118
+ total_affected_tokens += req_num_computed_tokens
3119
+ request.num_computed_tokens = 0
3120
+ if evict_blocks:
3121
+ for req_block_ids in req_block_ids_per_group:
3122
+ blocks_to_evict.update(
3123
+ block_id
3124
+ for block_id in req_block_ids
3125
+ if block_id != null_block_id
3126
+ )
3127
+ affected_req_ids.add(req_id)
3128
+ continue
3129
+
3130
+ (req_block_ids,) = req_block_ids_per_group
3131
+
3132
req_num_computed_blocks = (
3133
req_num_computed_tokens + self.block_size - 1
3134
) // self.block_size