[Bugfix][KV Connector] Fix DecodeBench DCP block selection (#54679)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: OpenAI Codex <noreply@openai.com>
Jeff (Junze) Ma committed
Sep 1, 2026 at 16:31 UTC
18c53727cebb588771c5e2f62a207137b0d6dffd
2 files changed
+90
-13
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+66
-1
@@ -22,18 +22,26 @@ from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import
22
)
23
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
25
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
26
+from vllm.v1.core.kv_cache_utils import (
27
+ KVCacheBlock,
28
+ get_request_block_hasher,
29
+ init_none_hash,
30
+)
31
+from vllm.v1.core.sched.output import SchedulerOutput
32
from vllm.v1.core.sched.scheduler import Scheduler
33
from vllm.v1.kv_cache_interface import (
34
FullAttentionSpec,
35
KVCacheConfig,
36
KVCacheGroupSpec,
37
+ SlidingWindowSpec,
38
)
39
from vllm.v1.request import Request
40
41
from .utils import (
42
EOS_TOKEN_ID,
43
create_model_runner_output,
44
+ create_request,
45
create_scheduler,
46
create_vllm_config,
47
)
@@ -276,6 +284,63 @@ def test_decode_bench_connector_uses_per_group_block_sizes():
284
assert num_tokens == 17
285
286
287
+def test_decode_bench_connector_uses_per_group_dcp_block_sizes():
288
+ """DCP scales full-attention blocks but not sliding-window blocks."""
289
+ block_size = 16
290
+ dcp_world_size = 2
291
+ vllm_config = create_vllm_config(
292
+ block_size=block_size,
293
+ kv_connector="DecodeBenchConnector",
294
+ )
295
+ vllm_config.parallel_config.decode_context_parallel_size = dcp_world_size
296
+ full_attention = FullAttentionSpec(
297
+ block_size=block_size,
298
+ num_kv_heads=1,
299
+ head_size=1,
300
+ dtype=torch.float32,
301
+ )
302
+ sliding_window = SlidingWindowSpec(
303
+ block_size=block_size,
304
+ num_kv_heads=1,
305
+ head_size=1,
306
+ dtype=torch.float32,
307
+ sliding_window=2 * block_size,
308
+ )
309
+ connector = DecodeBenchConnector(
310
+ vllm_config,
311
+ KVConnectorRole.SCHEDULER,
312
+ KVCacheConfig(
313
+ num_blocks=8,
314
+ kv_cache_tensors=[],
315
+ kv_cache_groups=[
316
+ KVCacheGroupSpec(["full_attention"], full_attention),
317
+ KVCacheGroupSpec(["sliding_window"], sliding_window),
318
+ ],
319
+ ),
320
+ )
321
+ num_external_tokens = block_size * dcp_world_size + 1
322
+ request = create_request(
323
+ request_id=1,
324
+ num_tokens=num_external_tokens + 1,
325
+ block_size=block_size,
326
+ )
327
+ blocks = KVCacheBlocks(
328
+ (
329
+ [KVCacheBlock(block_id=i) for i in (1, 2, 3)],
330
+ [KVCacheBlock(block_id=i) for i in (5, 6, 7, 8)],
331
+ )
332
+ )
333
+
334
+ connector.update_state_after_alloc(request, blocks, num_external_tokens)
335
+ metadata = connector.build_connector_meta(SchedulerOutput.make_empty())
336
+
337
+ assert isinstance(metadata, DecodeBenchConnectorMetadata)
338
+ assert metadata.reqs_to_fill[request.request_id] == (
339
+ ([1, 2], [5, 6, 7]),
340
+ num_external_tokens,
341
+ )
342
+
343
+
344
def test_decode_bench_connector_no_refill():
345
"""Test that DecodeBenchConnector only fills once per request."""
346
block_size = 16
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+24
-12
@@ -47,6 +47,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
47
from vllm.logger import init_logger
48
from vllm.utils.math_utils import cdiv
49
from vllm.v1.attention.backend import AttentionMetadata
50
+from vllm.v1.core.kv_cache_utils import (
51
+ dcp_world_size_for_kv_cache_spec,
52
+ resolve_dcp_kv_block_size,
53
+)
54
55
if TYPE_CHECKING:
56
from vllm.config import VllmConfig
@@ -187,9 +191,16 @@ class DecodeBenchConnectorScheduler:
191
"""Scheduler-side implementation for DecodeBenchConnector."""
192
193
def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
190
- self.vllm_config = vllm_config
194
+ dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
195
self.group_block_sizes = tuple(
192
- group.kv_cache_spec.block_size for group in kv_cache_config.kv_cache_groups
196
+ resolve_dcp_kv_block_size(
197
+ group.kv_cache_spec,
198
+ dcp_world_size_for_kv_cache_spec(
199
+ group.kv_cache_spec,
200
+ dcp_world_size,
201
+ ),
202
+ )
203
+ for group in kv_cache_config.kv_cache_groups
204
)
205
206
# Track which requests have already been filled
@@ -253,7 +264,7 @@ class DecodeBenchConnectorScheduler:
264
# block_groups is a tuple of lists, one per KV cache group
265
block_groups = blocks.get_block_ids()
266
256
- # Extract the blocks covering num_external_tokens from each group
267
+ # Extract the blocks covering the external tokens from each group
268
block_ids_per_group = tuple(
269
group_blocks[: cdiv(num_external_tokens, group_block_size)]
270
for group_blocks, group_block_size in zip(
@@ -269,12 +280,14 @@ class DecodeBenchConnectorScheduler:
280
)
281
self._filled_requests.add(req_id)
282
283
+ block_counts = tuple(len(group) for group in block_ids_per_group)
284
logger.debug(
273
- "DecodeBenchConnector: Allocated %s blocks across %d KV cache groups "
274
- "for request %s",
275
- tuple(len(block_ids) for block_ids in block_ids_per_group),
285
+ "DecodeBenchConnector: Selected %d total blocks across %d KV cache "
286
+ "groups for request %s (per-group counts: %s)",
287
+ sum(block_counts),
288
len(block_groups),
289
req_id,
290
+ ", ".join(map(str, block_counts)),
291
)
292
293
def build_connector_meta(
@@ -302,9 +315,6 @@ class DecodeBenchConnectorWorker:
315
"""Worker-side implementation for DecodeBenchConnector."""
316
317
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
-
318
# Get fill parameters from extra config
319
kv_transfer_config = vllm_config.kv_transfer_config
320
assert kv_transfer_config is not None
@@ -348,13 +358,15 @@ class DecodeBenchConnectorWorker:
358
for group_idx, block_ids in enumerate(block_ids_per_group):
359
self._fill_blocks(group_idx, block_ids, num_tokens)
360
361
+ block_counts = tuple(len(group) for group in block_ids_per_group)
362
logger.debug(
352
- "DecodeBenchConnector: Filled %d blocks (%d tokens) across %d groups "
353
- "for request %s",
354
- len(block_ids_per_group[0]) if block_ids_per_group else 0,
363
+ "DecodeBenchConnector: Filled %d total blocks (%d tokens) across "
364
+ "%d groups for request %s (per-group counts: %s)",
365
+ sum(block_counts),
366
num_tokens,
367
len(block_ids_per_group),
368
req_id,
369
+ ", ".join(map(str, block_counts)),
370
)
371
372
def _fill_blocks(self, group_idx: int, block_ids: list[int], num_tokens: int):