diff --git a/tests/v1/kv_connector/unit/test_decode_bench_connector.py b/tests/v1/kv_connector/unit/test_decode_bench_connector.py index 6e7330606c..66e403a31f 100644 --- a/tests/v1/kv_connector/unit/test_decode_bench_connector.py +++ b/tests/v1/kv_connector/unit/test_decode_bench_connector.py @@ -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( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py index c7c1078db1..b9f4cc15fe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py @@ -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)