[Bugfix][KV Connector] Safely fill circular buffers in DecodeBench (#54879)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: OpenAI Codex <noreply@openai.com>
Jeff (Junze) Ma committed
Sep 3, 2026 at 11:16 UTC
31e9c13685b484d4ba67dc4e73019b23402609c7
2 files changed
+93
-19
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+45
@@ -31,6 +31,7 @@ from vllm.v1.core.kv_cache_utils import (
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
+ CircularBufferSpec,
35
FullAttentionSpec,
36
KVCacheConfig,
37
KVCacheGroupSpec,
@@ -243,6 +244,50 @@ def test_decode_bench_connector_fills_each_hma_group():
244
assert torch.allclose(group_1_cache[5], torch.tensor(0.015))
245
246
247
+def test_decode_bench_connector_zero_fills_circular_buffer_groups():
248
+ """Circular buffers are zeroed to preserve packed non-floating metadata."""
249
+ num_gpu_blocks = 8
250
+ vllm_config = create_vllm_config(
251
+ block_size=16,
252
+ max_num_batched_tokens=1000,
253
+ kv_connector="DecodeBenchConnector",
254
+ )
255
+ kv_cache_config = KVCacheConfig(
256
+ num_blocks=num_gpu_blocks,
257
+ kv_cache_tensors=[],
258
+ kv_cache_groups=[
259
+ KVCacheGroupSpec(
260
+ ["circular_layer"],
261
+ CircularBufferSpec(
262
+ block_size=2,
263
+ num_kv_heads=1,
264
+ head_size=16,
265
+ head_size_v=0,
266
+ dtype=torch.bfloat16,
267
+ ),
268
+ ),
269
+ ],
270
+ )
271
+ connector = DecodeBenchConnector(
272
+ vllm_config,
273
+ KVConnectorRole.WORKER,
274
+ kv_cache_config,
275
+ )
276
+ circular_cache = torch.ones(num_gpu_blocks, 16, dtype=torch.bfloat16)
277
+ connector.register_kv_caches({"circular_layer": circular_cache})
278
+ connector.bind_connector_metadata(
279
+ DecodeBenchConnectorMetadata(reqs_to_fill={"request": (([5],), 1)})
280
+ )
281
+
282
+ connector.start_load_kv(
283
+ ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
284
+ )
285
+
286
+ expected_circular = torch.ones_like(circular_cache)
287
+ expected_circular[5] = 0
288
+ assert torch.equal(circular_cache, expected_circular)
289
+
290
+
291
def test_decode_bench_connector_uses_per_group_block_sizes():
292
"""Scheduler selects only the blocks covering each group's token span."""
293
vllm_config = create_vllm_config(
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+48
-19
@@ -51,6 +51,7 @@ from vllm.v1.core.kv_cache_utils import (
51
dcp_world_size_for_kv_cache_spec,
52
resolve_dcp_kv_block_size,
53
)
54
+from vllm.v1.kv_cache_interface import CircularBufferSpec, iter_layer_specs
55
56
if TYPE_CHECKING:
57
from vllm.config import VllmConfig
@@ -82,9 +83,9 @@ class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA):
83
"""
84
A KV Connector for decode instance performance testing.
85
85
- This connector fills the KV cache with dummy (non-zero) values to
86
- emulate a prefill-decode disaggregated setting, enabling performance
87
- testing of the decoder with larger input sequence lengths.
86
+ This connector fills the KV cache with dummy values to emulate a
87
+ prefill-decode disaggregated setting, enabling performance testing of the
88
+ decoder with larger input sequence lengths.
89
"""
90
91
def __init__(
@@ -329,6 +330,14 @@ class DecodeBenchConnectorWorker:
330
group_idx: list(group.layer_names)
331
for group_idx, group in enumerate(kv_cache_config.kv_cache_groups)
332
}
333
+ self._zero_fill_group_ids = {
334
+ group_idx
335
+ for group_idx, group in enumerate(kv_cache_config.kv_cache_groups)
336
+ if any(
337
+ isinstance(spec, CircularBufferSpec)
338
+ for spec in iter_layer_specs(group.kv_cache_spec)
339
+ )
340
+ }
341
342
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
343
"""Store references to the KV cache tensors."""
@@ -341,7 +350,7 @@ class DecodeBenchConnectorWorker:
350
351
def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata):
352
"""
344
- Fill the allocated KV cache blocks with dummy (non-zero) values.
353
+ Fill the allocated KV cache blocks with dummy values.
354
355
This simulates having a populated KV cache from a prefill phase,
356
allowing decode performance testing with larger context sizes.
@@ -371,7 +380,7 @@ class DecodeBenchConnectorWorker:
380
381
def _fill_blocks(self, group_idx: int, block_ids: list[int], num_tokens: int):
382
"""
374
- Fill specified blocks with dummy non-zero values for a specific KV cache group.
383
+ Fill specified blocks with dummy values for a specific KV cache group.
384
385
Args:
386
group_idx: The KV cache group index to fill
@@ -383,6 +392,14 @@ class DecodeBenchConnectorWorker:
392
393
assert self.kv_caches is not None
394
395
+ # Circular buffers may pack non-floating metadata alongside their
396
+ # floating-point state, so arbitrary fill values are not representation-safe.
397
+ fill_mean, fill_std = (
398
+ (0.0, 0.0)
399
+ if group_idx in self._zero_fill_group_ids
400
+ else (self.fill_mean, self.fill_std)
401
+ )
402
+
403
# Get the layers that belong to this group
404
layer_names = self.group_to_layers.get(group_idx, [])
405
@@ -404,12 +421,12 @@ class DecodeBenchConnectorWorker:
421
# dimension — so fill each tensor in its entirety with the same
422
# dummy values.
423
if isinstance(kv_cache, torch.Tensor):
407
- self._fill_block_tensor(kv_cache, block_ids)
424
+ self._fill_block_tensor(kv_cache, block_ids, fill_mean, fill_std)
425
elif isinstance(kv_cache, (list, tuple)) and all(
426
isinstance(t, torch.Tensor) for t in kv_cache
427
):
428
for state_tensor in kv_cache:
412
- self._fill_state_tensor(state_tensor)
429
+ self._fill_state_tensor(state_tensor, fill_mean, fill_std)
430
else:
431
logger.warning_once(
432
"DecodeBenchConnector: skipping fill for layer %s whose KV "
@@ -424,18 +441,26 @@ class DecodeBenchConnectorWorker:
441
"(mean=%.3f, std=%.3f)",
442
len(block_ids),
443
group_idx,
427
- "random" if self.fill_std > 0 else "constant",
428
- self.fill_mean,
429
- self.fill_std,
444
+ "random" if fill_std > 0 else "constant",
445
+ fill_mean,
446
+ fill_std,
447
)
448
432
- def _fill_block_tensor(self, kv_cache: torch.Tensor, block_ids: list[int]):
449
+ def _fill_block_tensor(
450
+ self,
451
+ kv_cache: torch.Tensor,
452
+ block_ids: list[int],
453
+ fill_mean: float,
454
+ fill_std: float,
455
+ ):
456
"""Fill the requested block rows of a block-indexed KV cache tensor.
457
458
Args:
459
kv_cache: A KV cache tensor whose first dim is num_blocks.
460
block_ids: Block IDs to fill. IDs that are out of range for this
461
tensor's first dim are ignored.
462
+ fill_mean: Mean value for the fill.
463
+ fill_std: Standard deviation for the fill.
464
"""
465
# Convert block_ids to tensor on device
466
block_ids_tensor = torch.tensor(
@@ -451,11 +476,11 @@ class DecodeBenchConnectorWorker:
476
477
# Create fill values - either constant or random
478
block_shape = kv_cache.shape[1:]
454
- if self.fill_std > 0:
479
+ if fill_std > 0:
480
# Random normal sampling
481
fill_values = torch.normal(
457
- mean=self.fill_mean,
458
- std=self.fill_std,
482
+ mean=fill_mean,
483
+ std=fill_std,
484
size=(len(valid_block_ids),) + block_shape,
485
dtype=kv_cache.dtype,
486
device=kv_cache.device,
@@ -464,7 +489,7 @@ class DecodeBenchConnectorWorker:
489
# Constant fill value
490
fill_values = torch.full(
491
(len(valid_block_ids),) + block_shape,
467
- self.fill_mean,
492
+ fill_mean,
493
dtype=kv_cache.dtype,
494
device=kv_cache.device,
495
)
@@ -472,7 +497,9 @@ class DecodeBenchConnectorWorker:
497
# Batch fill operation
498
kv_cache[valid_block_ids] = fill_values
499
475
- def _fill_state_tensor(self, kv_cache: torch.Tensor):
500
+ def _fill_state_tensor(
501
+ self, kv_cache: torch.Tensor, fill_mean: float, fill_std: float
502
+ ):
503
"""Fill an entire non-block-indexed state tensor with dummy values.
504
505
Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention)
@@ -482,8 +509,10 @@ class DecodeBenchConnectorWorker:
509
510
Args:
511
kv_cache: A state tensor to fill in its entirety.
512
+ fill_mean: Mean value for the fill.
513
+ fill_std: Standard deviation for the fill.
514
"""
486
- if self.fill_std > 0:
487
- kv_cache.normal_(mean=self.fill_mean, std=self.fill_std)
515
+ if fill_std > 0:
516
+ kv_cache.normal_(mean=fill_mean, std=fill_std)
517
else:
489
- kv_cache.fill_(self.fill_mean)
518
+ kv_cache.fill_(fill_mean)