[Bugfix] Restore KV cache metadata GET method for external event consumer (#56925)
Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
Wei Zhao committed
Sep 17, 2026 at 19:01 UTC
ccde9c61e958d9dbfc21af6c62df6ff15881c811
2 files changed
+101
-1
tests/v1/engine/test_kv_cache_group_metadata.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.v1.core.kv_cache_manager import KVCacheManager
10
+from vllm.v1.engine.core import EngineCore
11
+from vllm.v1.kv_cache_interface import (
12
+ KVCacheConfig,
13
+ KVCacheGroupSpec,
14
+ MambaSpec,
15
+ MLAAttentionSpec,
16
+)
17
+
18
+pytestmark = pytest.mark.cpu_test
19
+
20
+BASE_BLOCK_SIZE = 1536
21
+DCP_WORLD_SIZE = 8
22
+
23
+
24
+def _make_engine_core_with_dcp_manager() -> EngineCore:
25
+ mla_spec = MLAAttentionSpec(
26
+ block_size=BASE_BLOCK_SIZE,
27
+ num_kv_heads=1,
28
+ head_size=576,
29
+ dtype=torch.bfloat16,
30
+ )
31
+ mamba_spec = MambaSpec(
32
+ block_size=BASE_BLOCK_SIZE,
33
+ shapes=((1, 1),),
34
+ dtypes=(torch.float32,),
35
+ mamba_cache_mode="align",
36
+ )
37
+ config = KVCacheConfig(
38
+ num_blocks=32,
39
+ kv_cache_tensors=[],
40
+ kv_cache_groups=[
41
+ KVCacheGroupSpec(["mla"], mla_spec),
42
+ KVCacheGroupSpec(["mamba"], mamba_spec),
43
+ ],
44
+ )
45
+ manager = KVCacheManager(
46
+ config,
47
+ max_model_len=BASE_BLOCK_SIZE * DCP_WORLD_SIZE * 2,
48
+ scheduler_block_size=BASE_BLOCK_SIZE * DCP_WORLD_SIZE,
49
+ hash_block_size=BASE_BLOCK_SIZE,
50
+ dcp_world_size=DCP_WORLD_SIZE,
51
+ )
52
+ scheduler = SimpleNamespace(kv_cache_config=config, kv_cache_manager=manager)
53
+ engine_core = object.__new__(EngineCore)
54
+ engine_core.scheduler = scheduler
55
+ return engine_core
56
+
57
+
58
+def test_kv_cache_group_metadata_uses_effective_dcp_block_sizes() -> None:
59
+ engine_core = _make_engine_core_with_dcp_manager()
60
+
61
+ metadata = engine_core.get_kv_cache_group_metadata()
62
+
63
+ # DCP scales attention blocks but does not scale replicated Mamba state.
64
+ assert [item["block_size"] for item in metadata] == [12288, 1536]
65
+ assert [item["kind"] for item in metadata] == ["mla_attention", "mamba"]
66
+
67
+ managers = engine_core.scheduler.kv_cache_manager.coordinator.single_type_managers
68
+ assert [item["block_size"] for item in metadata] == [
69
+ manager.block_size for manager in managers
70
+ ]
vllm/v1/engine/core.py
+31
-1
@@ -85,7 +85,11 @@ from vllm.v1.fault_tolerance.engine_core_sentinel import (
85
EngineCoreSentinel,
86
fault_tolerant_wrapper,
87
)
88
-from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec
88
+from vllm.v1.kv_cache_interface import (
89
+ KVCacheConfig,
90
+ get_kv_cache_spec_kind,
91
+ is_full_attention_spec,
92
+)
93
from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats
94
from vllm.v1.outputs import ModelRunnerOutput
95
from vllm.v1.request import Request, RequestStatus
@@ -457,6 +461,32 @@ class EngineCore:
461
supported_pooling_tasks,
462
)
463
464
+ def get_kv_cache_group_metadata(self) -> list[dict[str, int | str | None]]:
465
+ """Return serializable KV cache metadata for external event consumers
466
+ (e.g., Dynamo). This function has no call sites within vLLM but needs to
467
+ be kept.
468
+ """
469
+ kv_cache_config = getattr(self.scheduler, "kv_cache_config", None)
470
+ if kv_cache_config is None:
471
+ return []
472
+
473
+ kv_cache_manager = cast(Any, self.scheduler).kv_cache_manager
474
+ managers = kv_cache_manager.coordinator.single_type_managers
475
+ metadata: list[dict[str, int | str | None]] = []
476
+ for group_idx, (group, manager) in enumerate(
477
+ zip(kv_cache_config.kv_cache_groups, managers, strict=True)
478
+ ):
479
+ spec = group.kv_cache_spec
480
+ metadata.append(
481
+ {
482
+ "group_idx": group_idx,
483
+ "kind": get_kv_cache_spec_kind(spec).value,
484
+ "block_size": manager.block_size,
485
+ "sliding_window": getattr(spec, "sliding_window", None),
486
+ }
487
+ )
488
+ return metadata
489
+
490
def add_request(self, request: Request, request_wave: int = 0):
491
"""Add request to the scheduler.
492