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 0f835b1eeb..4485c2ecd0 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 @@ -387,40 +387,28 @@ class DecodeBenchConnectorWorker: kv_cache = self.kv_caches[layer_name] - # Convert block_ids to tensor on device - block_ids_tensor = torch.tensor( - block_ids, dtype=torch.long, device=kv_cache.device - ) - - # Filter invalid block IDs - valid_mask = block_ids_tensor < kv_cache.shape[0] - valid_block_ids = block_ids_tensor[valid_mask] - - if len(valid_block_ids) == 0: - continue - - # Create fill values - either constant or random - block_shape = kv_cache.shape[1:] - if self.fill_std > 0: - # Random normal sampling - fill_values = torch.normal( - mean=self.fill_mean, - std=self.fill_std, - size=(len(valid_block_ids),) + block_shape, - dtype=kv_cache.dtype, - device=kv_cache.device, - ) + # Attention layers store KV as a single block-indexed tensor whose + # first dim is num_blocks; fill the requested block rows. Hybrid / + # linear-attention layers (e.g. Mamba, Kimi Delta Attention) store + # their state as a list/tuple of tensors that are NOT block-indexed + # — each tensor is a single state buffer with no num_blocks + # 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) + 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) else: - # Constant fill value - fill_values = torch.full( - (len(valid_block_ids),) + block_shape, - self.fill_mean, - dtype=kv_cache.dtype, - device=kv_cache.device, + logger.warning_once( + "DecodeBenchConnector: skipping fill for layer %s whose KV " + "cache is %s, not a tensor or a list/tuple of tensors.", + layer_name, + type(kv_cache).__name__, ) - - # Batch fill operation - kv_cache[valid_block_ids] = fill_values + continue logger.debug( "DecodeBenchConnector: Filled %d blocks in group %d with %s values " @@ -431,3 +419,62 @@ class DecodeBenchConnectorWorker: self.fill_mean, self.fill_std, ) + + def _fill_block_tensor(self, kv_cache: torch.Tensor, block_ids: list[int]): + """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. + """ + # Convert block_ids to tensor on device + block_ids_tensor = torch.tensor( + block_ids, dtype=torch.long, device=kv_cache.device + ) + + # Filter invalid block IDs + valid_mask = block_ids_tensor < kv_cache.shape[0] + valid_block_ids = block_ids_tensor[valid_mask] + + if len(valid_block_ids) == 0: + return + + # Create fill values - either constant or random + block_shape = kv_cache.shape[1:] + if self.fill_std > 0: + # Random normal sampling + fill_values = torch.normal( + mean=self.fill_mean, + std=self.fill_std, + size=(len(valid_block_ids),) + block_shape, + dtype=kv_cache.dtype, + device=kv_cache.device, + ) + else: + # Constant fill value + fill_values = torch.full( + (len(valid_block_ids),) + block_shape, + self.fill_mean, + dtype=kv_cache.dtype, + device=kv_cache.device, + ) + + # Batch fill operation + kv_cache[valid_block_ids] = fill_values + + def _fill_state_tensor(self, kv_cache: torch.Tensor): + """Fill an entire non-block-indexed state tensor with dummy values. + + Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention) + store their per-layer state as tensors with no num_blocks dimension, + so the whole tensor is filled with the same constant or random values + used for block fills, rather than selected block rows. + + Args: + kv_cache: A state tensor to fill in its entirety. + """ + if self.fill_std > 0: + kv_cache.normal_(mean=self.fill_mean, std=self.fill_std) + else: + kv_cache.fill_(self.fill_mean)