[Bugfix][KV Connector] Fix DecodeBenchConnector prefix block selection (#54878)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: OpenAI Codex <noreply@openai.com>
Jeff (Junze) Ma committed
Sep 4, 2026 at 00:05 UTC
e35298628fceb490d4ca2a21c0a6e8f63d3c6740
2 files changed
+230
-28
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+178
-16
@@ -7,8 +7,6 @@ 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
-
10
import pytest
11
import torch
12
@@ -24,9 +22,11 @@ from vllm.forward_context import ForwardContext
22
from vllm.utils.hashing import sha256
23
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
24
from vllm.v1.core.kv_cache_utils import (
25
+ BlockHash,
26
KVCacheBlock,
27
get_request_block_hasher,
28
init_none_hash,
29
+ make_block_hash_with_group_id,
30
)
31
from vllm.v1.core.sched.output import SchedulerOutput
32
from vllm.v1.core.sched.scheduler import Scheduler
@@ -194,6 +194,43 @@ def test_decode_bench_connector_basic():
194
assert torch.allclose(block_data, torch.tensor(0.015))
195
196
197
+def test_decode_bench_connector_prefix_hit_fills_only_suffix_blocks():
198
+ """A local prefix-cache hit shifts the fill to the suffix blocks."""
199
+ block_size = 16
200
+ runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=100)
201
+
202
+ prefix = list(range(1, 1 + block_size * 3))
203
+ req1 = runner.new_request(prefix)
204
+ _, metadata = runner.run_single_step()
205
+ (prefix_block_ids,), _ = metadata.reqs_to_fill[req1.request_id]
206
+ assert len(prefix_block_ids) == 3
207
+
208
+ scheduler_output = runner.scheduler.schedule()
209
+ runner.scheduler.update_from_output(
210
+ scheduler_output,
211
+ create_model_runner_output(
212
+ reqs=runner.scheduler.running,
213
+ token_id=EOS_TOKEN_ID,
214
+ use_eos=True,
215
+ ),
216
+ )
217
+ assert not runner.scheduler.running
218
+
219
+ suffix = list(range(1000, 1000 + block_size * 2))
220
+ req2 = runner.new_request(prefix + suffix)
221
+ _, metadata = runner.run_single_step()
222
+ (fill_block_ids,), num_tokens_to_fill = metadata.reqs_to_fill[req2.request_id]
223
+
224
+ req2_block_ids = runner.scheduler.kv_cache_manager.get_block_ids(req2.request_id)[0]
225
+ assert req2_block_ids[:3] == prefix_block_ids
226
+ assert num_tokens_to_fill == len(suffix) - 1
227
+ assert fill_block_ids == req2_block_ids[3:5]
228
+ assert set(fill_block_ids).isdisjoint(prefix_block_ids)
229
+ for kv_cache in runner.kv_caches.values():
230
+ for block_id in fill_block_ids:
231
+ assert torch.allclose(kv_cache[block_id], torch.tensor(0.015))
232
+
233
+
234
def test_decode_bench_connector_fills_each_hma_group():
235
"""Each cache group is filled using its own block IDs."""
236
block_size = 16
@@ -288,8 +325,8 @@ def test_decode_bench_connector_zero_fills_circular_buffer_groups():
325
assert torch.equal(circular_cache, expected_circular)
326
327
291
-def test_decode_bench_connector_uses_per_group_block_sizes():
292
- """Scheduler selects only the blocks covering each group's token span."""
328
+def test_decode_bench_connector_selects_external_blocks_per_group():
329
+ """External fills follow the local prefix using each group's block size."""
330
vllm_config = create_vllm_config(
331
block_size=16,
332
max_num_batched_tokens=1000,
@@ -297,15 +334,23 @@ def test_decode_bench_connector_uses_per_group_block_sizes():
334
)
335
kv_cache_groups = [
336
KVCacheGroupSpec(
300
- [f"group_{group_idx}_layer"],
337
+ ["full_attention"],
338
FullAttentionSpec(
302
- block_size=block_size,
339
+ block_size=16,
340
num_kv_heads=1,
341
head_size=1,
342
dtype=torch.float32,
343
),
307
- )
308
- for group_idx, block_size in enumerate((16, 32))
344
+ ),
345
+ KVCacheGroupSpec(
346
+ ["compressor_ring"],
347
+ CircularBufferSpec(
348
+ block_size=32,
349
+ num_kv_heads=1,
350
+ head_size=1,
351
+ dtype=torch.float32,
352
+ ),
353
+ ),
354
]
355
connector = DecodeBenchConnector(
356
vllm_config,
@@ -316,17 +361,134 @@ def test_decode_bench_connector_uses_per_group_block_sizes():
361
kv_cache_groups=kv_cache_groups,
362
),
363
)
319
- request = MagicMock(request_id="request")
320
- blocks = MagicMock()
321
- blocks.get_block_ids.return_value = ([0, 1], [2, 3])
364
+ num_local_tokens = 32
365
+ num_external_tokens = 17
366
+ request = create_request(
367
+ request_id=1,
368
+ num_tokens=num_local_tokens + num_external_tokens + 1,
369
+ block_size=16,
370
+ )
371
+ blocks = KVCacheBlocks(
372
+ (
373
+ [
374
+ KVCacheBlock(block_id=0),
375
+ KVCacheBlock(block_id=1),
376
+ # Full external blocks may already be hashed by allocate_slots().
377
+ KVCacheBlock(
378
+ block_id=2,
379
+ _block_hash=make_block_hash_with_group_id(
380
+ BlockHash(b"external-full-block"), 0
381
+ ),
382
+ _block_hash_num_tokens=48,
383
+ ),
384
+ KVCacheBlock(block_id=3),
385
+ ],
386
+ [KVCacheBlock(block_id=4)],
387
+ )
388
+ )
389
323
- connector.update_state_after_alloc(request, blocks, num_external_tokens=17)
324
- metadata = connector.build_connector_meta(MagicMock())
390
+ connector.update_state_after_alloc(request, blocks, num_external_tokens)
391
+ metadata = connector.build_connector_meta(SchedulerOutput.make_empty())
392
393
assert isinstance(metadata, DecodeBenchConnectorMetadata)
327
- block_ids_per_group, num_tokens = metadata.reqs_to_fill["request"]
328
- assert block_ids_per_group == ([0, 1], [2])
329
- assert num_tokens == 17
394
+ block_ids_per_group, num_tokens = metadata.reqs_to_fill[request.request_id]
395
+ assert block_ids_per_group == ([2, 3], [4])
396
+ assert num_tokens == num_external_tokens
397
+
398
+
399
+def test_decode_bench_connector_selects_unaligned_external_block():
400
+ """An external interval includes the block containing its first token."""
401
+ block_size = 32
402
+ vllm_config = create_vllm_config(
403
+ block_size=block_size,
404
+ max_num_batched_tokens=1000,
405
+ kv_connector="DecodeBenchConnector",
406
+ )
407
+ kv_cache_spec = FullAttentionSpec(
408
+ block_size=block_size,
409
+ num_kv_heads=1,
410
+ head_size=1,
411
+ dtype=torch.float32,
412
+ )
413
+ connector = DecodeBenchConnector(
414
+ vllm_config,
415
+ KVConnectorRole.SCHEDULER,
416
+ KVCacheConfig(
417
+ num_blocks=4,
418
+ kv_cache_tensors=[],
419
+ kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
420
+ ),
421
+ )
422
+ num_local_tokens = 23
423
+ num_external_tokens = 17
424
+ request = create_request(
425
+ request_id=1,
426
+ num_tokens=num_local_tokens + num_external_tokens + 1,
427
+ block_size=block_size,
428
+ )
429
+ blocks = KVCacheBlocks(([KVCacheBlock(block_id=1), KVCacheBlock(block_id=2)],))
430
+
431
+ connector.update_state_after_alloc(request, blocks, num_external_tokens)
432
+ metadata = connector.build_connector_meta(SchedulerOutput.make_empty())
433
+
434
+ assert isinstance(metadata, DecodeBenchConnectorMetadata)
435
+ assert metadata.reqs_to_fill[request.request_id] == (
436
+ ([1, 2],),
437
+ num_external_tokens,
438
+ )
439
+
440
+
441
+def test_decode_bench_connector_skips_null_padding():
442
+ """External fills exclude null padding from a sliding-window group."""
443
+ block_size = 16
444
+ vllm_config = create_vllm_config(
445
+ block_size=block_size,
446
+ max_num_batched_tokens=1000,
447
+ kv_connector="DecodeBenchConnector",
448
+ )
449
+ kv_cache_spec = SlidingWindowSpec(
450
+ block_size=block_size,
451
+ num_kv_heads=1,
452
+ head_size=1,
453
+ dtype=torch.float32,
454
+ sliding_window=2 * block_size,
455
+ )
456
+ connector = DecodeBenchConnector(
457
+ vllm_config,
458
+ KVConnectorRole.SCHEDULER,
459
+ KVCacheConfig(
460
+ num_blocks=8,
461
+ kv_cache_tensors=[],
462
+ kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
463
+ ),
464
+ )
465
+ num_external_tokens = 4 * block_size
466
+ request = create_request(
467
+ request_id=1,
468
+ num_tokens=num_external_tokens + 1,
469
+ block_size=block_size,
470
+ )
471
+ null_block = KVCacheBlock(block_id=0, is_null=True)
472
+ blocks = KVCacheBlocks(
473
+ (
474
+ [
475
+ null_block,
476
+ null_block,
477
+ KVCacheBlock(block_id=1),
478
+ KVCacheBlock(block_id=2),
479
+ KVCacheBlock(block_id=3),
480
+ ],
481
+ )
482
+ )
483
+
484
+ connector.update_state_after_alloc(request, blocks, num_external_tokens)
485
+ metadata = connector.build_connector_meta(SchedulerOutput.make_empty())
486
+
487
+ assert isinstance(metadata, DecodeBenchConnectorMetadata)
488
+ assert metadata.reqs_to_fill[request.request_id] == (
489
+ ([1, 2],),
490
+ num_external_tokens,
491
+ )
492
493
494
def test_decode_bench_connector_uses_per_group_dcp_block_sizes():
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+52
-12
@@ -193,6 +193,7 @@ class DecodeBenchConnectorScheduler:
193
194
def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
195
dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
196
+ self.kv_cache_groups = kv_cache_config.kv_cache_groups
197
self.group_block_sizes = tuple(
198
resolve_dcp_kv_block_size(
199
group.kv_cache_spec,
@@ -201,7 +202,7 @@ class DecodeBenchConnectorScheduler:
202
dcp_world_size,
203
),
204
)
204
- for group in kv_cache_config.kv_cache_groups
205
+ for group in self.kv_cache_groups
206
)
207
208
# Track which requests have already been filled
@@ -261,17 +262,56 @@ class DecodeBenchConnectorScheduler:
262
if num_external_tokens == 0:
263
return
264
264
- # Get the block IDs that were allocated
265
- # block_groups is a tuple of lists, one per KV cache group
266
- block_groups = blocks.get_block_ids()
267
-
268
- # Extract the blocks covering the external tokens from each group
269
- block_ids_per_group = tuple(
270
- group_blocks[: cdiv(num_external_tokens, group_block_size)]
271
- for group_blocks, group_block_size in zip(
272
- block_groups, self.group_block_sizes, strict=True
265
+ total_computed_tokens = request.num_tokens - 1
266
+ num_local_computed_tokens = total_computed_tokens - num_external_tokens
267
+ block_ids_per_group_list: list[list[int]] = []
268
+ for group_idx, (
269
+ group,
270
+ group_blocks,
271
+ group_block_size,
272
+ ) in enumerate(
273
+ zip(
274
+ self.kv_cache_groups,
275
+ blocks.blocks,
276
+ self.group_block_sizes,
277
+ strict=True,
278
)
274
- )
279
+ ):
280
+ is_circular_buffer = all(
281
+ isinstance(spec, CircularBufferSpec)
282
+ for spec in iter_layer_specs(group.kv_cache_spec)
283
+ )
284
+ if not is_circular_buffer:
285
+ num_computed_blocks = cdiv(total_computed_tokens, group_block_size)
286
+ external_block_start = num_local_computed_tokens // group_block_size
287
+ assert (
288
+ 0
289
+ <= external_block_start
290
+ <= num_computed_blocks
291
+ <= len(group_blocks)
292
+ ), (
293
+ "DecodeBenchConnector block range exceeds allocated blocks: "
294
+ f"request={req_id}, group={group_idx}, "
295
+ f"range=[{external_block_start}, {num_computed_blocks}), "
296
+ f"allocated={len(group_blocks)}"
297
+ )
298
+ selected_blocks = group_blocks[external_block_start:num_computed_blocks]
299
+ else:
300
+ selected_blocks = group_blocks
301
+
302
+ block_ids = [
303
+ block.block_id for block in selected_blocks if not block.is_null
304
+ ]
305
+ if not block_ids:
306
+ logger.warning(
307
+ "DecodeBenchConnector: No blocks selected for KV cache group "
308
+ "%d with %d external tokens for request %s",
309
+ group_idx,
310
+ num_external_tokens,
311
+ req_id,
312
+ )
313
+ block_ids_per_group_list.append(block_ids)
314
+ block_ids_per_group = tuple(block_ids_per_group_list)
315
316
# Store the blocks to fill for all group. _pending_fills doesn't need cleanup
317
# as it's cleared after build_connector_meta
@@ -286,7 +326,7 @@ class DecodeBenchConnectorScheduler:
326
"DecodeBenchConnector: Selected %d total blocks across %d KV cache "
327
"groups for request %s (per-group counts: %s)",
328
sum(block_counts),
289
- len(block_groups),
329
+ len(blocks.blocks),
330
req_id,
331
", ".join(map(str, block_counts)),
332
)