[Mooncake] Address review nits from #56855 (#57174)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Yifan Qiao <17067717+ivanium@users.noreply.github.com>
Nicolò Lucchesi committed
Sep 23, 2026 at 10:15 UTC
9488318c6a40d1923614fb2d189ca93dbeb25920
5 files changed
+81
-40
tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
+44
-12
@@ -10,6 +10,7 @@ import asyncio
10
from unittest.mock import patch
11
12
import pytest
13
+import torch
14
15
from vllm.config import set_current_vllm_config
16
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
@@ -22,6 +23,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector im
23
SendBlockMeta,
24
TransferRegion,
25
)
26
+from vllm.v1.kv_cache_interface import (
27
+ FullAttentionSpec,
28
+ KVCacheConfig,
29
+ KVCacheGroupSpec,
30
+)
31
32
from .test_mooncake_connector import FakeMooncakeWrapper, patch_worker_dependencies
33
from .utils import create_request, create_vllm_config, make_kv_cache_config
@@ -440,17 +446,19 @@ def test_request_finished_with_hma_groups():
446
# ---------------------------------------------------------------------------
447
# Worker-side load-failure reporting (HMA vs non-HMA)
448
# ---------------------------------------------------------------------------
443
-def _make_kv_consumer_worker(swa_enabled: bool, disable_hma: bool = False):
449
+def _make_kv_consumer_worker(
450
+ swa_enabled: bool = False, kv_cache_config: KVCacheConfig | None = None
451
+):
452
block_size = 16
453
vllm_config = create_vllm_config(
454
kv_connector="MooncakeConnector",
455
kv_role="kv_consumer",
456
block_size=block_size,
457
)
450
- vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma
451
- kv_cache_config = make_kv_cache_config(
452
- block_size=block_size, swa_enabled=swa_enabled
453
- )
458
+ if kv_cache_config is None:
459
+ kv_cache_config = make_kv_cache_config(
460
+ block_size=block_size, swa_enabled=swa_enabled
461
+ )
462
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
463
connector = MooncakeConnector(
464
vllm_config, KVConnectorRole.WORKER, kv_cache_config
@@ -472,19 +480,43 @@ def _make_pull_meta(d_req_id: str, local_block_ids: list[list[int]]) -> PullReqM
480
# requires an accelerator device, so (like the worker tests in
481
# test_mooncake_connector.py) they run on the GPU KV-connectors job only.
482
@pytest.mark.parametrize(
475
- "swa_enabled,disable_hma,expected_is_hma",
483
+ "swa_enabled,expected_is_hma",
484
[
477
- (True, False, True), # SWA group present, HMA enabled
478
- (True, True, False), # SWA group present, but HMA disabled
479
- (False, False, False), # FA only, HMA not needed
485
+ (True, True), # multiple groups (FA + SWA)
486
+ (False, False), # single group, block-level reporting is unambiguous
487
],
488
)
482
-def test_worker_is_hma_required(swa_enabled, disable_hma, expected_is_hma):
483
- """Worker-side _is_hma_required mirrors the scheduler's derivation."""
484
- worker = _make_kv_consumer_worker(swa_enabled, disable_hma)
489
+def test_worker_is_hma_required(swa_enabled, expected_is_hma):
490
+ """Worker-side _is_hma_required mirrors the scheduler's condition."""
491
+ worker = _make_kv_consumer_worker(swa_enabled=swa_enabled)
492
assert worker._is_hma_required is expected_is_hma
493
494
495
+def test_worker_is_hma_required_multiple_full_attention_groups():
496
+ """Two full-attention groups with different block sizes also select
497
+ request-level failure reporting, matching the scheduler."""
498
+ kv_cache_config = KVCacheConfig(
499
+ num_blocks=100,
500
+ kv_cache_tensors=[],
501
+ kv_cache_groups=[
502
+ KVCacheGroupSpec(
503
+ ["layer0"],
504
+ FullAttentionSpec(
505
+ block_size=16, num_kv_heads=4, head_size=16, dtype=torch.float16
506
+ ),
507
+ ),
508
+ KVCacheGroupSpec(
509
+ ["layer1"],
510
+ FullAttentionSpec(
511
+ block_size=32, num_kv_heads=4, head_size=16, dtype=torch.float16
512
+ ),
513
+ ),
514
+ ],
515
+ )
516
+ worker = _make_kv_consumer_worker(kv_cache_config=kv_cache_config)
517
+ assert worker._is_hma_required
518
+
519
+
520
def test_worker_failed_recv_reports_request_level_failure_with_hma():
521
"""With HMA, load failures report the request, not ambiguous block IDs."""
522
worker = _make_kv_consumer_worker(swa_enabled=True)
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+31
-10
@@ -2222,15 +2222,40 @@ def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch):
2222
assert [db.block_size for db in store_worker.token_dbs] == [800, 800]
2223
2224
2225
+def _make_two_full_attention_groups_kv_cache_config():
2226
+ """Two full-attention groups with different block sizes: no hybrid
2227
+ layers, but block-level failure reporting is still unsupported."""
2228
+ return KVCacheConfig(
2229
+ num_blocks=10,
2230
+ kv_cache_tensors=[],
2231
+ kv_cache_groups=[
2232
+ KVCacheGroupSpec(
2233
+ ["full0"],
2234
+ FullAttentionSpec(
2235
+ block_size=800, num_kv_heads=8, head_size=64, dtype=None
2236
+ ),
2237
+ ),
2238
+ KVCacheGroupSpec(
2239
+ ["full1"],
2240
+ FullAttentionSpec(
2241
+ block_size=1600, num_kv_heads=8, head_size=64, dtype=None
2242
+ ),
2243
+ ),
2244
+ ],
2245
+ )
2246
+
2247
+
2248
@pytest.mark.parametrize(
2226
- ("disable_hybrid_kv_cache_manager", "expected_is_hma_required"),
2249
+ ("make_config", "expected_is_hma_required"),
2250
[
2228
- (True, False),
2229
- (False, True),
2251
+ (lambda: _make_qsa_hybrid_kv_cache_config(), True),
2252
+ (lambda: _make_two_full_attention_groups_kv_cache_config(), True),
2253
+ (lambda: _make_kv_cache_config(block_size=800), False),
2254
],
2255
+ ids=["hybrid", "two-full-attention-groups", "single-group"],
2256
)
2257
def test_worker_is_hma_required_from_kv_cache_groups(
2233
- monkeypatch, disable_hybrid_kv_cache_manager, expected_is_hma_required
2258
+ monkeypatch, make_config, expected_is_hma_required
2259
):
2260
store = MagicMock()
2261
store.setup.return_value = 0
@@ -2241,16 +2266,12 @@ def test_worker_is_hma_required_from_kv_cache_groups(
2266
"load_from_config",
2267
staticmethod(lambda: _make_config()),
2268
)
2244
- vllm_config = _make_vllm_config(
2245
- disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager
2246
- )
2269
+ vllm_config = _make_vllm_config()
2270
vllm_config.cache_config.block_size = 800
2271
vllm_config.cache_config.enable_prefix_caching = True
2272
vllm_config.cache_config.prefix_match_unit = None
2273
2251
- store_worker = worker.MooncakeStoreWorker(
2252
- vllm_config, _make_qsa_hybrid_kv_cache_config()
2253
- )
2274
+ store_worker = worker.MooncakeStoreWorker(vllm_config, make_config())
2275
2276
assert store_worker._is_hma_required is expected_is_hma_required
2277
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+3
-9
@@ -1051,15 +1051,9 @@ class MooncakeConnectorWorker:
1051
self.finished_recving_reqs: set[ReqId] = set()
1052
# Written from the receiver loop, drained from the worker thread.
1053
self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()
1054
- self._is_hma_required = (
1055
- not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
1056
- and any(
1057
- not isinstance(g.kv_cache_spec, FullAttentionSpec)
1058
- for g in kv_cache_config.transfer_groups
1059
- )
1060
- )
1061
- # Block IDs are only unique within a group; with HMA the scheduler
1062
- # tracks a single merged group, so failures are reported per request.
1054
+ self._is_hma_required = len(kv_cache_config.kv_cache_groups) > 1
1055
+ # Block IDs are only unique within a group, so with multiple groups
1056
+ # load failures are reported per request instead.
1057
self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
1058
1059
self.xfer_stats = MooncakeKVConnectorStats()
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+1
-7
@@ -1461,13 +1461,7 @@ class MooncakeStoreWorker:
1461
and not self.can_put
1462
)
1463
self.cache_config = vllm_config.cache_config
1464
- self._is_hma_required = (
1465
- not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
1466
- and any(
1467
- not isinstance(g.kv_cache_spec, FullAttentionSpec)
1468
- for g in kv_cache_config.transfer_groups
1469
- )
1470
- )
1464
+ self._is_hma_required = len(kv_cache_config.kv_cache_groups) > 1
1465
self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
1466
kv_cache_config, vllm_config
1467
)
vllm/v1/core/sched/scheduler.py
+2
-2
@@ -3207,8 +3207,8 @@ class Scheduler(SchedulerInterface):
3207
# invalid block IDs cannot be mapped back to requests.
3208
raise RuntimeError(
3209
"A KV connector reported block-level load failures "
3210
- "(invalid_block_ids) on a layout with multiple KV cache "
3211
- "groups, where block IDs are only unique within a group. "
3210
+ "(invalid_block_ids) which is not supported for models "
3211
+ "with multiple KV cache groups. "
3212
"Connectors must report failed requests via "
3213
"KVConnectorTransferResults.failed_recving instead."
3214
)