iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][KV Connector] Safely fill circular buffers in DecodeBench (#54879)

Browse Source 
Signed-off-by: Jeff Ma <jeffjma@umich.edu>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Jeff (Junze) Ma
2026-09-03 15:16:22 +00:00
committed by GitHub
co-authored by OpenAI Codex
parent edc0fb7e03
commit 31e9c13685
2 changed files with 93 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+45
View File
@@ -31,6 +31,7 @@ from vllm.v1.core.kv_cache_utils import (
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.kv_cache_interface import (
CircularBufferSpec,
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
@@ -243,6 +244,50 @@ def test_decode_bench_connector_fills_each_hma_group():
assert torch.allclose(group_1_cache[5], torch.tensor(0.015))
def test_decode_bench_connector_zero_fills_circular_buffer_groups():
"""Circular buffers are zeroed to preserve packed non-floating metadata."""
num_gpu_blocks = 8
vllm_config = create_vllm_config(
block_size=16,
max_num_batched_tokens=1000,
kv_connector="DecodeBenchConnector",
)
kv_cache_config = KVCacheConfig(
num_blocks=num_gpu_blocks,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["circular_layer"],
CircularBufferSpec(
block_size=2,
num_kv_heads=1,
head_size=16,
head_size_v=0,
dtype=torch.bfloat16,
),
),
],
)
connector = DecodeBenchConnector(
vllm_config,
KVConnectorRole.WORKER,
kv_cache_config,
)
circular_cache = torch.ones(num_gpu_blocks, 16, dtype=torch.bfloat16)
connector.register_kv_caches({"circular_layer": circular_cache})
connector.bind_connector_metadata(
DecodeBenchConnectorMetadata(reqs_to_fill={"request": (([5],), 1)})
)
connector.start_load_kv(
ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
)
expected_circular = torch.ones_like(circular_cache)
expected_circular[5] = 0
assert torch.equal(circular_cache, expected_circular)
def test_decode_bench_connector_uses_per_group_block_sizes():
"""Scheduler selects only the blocks covering each group's token span."""
vllm_config = create_vllm_config(
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+48 -19
View File
@@ -51,6 +51,7 @@ from vllm.v1.core.kv_cache_utils import (
dcp_world_size_for_kv_cache_spec,
resolve_dcp_kv_block_size,
)
from vllm.v1.kv_cache_interface import CircularBufferSpec, iter_layer_specs
if TYPE_CHECKING:
from vllm.config import VllmConfig
@@ -82,9 +83,9 @@ class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA):
"""
A KV Connector for decode instance performance testing.
This connector fills the KV cache with dummy (non-zero) values to
emulate a prefill-decode disaggregated setting, enabling performance
testing of the decoder with larger input sequence lengths.
This connector fills the KV cache with dummy values to emulate a
prefill-decode disaggregated setting, enabling performance testing of the
decoder with larger input sequence lengths.
"""
def __init__(
@@ -329,6 +330,14 @@ class DecodeBenchConnectorWorker:
group_idx: list(group.layer_names)
for group_idx, group in enumerate(kv_cache_config.kv_cache_groups)
}
self._zero_fill_group_ids = {
group_idx
for group_idx, group in enumerate(kv_cache_config.kv_cache_groups)
if any(
isinstance(spec, CircularBufferSpec)
for spec in iter_layer_specs(group.kv_cache_spec)
)
}
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
"""Store references to the KV cache tensors."""
@@ -341,7 +350,7 @@ class DecodeBenchConnectorWorker:
def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata):
"""
Fill the allocated KV cache blocks with dummy (non-zero) values.
Fill the allocated KV cache blocks with dummy values.
This simulates having a populated KV cache from a prefill phase,
allowing decode performance testing with larger context sizes.
@@ -371,7 +380,7 @@ class DecodeBenchConnectorWorker:
def _fill_blocks(self, group_idx: int, block_ids: list[int], num_tokens: int):
"""
Fill specified blocks with dummy non-zero values for a specific KV cache group.
Fill specified blocks with dummy values for a specific KV cache group.
Args:
group_idx: The KV cache group index to fill
@@ -383,6 +392,14 @@ class DecodeBenchConnectorWorker:
assert self.kv_caches is not None
# Circular buffers may pack non-floating metadata alongside their
# floating-point state, so arbitrary fill values are not representation-safe.
fill_mean, fill_std = (
(0.0, 0.0)
if group_idx in self._zero_fill_group_ids
else (self.fill_mean, self.fill_std)
)
# Get the layers that belong to this group
layer_names = self.group_to_layers.get(group_idx, [])
@@ -404,12 +421,12 @@ class DecodeBenchConnectorWorker:
# dimension — so fill each tensor in its entirety with the same
# dummy values.
if isinstance(kv_cache, torch.Tensor):
self._fill_block_tensor(kv_cache, block_ids)
self._fill_block_tensor(kv_cache, block_ids, fill_mean, fill_std)
elif isinstance(kv_cache, (list, tuple)) and all(
isinstance(t, torch.Tensor) for t in kv_cache
):
for state_tensor in kv_cache:
self._fill_state_tensor(state_tensor)
self._fill_state_tensor(state_tensor, fill_mean, fill_std)
else:
logger.warning_once(
"DecodeBenchConnector: skipping fill for layer %s whose KV "
@@ -424,18 +441,26 @@ class DecodeBenchConnectorWorker:
"(mean=%.3f, std=%.3f)",
len(block_ids),
group_idx,
"random" if self.fill_std > 0 else "constant",
self.fill_mean,
self.fill_std,
"random" if fill_std > 0 else "constant",
fill_mean,
fill_std,
)
def _fill_block_tensor(self, kv_cache: torch.Tensor, block_ids: list[int]):
def _fill_block_tensor(
self,
kv_cache: torch.Tensor,
block_ids: list[int],
fill_mean: float,
fill_std: float,
):
"""Fill the requested block rows of a block-indexed KV cache tensor.
Args:
kv_cache: A KV cache tensor whose first dim is num_blocks.
block_ids: Block IDs to fill. IDs that are out of range for this
tensor's first dim are ignored.
fill_mean: Mean value for the fill.
fill_std: Standard deviation for the fill.
"""
# Convert block_ids to tensor on device
block_ids_tensor = torch.tensor(
@@ -451,11 +476,11 @@ class DecodeBenchConnectorWorker:
# Create fill values - either constant or random
block_shape = kv_cache.shape[1:]
if self.fill_std > 0:
if fill_std > 0:
# Random normal sampling
fill_values = torch.normal(
mean=self.fill_mean,
std=self.fill_std,
mean=fill_mean,
std=fill_std,
size=(len(valid_block_ids),) + block_shape,
dtype=kv_cache.dtype,
device=kv_cache.device,
@@ -464,7 +489,7 @@ class DecodeBenchConnectorWorker:
# Constant fill value
fill_values = torch.full(
(len(valid_block_ids),) + block_shape,
self.fill_mean,
fill_mean,
dtype=kv_cache.dtype,
device=kv_cache.device,
)
@@ -472,7 +497,9 @@ class DecodeBenchConnectorWorker:
# Batch fill operation
kv_cache[valid_block_ids] = fill_values
def _fill_state_tensor(self, kv_cache: torch.Tensor):
def _fill_state_tensor(
self, kv_cache: torch.Tensor, fill_mean: float, fill_std: float
):
"""Fill an entire non-block-indexed state tensor with dummy values.
Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention)
@@ -482,8 +509,10 @@ class DecodeBenchConnectorWorker:
Args:
kv_cache: A state tensor to fill in its entirety.
fill_mean: Mean value for the fill.
fill_std: Standard deviation for the fill.
"""
if self.fill_std > 0:
kv_cache.normal_(mean=self.fill_mean, std=self.fill_std)
if fill_std > 0:
kv_cache.normal_(mean=fill_mean, std=fill_std)
else:
kv_cache.fill_(self.fill_mean)
kv_cache.fill_(fill_mean)