@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8b1d188046

[Bugfix][Mooncake] Report request-level KV load failures under HMA (#56855)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Nicolò Lucchesi committed Sep 16, 2026 at 14:55 UTC 8b1d188046034b42024d22ee52fa6c3948b4acbd
9 files changed +329 -377
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+13 -332
@@ -2,21 +2,12 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 -from typing import Literal
6 -from unittest.mock import Mock, patch
5 +from unittest.mock import Mock
6
7 import pytest
9 -import torch
8
11 -from vllm.v1.core.sched.output import SchedulerOutput
9 from vllm.v1.core.sched.scheduler import Scheduler
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
10 +from vllm.v1.request import Request, RequestStatus
11
12 from .utils import (
13 create_model_runner_output,
@@ -44,18 +35,6 @@ def scheduler():
35 return create_scheduler(vllm_config)
36
37
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 -
38 @pytest.mark.parametrize(
39 "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs",
40 [
@@ -361,317 +340,19 @@ def test_async_progressive_load_failure(
340 assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
341
342
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 - )
343 +@pytest.mark.cpu_test
344 +def test_block_level_load_failure_fails_closed_with_multiple_groups():
345 + """Block-level invalid-block reporting is rejected on multi-group layouts.
346
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.
347 + Block IDs are only unique within a group, so the scheduler cannot map
348 + them back to requests; connectors must report failed_recving instead.
349 """
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(
350 + vllm_config = create_vllm_config()
351 + scheduler = create_scheduler(
352 vllm_config,
523 - num_blocks=num_blocks,
524 - kv_cache_config=_make_deepseek_v4_five_group_config(num_blocks),
525 - hash_block_size=4,
353 + kv_cache_config=make_kv_cache_config(block_size=16, swa_enabled=True),
354 )
355 + assert len(scheduler.kv_cache_config.kv_cache_groups) == 2
356
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
357 + with pytest.raises(RuntimeError, match="failed_recving"):
358 + scheduler._handle_invalid_blocks({1}, {})
tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
+77
@@ -18,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector im
18 MooncakeConnectorMetadata,
19 MooncakeConnectorScheduler,
20 MooncakeXferMetadata,
21 + PullReqMeta,
22 SendBlockMeta,
23 TransferRegion,
24 )
@@ -436,3 +437,79 @@ def test_request_finished_with_hma_groups():
437 assert stored_blocks[0] == fa_blocks
438 # SW: clipped to last 9 blocks (sw_size=128, block_size=16 → 8+1=9)
439 assert stored_blocks[1] == sw_blocks[-9:]
440 +
441 +
442 +# ---------------------------------------------------------------------------
443 +# Worker-side load-failure reporting (HMA vs non-HMA)
444 +# ---------------------------------------------------------------------------
445 +def _make_kv_consumer_worker(swa_enabled: bool, disable_hma: bool = False):
446 + block_size = 16
447 + vllm_config = create_vllm_config(
448 + kv_connector="MooncakeConnector",
449 + kv_role="kv_consumer",
450 + block_size=block_size,
451 + )
452 + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma
453 + kv_cache_config = make_kv_cache_config(
454 + block_size=block_size, swa_enabled=swa_enabled
455 + )
456 + with set_current_vllm_config(vllm_config), patch_worker_dependencies():
457 + connector = MooncakeConnector(
458 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
459 + )
460 + return connector.connector_worker
461 +
462 +
463 +def _make_pull_meta(d_req_id: str, local_block_ids: list[list[int]]) -> PullReqMeta:
464 + return PullReqMeta(
465 + d_req_id=d_req_id,
466 + transfer_id="tx-1",
467 + local_block_ids=local_block_ids,
468 + remote_engine_id="engine-0",
469 + remote_bootstrap_addr="127.0.0.1:1234",
470 + )
471 +
472 +
473 +# The tests below construct a MooncakeConnectorWorker, whose __init__
474 +# requires an accelerator device, so (like the worker tests in
475 +# test_mooncake_connector.py) they run on the GPU KV-connectors job only.
476 +@pytest.mark.parametrize(
477 + "swa_enabled,disable_hma,expected_is_hma",
478 + [
479 + (True, False, True), # SWA group present, HMA enabled
480 + (True, True, False), # SWA group present, but HMA disabled
481 + (False, False, False), # FA only, HMA not needed
482 + ],
483 +)
484 +def test_worker_is_hma_required(swa_enabled, disable_hma, expected_is_hma):
485 + """Worker-side _is_hma_required mirrors the scheduler's derivation."""
486 + worker = _make_kv_consumer_worker(swa_enabled, disable_hma)
487 + assert worker._is_hma_required is expected_is_hma
488 +
489 +
490 +def test_worker_failed_recv_reports_request_level_failure_with_hma():
491 + """With HMA, load failures report the request, not ambiguous block IDs."""
492 + worker = _make_kv_consumer_worker(swa_enabled=True)
493 + assert worker._is_hma_required
494 +
495 + pull_meta = _make_pull_meta("d-req-1", [[1, 2], [3, 4]])
496 + worker._handle_failed_recv({"p-req-1": pull_meta}, ["p-req-1"], "boom")
497 +
498 + assert worker.get_block_ids_with_load_errors() == set()
499 + results = worker.get_transfer_results()
500 + assert results.failed_recving == {"d-req-1"}
501 + assert results.finished_recving == {"d-req-1"}
502 +
503 +
504 +def test_worker_failed_recv_reports_block_ids_without_hma():
505 + """Without HMA, load failures keep reporting block-level errors."""
506 + worker = _make_kv_consumer_worker(swa_enabled=False)
507 + assert not worker._is_hma_required
508 +
509 + pull_meta = _make_pull_meta("d-req-1", [[1, 2], [3, 4]])
510 + worker._handle_failed_recv({"p-req-1": pull_meta}, ["p-req-1"], "boom")
511 +
512 + assert worker.get_block_ids_with_load_errors() == {1, 2, 3, 4}
513 + results = worker.get_transfer_results()
514 + assert results.failed_recving == set()
515 + assert results.finished_recving == {"d-req-1"}
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+33
@@ -12,6 +12,7 @@ from vllm.config import set_current_vllm_config
12 from vllm.distributed.kv_events import BlockStored
13 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
14 KVConnectorRole,
15 + KVConnectorTransferResults,
16 )
17 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
18 connector as mooncake_store_connector,
@@ -245,6 +246,38 @@ def test_worker_methods_delegate_to_store_worker():
246 assert invalid_block_ids == {3, 4}
247
248
249 +def test_get_transfer_results_delegates_to_store_worker():
250 + vllm_config = _make_vllm_config()
251 + kv_cache_config = _make_kv_cache_config()
252 + metadata = MooncakeStoreConnectorMetadata(set(), set())
253 + finished_req_ids = {"req-1"}
254 +
255 + with (
256 + set_current_vllm_config(vllm_config),
257 + patch(
258 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
259 + "connector.MooncakeStoreWorker"
260 + ) as mock_worker_cls,
261 + ):
262 + connector = mooncake_store_connector.MooncakeStoreConnector(
263 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
264 + )
265 +
266 + worker = mock_worker_cls.return_value
267 + transfer_results = KVConnectorTransferResults(
268 + finished_sending={"req-1"},
269 + finished_recving={"req-2"},
270 + failed_recving={"req-3"},
271 + )
272 + worker.get_transfer_results.return_value = transfer_results
273 + connector.bind_connector_metadata(metadata)
274 +
275 + result = connector.get_transfer_results(finished_req_ids)
276 +
277 + worker.get_transfer_results.assert_called_once_with(finished_req_ids, metadata)
278 + assert result is transfer_results
279 +
280 +
281 def test_get_kv_connector_kv_cache_events_returns_none_when_disabled():
282 vllm_config = _make_vllm_config()
283 kv_cache_config = _make_kv_cache_config()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+82
@@ -177,6 +177,7 @@ def _make_store_recving_thread(
177 *,
178 tp_rank: int = 0,
179 disk_offload_buffer_budget_bytes: int | None = None,
180 + is_hma_required: bool = False,
181 ) -> mooncake_store_worker.KVCacheStoreRecvingThread:
182 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
183
@@ -199,6 +200,7 @@ def _make_store_recving_thread(
200 ready_event=threading.Event(),
201 coord=coord,
202 disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,
203 + is_hma_required=is_hma_required,
204 )
205 thread.request_queue.task_done = MagicMock()
206 return thread
@@ -303,6 +305,7 @@ def _make_vllm_config(
305 kv_role: str = "kv_both",
306 pipeline_parallel_size: int = 1,
307 kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC,
308 + disable_hybrid_kv_cache_manager: bool = True,
309 ) -> SimpleNamespace:
310 cache_config = SimpleNamespace(block_size=16, num_gpu_blocks=10)
311 cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout
@@ -319,6 +322,9 @@ def _make_vllm_config(
322 kv_role=kv_role, extra_config=extra_config
323 ),
324 cache_config=cache_config,
325 + scheduler_config=SimpleNamespace(
326 + disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager
327 + ),
328 kv_events_config=SimpleNamespace(enable_kv_cache_events=False),
329 speculative_config=None,
330 )
@@ -2147,6 +2153,48 @@ def test_recv_thread_reports_unsplittable_key_larger_than_budget():
2153 assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}
2154
2155
2156 +def test_recv_thread_partial_load_failure_reports_block_ids_without_hma():
2157 + store = MagicMock()
2158 + # First key fails, second succeeds.
2159 + store.batch_get_into_multi_buffers.return_value = [-1, 0]
2160 + thread = _make_store_recving_thread(store)
2161 +
2162 + thread._handle_request(_make_load_req("req-a", [b"h0", b"h1"], token_len=32))
2163 +
2164 + assert thread.get_and_clear_block_ids_with_load_errors() == {0}
2165 + assert thread.get_and_clear_failed_requests() == set()
2166 + assert thread.get_and_clear_finished_requests() == {"req-a"}
2167 +
2168 +
2169 +def test_recv_thread_partial_load_failure_reports_request_with_hma():
2170 + store = MagicMock()
2171 + store.batch_get_into_multi_buffers.return_value = [-1, 0]
2172 + thread = _make_store_recving_thread(store, is_hma_required=True)
2173 +
2174 + thread._handle_request(_make_load_req("req-a", [b"h0", b"h1"], token_len=32))
2175 +
2176 + assert thread.get_and_clear_block_ids_with_load_errors() == set()
2177 + assert thread.get_and_clear_failed_requests() == {"req-a"}
2178 + assert thread.get_and_clear_finished_requests() == {"req-a"}
2179 +
2180 +
2181 +def test_recv_thread_oversized_key_reports_request_with_hma():
2182 + store = MagicMock()
2183 + thread = _make_store_recving_thread(
2184 + store,
2185 + tp_rank=2,
2186 + disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL,
2187 + is_hma_required=True,
2188 + )
2189 +
2190 + thread._handle_request(_make_load_req("req-a", [b"a0", b"a1", b"a2"], token_len=48))
2191 +
2192 + assert store.batch_get_into_multi_buffers.call_count == 0
2193 + assert thread.get_and_clear_block_ids_with_load_errors() == set()
2194 + assert thread.get_and_clear_failed_requests() == {"req-a"}
2195 + assert thread.get_and_clear_finished_requests() == {"req-a"}
2196 +
2197 +
2198 def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch):
2199 store = MagicMock()
2200 store.setup.return_value = 0
@@ -2172,6 +2220,39 @@ def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch):
2220 assert [db.block_size for db in store_worker.token_dbs] == [800, 800]
2221
2222
2223 +@pytest.mark.parametrize(
2224 + ("disable_hybrid_kv_cache_manager", "expected_is_hma_required"),
2225 + [
2226 + (True, False),
2227 + (False, True),
2228 + ],
2229 +)
2230 +def test_worker_is_hma_required_from_kv_cache_groups(
2231 + monkeypatch, disable_hybrid_kv_cache_manager, expected_is_hma_required
2232 +):
2233 + store = MagicMock()
2234 + store.setup.return_value = 0
2235 + _install_fake_mooncake(monkeypatch, store)
2236 + _patch_worker_runtime(monkeypatch)
2237 + monkeypatch.setattr(
2238 + worker.MooncakeStoreConfig,
2239 + "load_from_config",
2240 + staticmethod(lambda: _make_config()),
2241 + )
2242 + vllm_config = _make_vllm_config(
2243 + disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager
2244 + )
2245 + vllm_config.cache_config.block_size = 800
2246 + vllm_config.cache_config.enable_prefix_caching = True
2247 + vllm_config.cache_config.prefix_match_unit = None
2248 +
2249 + store_worker = worker.MooncakeStoreWorker(
2250 + vllm_config, _make_qsa_hybrid_kv_cache_config()
2251 + )
2252 +
2253 + assert store_worker._is_hma_required is expected_is_hma_required
2254 +
2255 +
2256 def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
2257 store = MagicMock()
2258 store.setup.return_value = 0
@@ -3345,6 +3426,7 @@ def _make_bare_worker(
3426 worker.can_put = kv_role in ("kv_producer", "kv_both") or save_decode_cache
3427 worker._capacity_only = False
3428 worker.block_size = block_size
3429 + worker._is_hma_required = False
3430 worker.tp_rank = 0
3431 worker.enable_kv_events = False
3432 worker.load_async = True
tests/v1/kv_connector/unit/utils.py
-2
@@ -155,7 +155,6 @@ 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,
158 ) -> Scheduler | AsyncScheduler:
159 """Initialize Scheduler For Testing."""
160 block_size = vllm_config.cache_config.block_size
@@ -186,7 +185,6 @@ def create_scheduler(
185 log_stats=True,
186 structured_output_manager=StructuredOutputManager(vllm_config),
187 block_size=block_size,
189 - hash_block_size=hash_block_size,
188 )
189
190
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+45 -1
@@ -30,6 +30,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
30 KVConnectorBase_V1,
31 KVConnectorMetadata,
32 KVConnectorRole,
33 + KVConnectorTransferResults,
34 SupportsHMA,
35 )
36 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
@@ -587,6 +588,13 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
588 assert self.connector_worker is not None
589 return self.connector_worker.get_finished()
590
591 + def get_transfer_results(
592 + self, finished_req_ids: set[str]
593 + ) -> KVConnectorTransferResults:
594 + """Get finished and failed sends/recvs from the worker."""
595 + assert self.connector_worker is not None
596 + return self.connector_worker.get_transfer_results()
597 +
598 def get_block_ids_with_load_errors(self) -> set[int]:
599 """Get block IDs whose remote KV load failed."""
600 assert self.connector_worker is not None
@@ -1046,6 +1054,16 @@ class MooncakeConnectorWorker:
1054 self.finished_recving_reqs: set[ReqId] = set()
1055 # Written from the receiver loop, drained from the worker thread.
1056 self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()
1057 + self._is_hma_required = (
1058 + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
1059 + and any(
1060 + not isinstance(g.kv_cache_spec, FullAttentionSpec)
1061 + for g in kv_cache_config.transfer_groups
1062 + )
1063 + )
1064 + # Block IDs are only unique within a group; with HMA the scheduler
1065 + # tracks a single merged group, so failures are reported per request.
1066 + self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
1067
1068 self.xfer_stats = MooncakeKVConnectorStats()
1069
@@ -1866,6 +1884,29 @@ class MooncakeConnectorWorker:
1884
1885 return finished_sending_reqs or None, finished_recving_reqs or None
1886
1887 + def get_transfer_results(self) -> KVConnectorTransferResults:
1888 + """
1889 + Get transfers that completed on this specific worker, including
1890 + requests whose remote KV load failed.
1891 +
1892 + The scheduler process (via the MultiprocExecutor) will use this output
1893 + to track which workers are done.
1894 + """
1895 + finished_sending_reqs, finished_recving_reqs = self.get_finished()
1896 +
1897 + failed_recving_reqs: set[ReqId] = set()
1898 + while True:
1899 + try:
1900 + failed_recving_reqs.add(self._failed_recv_reqs.get_nowait())
1901 + except queue.Empty:
1902 + break
1903 +
1904 + return KVConnectorTransferResults(
1905 + finished_sending=set(finished_sending_reqs or ()),
1906 + finished_recving=set(finished_recving_reqs or ()),
1907 + failed_recving=failed_recving_reqs,
1908 + )
1909 +
1910 def get_kv_connector_stats(self) -> KVConnectorStats | None:
1911 """Return transfer stats collected since the last call, or None
1912 if nothing has been recorded in this interval."""
@@ -1959,7 +2000,10 @@ class MooncakeConnectorWorker:
2000 # request is waiting on a load, and reporting one here would trip
2001 # the scheduler's `assert req_id in self.requests`.
2002 continue
1962 - self._invalid_block_ids.put(invalid)
2003 + if self._is_hma_required:
2004 + self._failed_recv_reqs.put(pull_meta.d_req_id)
2005 + else:
2006 + self._invalid_block_ids.put(invalid)
2007 self.finished_recving_reqs.add(pull_meta.d_req_id)
2008
2009 if failed:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+9
@@ -27,6 +27,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
27 KVConnectorBase_V1,
28 KVConnectorMetadata,
29 KVConnectorRole,
30 + KVConnectorTransferResults,
31 KVConnectorWorkerMetadata,
32 SupportsHMA,
33 )
@@ -359,6 +360,14 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
360 assert isinstance(metadata, MooncakeStoreConnectorMetadata)
361 return self.connector_worker.get_finished(finished_req_ids, metadata)
362
363 + def get_transfer_results(
364 + self, finished_req_ids: set[str]
365 + ) -> KVConnectorTransferResults:
366 + assert self.connector_worker is not None
367 + metadata = self._get_connector_metadata()
368 + assert isinstance(metadata, MooncakeStoreConnectorMetadata)
369 + return self.connector_worker.get_transfer_results(finished_req_ids, metadata)
370 +
371 def get_block_ids_with_load_errors(self) -> set[int]:
372 assert self.connector_worker is not None
373 return self.connector_worker.get_block_ids_with_load_errors()
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+56 -4
@@ -36,6 +36,9 @@ from vllm.distributed import (
36 get_tensor_model_parallel_world_size,
37 )
38 from vllm.distributed.kv_events import BlockStored
39 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
40 + KVConnectorTransferResults,
41 +)
42 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
43 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501
44 ExternalCachedBlockPool,
@@ -1265,6 +1268,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1268 record_operation: Callable[..., None] | None = None,
1269 request_queue: queue.Queue[Any] | None = None,
1270 group_participates: Sequence[bool] | None = None,
1271 + is_hma_required: bool = False,
1272 ):
1273 super().__init__(
1274 store,
@@ -1284,6 +1288,10 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1288 # _invalid_block_ids can be access by both the Worker and RecvingThread
1289 self._invalid_block_ids_lock = threading.Lock()
1290 self._invalid_block_ids: set[int] = set()
1291 + # With HMA the scheduler tracks a single merged group while block IDs
1292 + # are only unique within a group, so failures are reported per request.
1293 + self._is_hma_required = is_hma_required
1294 + self._failed_requests: set[str] = set()
1295 self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes
1296 self.usable_disk_offload_buffer_budget_bytes = (
1297 None
@@ -1298,6 +1306,22 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1306 with self._invalid_block_ids_lock:
1307 self._invalid_block_ids.update(block_ids)
1308
1309 + def set_failed_request(self, req_id: str):
1310 + with self.done_task_lock:
1311 + self._failed_requests.add(req_id)
1312 +
1313 + def get_and_clear_failed_requests(self) -> set[str]:
1314 + with self.done_task_lock:
1315 + failed = self._failed_requests.copy()
1316 + self._failed_requests.clear()
1317 + return failed
1318 +
1319 + def _report_load_error(self, req_id: str, block_ids: list[int]) -> None:
1320 + if self._is_hma_required:
1321 + self.set_failed_request(req_id)
1322 + else:
1323 + self._add_load_error_block_ids(block_ids)
1324 +
1325 def get_and_clear_block_ids_with_load_errors(self) -> set[int]:
1326 with self._invalid_block_ids_lock:
1327 invalid_block_ids = self._invalid_block_ids.copy()
@@ -1392,7 +1416,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1416 # Mark every block: we skip the whole request, and the
1417 # tp_rank rotation means oversized_key isn't necessarily
1418 # the first block in the request's original order.
1395 - self._add_load_error_block_ids(block_id_list_c)
1419 + self._report_load_error(req_id, block_id_list_c)
1420 oversized_key_bytes = _estimate_disk_offload_staging_bytes(
1421 size_list_c[oversized_key_index]
1422 )
@@ -1465,8 +1489,8 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1489 num_failed_keys=len(failed),
1490 )
1491 if failed:
1468 - self._add_load_error_block_ids(
1469 - [block_id for _, _, block_id in failed]
1492 + self._report_load_error(
1493 + req_id, [block_id for _, _, block_id in failed]
1494 )
1495 logger.warning(
1496 "Failed to get %d Mooncake keys from sub-batch "
@@ -1477,7 +1501,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1501 )
1502 break
1503 except Exception as e:
1480 - self._add_load_error_block_ids(current_batch_block_ids)
1504 + self._report_load_error(req_id, current_batch_block_ids)
1505 self._record_operation(
1506 "load_get",
1507 load_get_start,
@@ -1551,6 +1575,13 @@ class MooncakeStoreWorker:
1575 and not self.can_put
1576 )
1577 self.cache_config = vllm_config.cache_config
1578 + self._is_hma_required = (
1579 + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
1580 + and any(
1581 + not isinstance(g.kv_cache_spec, FullAttentionSpec)
1582 + for g in kv_cache_config.transfer_groups
1583 + )
1584 + )
1585 self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
1586 kv_cache_config, vllm_config
1587 )
@@ -2096,6 +2127,7 @@ class MooncakeStoreWorker:
2127 group.kv_cache_spec.prefix_cacheable
2128 for group in self._kv_cache_groups
2129 ],
2130 + is_hma_required=self._is_hma_required,
2131 )
2132 recv_thread.name = f"KVCacheStoreRecvingThread-{i}"
2133 recv_thread.start()
@@ -2185,6 +2217,26 @@ class MooncakeStoreWorker:
2217 block_ids |= recv_thread.get_and_clear_block_ids_with_load_errors()
2218 return block_ids
2219
2220 + def get_transfer_results(
2221 + self, finished_req_ids: set[str], meta: MooncakeStoreConnectorMetadata
2222 + ) -> KVConnectorTransferResults:
2223 + """Get completed sends/recvs plus requests whose remote KV load failed."""
2224 + done_sending, done_recving = self.get_finished(finished_req_ids, meta)
2225 +
2226 + if self._capacity_only:
2227 + return KVConnectorTransferResults(done_sending, done_recving)
2228 +
2229 + failed_recving: set[str] = set()
2230 + if self.load_async:
2231 + for recv_thread in self.kv_recv_threads:
2232 + failed_recving |= recv_thread.get_and_clear_failed_requests()
2233 +
2234 + return KVConnectorTransferResults(
2235 + finished_sending=done_sending,
2236 + finished_recving=done_recving,
2237 + failed_recving=failed_recving,
2238 + )
2239 +
2240 def _record_kv_connector_operation(
2241 self,
2242 operation: str,
vllm/v1/core/sched/scheduler.py
+14 -38
@@ -3027,11 +3027,6 @@ class Scheduler(SchedulerInterface):
3027 For observability, it also accumulates the total number of tokens that
3028 will need to be recomputed across all affected requests.
3029
3030 - A hybrid model spreads one token span over KV cache groups that disagree
3031 - on block size, so a block index maps to a different token position in
3032 - each group and there is no single longest valid prefix. Those requests
3033 - recompute in full instead.
3034 -
3030 Args:
3031 requests: The set of requests to scan for invalid blocks.
3032 invalid_block_ids: IDs of invalid blocks.
@@ -3055,49 +3050,19 @@ class Scheduler(SchedulerInterface):
3050 # these requests must be rescheduled, but only the first will recompute
3051 # it. This set tracks blocks already marked for recomputation.
3052 marked_invalid_block_ids: set[int] = set()
3058 - # Positions a group holds no KV for -- skipped sliding-window blocks,
3059 - # and every Mamba block but the aligned snapshot -- all point at this
3060 - # one block, shared by every request. It carries no request's data, so
3061 - # it can neither fail for a request nor be evicted on its behalf.
3062 - null_block_id = self.kv_cache_manager.block_pool.null_block.block_id
3053 for request in requests:
3054 is_affected = False
3055 marked_invalid_block = False
3056 req_id = request.request_id
3067 - # One block-id list per KV cache group.
3068 - req_block_ids_per_group = self.kv_cache_manager.get_block_ids(req_id)
3057 + # One block-id list per KV cache group; single-group layouts are
3058 + # enforced by _handle_invalid_blocks.
3059 + (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id)
3060 # We iterate only over blocks that may contain externally computed
3061 # tokens
3062 req_num_computed_tokens = (
3063 request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0)
3064 )
3065
3075 - if len(req_block_ids_per_group) > 1:
3076 - # No single truncation point exists across groups, so the whole
3077 - # request is recomputed. Blocks are still marked, so a request
3078 - # sharing one of them is rescheduled too.
3079 - request_invalid_block_ids = {
3080 - block_id
3081 - for req_block_ids in req_block_ids_per_group
3082 - for block_id in req_block_ids
3083 - if block_id != null_block_id and block_id in invalid_block_ids
3084 - }
3085 - if request_invalid_block_ids:
3086 - marked_invalid_block_ids |= request_invalid_block_ids
3087 - total_affected_tokens += req_num_computed_tokens
3088 - request.num_computed_tokens = 0
3089 - if evict_blocks:
3090 - for req_block_ids in req_block_ids_per_group:
3091 - blocks_to_evict.update(
3092 - block_id
3093 - for block_id in req_block_ids
3094 - if block_id != null_block_id
3095 - )
3096 - affected_req_ids.add(req_id)
3097 - continue
3098 -
3099 - (req_block_ids,) = req_block_ids_per_group
3100 -
3066 req_num_computed_blocks = (
3067 req_num_computed_tokens + self.block_size - 1
3068 ) // self.block_size
@@ -3175,6 +3140,17 @@ class Scheduler(SchedulerInterface):
3140 Returns:
3141 Set of affected request IDs to skip in update_from_output main loop.
3142 """
3143 + if len(self.kv_cache_config.kv_cache_groups) > 1:
3144 + # Block IDs are only unique within a group, so a flat set of
3145 + # invalid block IDs cannot be mapped back to requests.
3146 + raise RuntimeError(
3147 + "A KV connector reported block-level load failures "
3148 + "(invalid_block_ids) on a layout with multiple KV cache "
3149 + "groups, where block IDs are only unique within a group. "
3150 + "Connectors must report failed requests via "
3151 + "KVConnectorTransferResults.failed_recving instead."
3152 + )
3153 +
3154 should_fail = not self.recompute_kv_load_failures
3155
3156 # handle async KV loads (not cached yet, evict_blocks=False)