[DecodeBenchConnector] Fix HMA cache-group mapping (#54647)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: OpenAI Codex <noreply@openai.com>
Jeff (Junze) Ma committed
Sep 1, 2026 at 01:52 UTC
30dd1a7954ff91ae0359806ec56b22b6be0c529f
2 files changed
+128
-34
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+101
-4
@@ -7,6 +7,8 @@ Tests the functionality of the DecodeBenchConnector which fills KV cache
7
with dummy values for decode performance benchmarking.
8
"""
9
10
+from unittest.mock import MagicMock
11
+
12
import pytest
13
import torch
14
@@ -22,6 +24,11 @@ from vllm.forward_context import ForwardContext
24
from vllm.utils.hashing import sha256
25
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
26
from vllm.v1.core.sched.scheduler import Scheduler
27
+from vllm.v1.kv_cache_interface import (
28
+ FullAttentionSpec,
29
+ KVCacheConfig,
30
+ KVCacheGroupSpec,
31
+)
32
from vllm.v1.request import Request
33
34
from .utils import (
@@ -65,11 +72,10 @@ class DecodeBenchTestRunner:
72
# Using simplified shape for testing
73
num_heads = 4
74
head_dim = 64
75
+ layer_names = self.scheduler.kv_cache_config.kv_cache_groups[0].layer_names
76
self.kv_caches = {
69
- f"layer_{i}": torch.zeros(
70
- num_gpu_blocks, 2, num_heads, block_size, head_dim
71
- )
72
- for i in range(2) # 2 layers for testing
77
+ layer_name: torch.zeros(num_gpu_blocks, 2, num_heads, block_size, head_dim)
78
+ for layer_name in layer_names
79
}
80
81
# Register KV caches with worker connector
@@ -179,6 +185,97 @@ def test_decode_bench_connector_basic():
185
assert torch.allclose(block_data, torch.tensor(0.015))
186
187
188
+def test_decode_bench_connector_fills_each_hma_group():
189
+ """Each cache group is filled using its own block IDs."""
190
+ block_size = 16
191
+ num_gpu_blocks = 8
192
+ vllm_config = create_vllm_config(
193
+ block_size=block_size,
194
+ max_num_batched_tokens=1000,
195
+ kv_connector="DecodeBenchConnector",
196
+ )
197
+ kv_cache_spec = FullAttentionSpec(
198
+ block_size=block_size,
199
+ num_kv_heads=1,
200
+ head_size=1,
201
+ dtype=torch.float32,
202
+ )
203
+ kv_cache_config = KVCacheConfig(
204
+ num_blocks=num_gpu_blocks,
205
+ kv_cache_tensors=[],
206
+ kv_cache_groups=[
207
+ KVCacheGroupSpec(["group_0_layer"], kv_cache_spec),
208
+ KVCacheGroupSpec(["group_1_layer"], kv_cache_spec),
209
+ ],
210
+ )
211
+ connector = DecodeBenchConnector(
212
+ vllm_config,
213
+ KVConnectorRole.WORKER,
214
+ kv_cache_config,
215
+ )
216
+ group_0_cache = torch.zeros(num_gpu_blocks, 2)
217
+ group_1_cache = torch.zeros(num_gpu_blocks, 2)
218
+ connector.register_kv_caches(
219
+ {
220
+ "group_0_layer": group_0_cache,
221
+ "group_1_layer": group_1_cache,
222
+ }
223
+ )
224
+ connector.bind_connector_metadata(
225
+ DecodeBenchConnectorMetadata(reqs_to_fill={"request": (([1], [5]), block_size)})
226
+ )
227
+
228
+ connector.start_load_kv(
229
+ ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
230
+ )
231
+
232
+ assert torch.allclose(group_0_cache[1], torch.tensor(0.015))
233
+ assert torch.count_nonzero(group_0_cache[5]) == 0
234
+ assert torch.count_nonzero(group_1_cache[1]) == 0
235
+ assert torch.allclose(group_1_cache[5], torch.tensor(0.015))
236
+
237
+
238
+def test_decode_bench_connector_uses_per_group_block_sizes():
239
+ """Scheduler selects only the blocks covering each group's token span."""
240
+ vllm_config = create_vllm_config(
241
+ block_size=16,
242
+ max_num_batched_tokens=1000,
243
+ kv_connector="DecodeBenchConnector",
244
+ )
245
+ kv_cache_groups = [
246
+ KVCacheGroupSpec(
247
+ [f"group_{group_idx}_layer"],
248
+ FullAttentionSpec(
249
+ block_size=block_size,
250
+ num_kv_heads=1,
251
+ head_size=1,
252
+ dtype=torch.float32,
253
+ ),
254
+ )
255
+ for group_idx, block_size in enumerate((16, 32))
256
+ ]
257
+ connector = DecodeBenchConnector(
258
+ vllm_config,
259
+ KVConnectorRole.SCHEDULER,
260
+ KVCacheConfig(
261
+ num_blocks=8,
262
+ kv_cache_tensors=[],
263
+ kv_cache_groups=kv_cache_groups,
264
+ ),
265
+ )
266
+ request = MagicMock(request_id="request")
267
+ blocks = MagicMock()
268
+ blocks.get_block_ids.return_value = ([0, 1], [2, 3])
269
+
270
+ connector.update_state_after_alloc(request, blocks, num_external_tokens=17)
271
+ metadata = connector.build_connector_meta(MagicMock())
272
+
273
+ assert isinstance(metadata, DecodeBenchConnectorMetadata)
274
+ block_ids_per_group, num_tokens = metadata.reqs_to_fill["request"]
275
+ assert block_ids_per_group == ([0, 1], [2])
276
+ assert num_tokens == 17
277
+
278
+
279
def test_decode_bench_connector_no_refill():
280
"""Test that DecodeBenchConnector only fills once per request."""
281
block_size = 16
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+27
-30
@@ -69,8 +69,8 @@ class DecodeBenchConnectorMetadata(KVConnectorMetadata):
69
70
# request_id -> (block_ids_per_group, num_tokens_to_fill)
71
# block_ids_per_group is a tuple of lists, one per KV cache group
72
- # For standard attention: single group, e.g., ([1, 2, 3],)
73
- # For MLA: multiple groups, e.g., ([1, 2], [1, 2])
72
+ # One group: ([1, 2, 3],)
73
+ # Multiple groups: ([1, 2], [5, 6])
74
reqs_to_fill: dict[str, tuple[tuple[list[int], ...], int]]
75
76
@@ -95,9 +95,13 @@ class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA):
95
self.connector_worker: DecodeBenchConnectorWorker | None = None
96
97
if role == KVConnectorRole.SCHEDULER:
98
- self.connector_scheduler = DecodeBenchConnectorScheduler(vllm_config)
98
+ self.connector_scheduler = DecodeBenchConnectorScheduler(
99
+ vllm_config, kv_cache_config
100
+ )
101
elif role == KVConnectorRole.WORKER:
100
- self.connector_worker = DecodeBenchConnectorWorker(vllm_config)
102
+ self.connector_worker = DecodeBenchConnectorWorker(
103
+ vllm_config, kv_cache_config
104
+ )
105
106
# ==============================
107
# Worker-side methods
@@ -182,9 +186,11 @@ class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA):
186
class DecodeBenchConnectorScheduler:
187
"""Scheduler-side implementation for DecodeBenchConnector."""
188
185
- def __init__(self, vllm_config: "VllmConfig"):
189
+ def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
190
self.vllm_config = vllm_config
187
- self.block_size = vllm_config.cache_config.block_size
191
+ self.group_block_sizes = tuple(
192
+ group.kv_cache_spec.block_size for group in kv_cache_config.kv_cache_groups
193
+ )
194
195
# Track which requests have already been filled
196
self._filled_requests: set[str] = set()
@@ -236,8 +242,7 @@ class DecodeBenchConnectorScheduler:
242
Called after blocks are allocated. Store the block IDs so we can
243
fill them with dummy values.
244
239
- Supports both standard attention (single KV cache group) and MLA
240
- (multiple KV cache groups).
245
+ Supports both single- and multi-group KV cache configurations.
246
"""
247
req_id = request.request_id
248
@@ -246,18 +251,14 @@ class DecodeBenchConnectorScheduler:
251
252
# Get the block IDs that were allocated
253
# block_groups is a tuple of lists, one per KV cache group
249
- # For standard attention: 1 group
250
- # For MLA: multiple groups (one per attention type)
254
block_groups = blocks.get_block_ids()
255
253
- # Calculate how many blocks we need to fill
254
- # num_external_tokens are the tokens we said we'd provide
255
- num_blocks_to_fill = cdiv(num_external_tokens, self.block_size)
256
-
257
- # Extract the first num_blocks_to_fill blocks from each group
258
- # All groups should have the same block IDs for the same request
256
+ # Extract the blocks covering num_external_tokens from each group
257
block_ids_per_group = tuple(
260
- group_blocks[:num_blocks_to_fill] for group_blocks in block_groups
258
+ group_blocks[: cdiv(num_external_tokens, group_block_size)]
259
+ for group_blocks, group_block_size in zip(
260
+ block_groups, self.group_block_sizes, strict=True
261
+ )
262
)
263
264
# Store the blocks to fill for all group. _pending_fills doesn't need cleanup
@@ -269,9 +270,9 @@ class DecodeBenchConnectorScheduler:
270
self._filled_requests.add(req_id)
271
272
logger.debug(
272
- "DecodeBenchConnector: Allocated %d blocks across %d KV cache groups "
273
+ "DecodeBenchConnector: Allocated %s blocks across %d KV cache groups "
274
"for request %s",
274
- num_blocks_to_fill,
275
+ tuple(len(block_ids) for block_ids in block_ids_per_group),
276
len(block_groups),
277
req_id,
278
)
@@ -300,7 +301,7 @@ class DecodeBenchConnectorScheduler:
301
class DecodeBenchConnectorWorker:
302
"""Worker-side implementation for DecodeBenchConnector."""
303
303
- def __init__(self, vllm_config: "VllmConfig"):
304
+ def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
305
self.vllm_config = vllm_config
306
self.block_size = vllm_config.cache_config.block_size
307
@@ -314,17 +315,15 @@ class DecodeBenchConnectorWorker:
315
self.kv_caches: dict[str, torch.Tensor] | None = None
316
317
# Mapping from KV cache group index to list of layer names in that group
317
- self.group_to_layers: dict[int, list[str]] | None = None
318
+ self.group_to_layers = {
319
+ group_idx: list(group.layer_names)
320
+ for group_idx, group in enumerate(kv_cache_config.kv_cache_groups)
321
+ }
322
323
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
320
- """Store references to the KV cache tensors and build group mapping."""
324
+ """Store references to the KV cache tensors."""
325
self.kv_caches = kv_caches
326
323
- # For simplicity, assume all layers belong to group 0 (standard attention)
324
- # For MLA models with multiple groups, the metadata will handle the mapping
325
- # We just need to fill the blocks specified in the metadata
326
- self.group_to_layers = {0: list(kv_caches.keys())}
327
-
327
logger.debug(
328
"DecodeBenchConnector: Registered %d KV cache layers",
329
len(kv_caches),
@@ -337,13 +336,12 @@ class DecodeBenchConnectorWorker:
336
This simulates having a populated KV cache from a prefill phase,
337
allowing decode performance testing with larger context sizes.
338
340
- Supports both standard attention (single group) and MLA (multiple groups).
339
+ Supports both single- and multi-group KV cache configurations.
340
"""
341
if not metadata.reqs_to_fill:
342
return
343
344
assert self.kv_caches is not None, "KV caches must be registered before filling"
346
- assert self.group_to_layers is not None, "Group mapping must be initialized"
345
346
for req_id, (block_ids_per_group, num_tokens) in metadata.reqs_to_fill.items():
347
# Fill blocks for each KV cache group
@@ -372,7 +370,6 @@ class DecodeBenchConnectorWorker:
370
return
371
372
assert self.kv_caches is not None
375
- assert self.group_to_layers is not None
373
374
# Get the layers that belong to this group
375
layer_names = self.group_to_layers.get(group_idx, [])