diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index 2b35718e7..1f3d4acab 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -105,6 +105,98 @@ def test_manager_extraction(kv_config, mock_connector, common_constants): assert data["layer_blocks"]["key_cache"][0].shape == expected_shape +def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants): + """Test extraction with tuple layout.""" + num_layers = common_constants["num_layers"] + block_size = common_constants["block_size"] + num_heads = common_constants["num_heads"] + head_dim = common_constants["head_dim"] + seq_len = common_constants["seq_len"] + req_id = common_constants["req_id"] + + num_blocks = 10 + kv_caches = [] + for _ in range(num_layers): + k_cache = torch.randn(num_blocks, block_size, num_heads, head_dim) + v_cache = torch.randn(num_blocks, block_size, num_heads, head_dim) + kv_caches.append((k_cache, v_cache)) + + block_ids = [1, 3, 5] + finished_reqs = {req_id: {"block_ids": block_ids, "seq_len": seq_len}} + + manager = OmniKVTransferManager(kv_config) + manager._connector = mock_connector + + processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32") + assert req_id in processed + + full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}" + expected_key = f"stage1->stage2:{full_request_id}" + assert expected_key in mock_connector.store + + data = mock_connector.store[expected_key] + expected_shape = (seq_len, num_heads, head_dim) + for idx in range(len(kv_caches)): + assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape + assert data["layer_blocks"]["value_cache"][idx].shape == expected_shape + + +def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector, common_constants): + """Mismatched key/value block counts should not crash extraction.""" + block_size = common_constants["block_size"] + num_heads = common_constants["num_heads"] + head_dim = common_constants["head_dim"] + req_id = common_constants["req_id"] + + key_blocks = torch.randn(3, block_size, num_heads, head_dim) + value_blocks = torch.randn(2, block_size, num_heads, head_dim) + kv_caches = [(key_blocks, value_blocks)] + + finished_reqs = {req_id: {"block_ids": [0, 1, 2], "seq_len": 32}} + + manager = OmniKVTransferManager(kv_config) + manager._connector = mock_connector + + processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32") + assert req_id in processed + + full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}" + expected_key = f"stage1->stage2:{full_request_id}" + assert expected_key in mock_connector.store + + data = mock_connector.store[expected_key] + expected_shape = (2 * block_size, num_heads, head_dim) + assert data["layer_blocks"]["key_cache"][0].shape == expected_shape + assert data["layer_blocks"]["value_cache"][0].shape == expected_shape + + +@pytest.mark.parametrize( + "invalid_case", + ["invalid_stacked_shape", "invalid_tuple_length", "non_tensor_entries"], +) +def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants, invalid_case): + """_normalize_layer_kv should reject malformed KV representations.""" + block_size = common_constants["block_size"] + num_heads = common_constants["num_heads"] + head_dim = common_constants["head_dim"] + req_id = common_constants["req_id"] + + if invalid_case == "invalid_stacked_shape": + layer_kv = torch.randn(3, block_size, num_heads, head_dim) + elif invalid_case == "invalid_tuple_length": + layer_kv = ( + torch.randn(2, block_size, num_heads, head_dim), + torch.randn(2, block_size, num_heads, head_dim), + torch.randn(2, block_size, num_heads, head_dim), + ) + else: + layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor") + + manager = OmniKVTransferManager(kv_config) + normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0) + assert normalized is None + + def test_manager_reception(kv_config, mock_connector, common_constants): """Test reception and injection logic in OmniKVTransferManager.""" num_layers = common_constants["num_layers"] diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 82c06fdaf..7c7d90ca1 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -15,6 +15,8 @@ from .utils.config import ConnectorSpec logger = init_logger(__name__) +LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] + @dataclass class OmniKVCacheConfig: @@ -157,7 +159,7 @@ class OmniKVTransferManager: def handle_finished_requests_kv_transfer( self, finished_reqs: dict[str, dict[str, Any]], - kv_caches: list[torch.Tensor], + kv_caches: list[LayerKV], block_size: int, cache_dtype: str, request_id_resolver: Callable[[str], str] | None = None, @@ -169,7 +171,7 @@ class OmniKVTransferManager: Args: finished_reqs: Dict mapping request_id to {block_ids, seq_len} - kv_caches: List of KV cache tensors per layer + kv_caches: List of KV cache (tensor or tuple) per layer block_size: Size of each cache block cache_dtype: Data type of the cache request_id_resolver: Optional function to resolve global request ID @@ -219,7 +221,7 @@ class OmniKVTransferManager: req_id: str, block_ids: list[int], seq_len: int, - kv_caches: list[torch.Tensor], + kv_caches: list[LayerKV], block_size: int, cache_dtype: str, ) -> KVCacheTransferData | None: @@ -229,10 +231,13 @@ class OmniKVTransferManager: req_id: Request identifier block_ids: List of block IDs to extract seq_len: Sequence length - kv_caches: List of KV cache tensors per layer + kv_caches: List of KV cache (tensor or tuple) per layer block_size: Size of each cache block cache_dtype: Data type of the cache + Note: If key/value block counts differ, extraction uses only the overlapping + block range. Extra key/value blocks are ignored, so returned KV may be partial. + Returns: KVCacheTransferData if extraction successful, None otherwise """ @@ -240,25 +245,37 @@ class OmniKVTransferManager: key_cache: list[torch.Tensor | None] = [None] * num_layers value_cache: list[torch.Tensor | None] = [None] * num_layers - for layer_idx, kv_tensor in enumerate(kv_caches): - # Validate block IDs - shape: [2, num_blocks, block_size, n_heads, head_dim] - max_block = kv_tensor.shape[1] - 1 + for layer_idx, layer_kv in enumerate(kv_caches): + kv_pair = self._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx) + if kv_pair is None: + continue + key_blocks, value_blocks = kv_pair + + if key_blocks.shape[0] != value_blocks.shape[0]: + logger.warning( + f"Layer {layer_idx} for request {req_id} has mismatched KV block counts: " + f"key={key_blocks.shape[0]}, value={value_blocks.shape[0]}; using shared range" + ) + + # Validate block IDs - shape: [num_blocks, block_size, n_heads, head_dim] + max_block = min(key_blocks.shape[0], value_blocks.shape[0]) - 1 valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block] if not valid_ids: continue - # Extract and reshape: [2, n_blocks, block_size, n_heads, head_dim] - # -> [2, seq_len, n_heads, head_dim] - selected = kv_tensor[:, valid_ids] # [2, n_valid, block_size, n_heads, head_dim] - n_kv, n_blks, blk_sz, n_heads, d_head = selected.shape - flat = selected.reshape(n_kv, n_blks * blk_sz, n_heads, d_head) - if seq_len < flat.shape[1]: - flat = flat[:, :seq_len] + # Extract and reshape: [n_blocks, block_size, n_heads, head_dim] + # -> [seq_len, n_heads, head_dim] + selected_k = key_blocks[valid_ids] + selected_v = value_blocks[valid_ids] + flat_k = selected_k.flatten(0, 1) + flat_v = selected_v.flatten(0, 1) + if seq_len < flat_k.shape[0]: + flat_k = flat_k[:seq_len] + flat_v = flat_v[:seq_len] # Move to CPU - flat_cpu = flat.detach().cpu().contiguous() - key_cache[layer_idx] = flat_cpu[0] - value_cache[layer_idx] = flat_cpu[1] + key_cache[layer_idx] = flat_k.detach().cpu().contiguous() + value_cache[layer_idx] = flat_v.detach().cpu().contiguous() if not any(k is not None for k in key_cache): return None @@ -275,6 +292,54 @@ class OmniKVTransferManager: }, ) + def _normalize_layer_kv( + self, + layer_kv: LayerKV, + req_id: str, + layer_idx: int, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Normalize one layer KV cache to a `(key_blocks, value_blocks)` tuple. + + Args: + layer_kv: The raw KV cache (tensor or tuple) for the layer + req_id: Request ID for logging + layer_idx: Layer index for logging + + Returns: + Tuple of (key_blocks, value_blocks) if valid, None otherwise + """ + if isinstance(layer_kv, torch.Tensor): + if layer_kv.ndim < 3 or layer_kv.shape[0] != 2: + logger.warning( + f"Layer {layer_idx} for request {req_id} has invalid stacked KV shape: " + f"expected [2, blocks, block_size, ...], got {tuple(layer_kv.shape)}" + ) + return None + key_blocks = layer_kv[0] + value_blocks = layer_kv[1] + elif isinstance(layer_kv, tuple): + if len(layer_kv) != 2: + logger.warning( + f"Layer {layer_idx} for request {req_id} has KV pair length {len(layer_kv)} (expected 2)" + ) + return None + key_blocks, value_blocks = layer_kv + if not isinstance(key_blocks, torch.Tensor) or not isinstance(value_blocks, torch.Tensor): + logger.warning(f"Layer {layer_idx} for request {req_id} has non-tensor KV pair entries") + return None + else: + logger.warning(f"Layer {layer_idx} for request {req_id} has unsupported KV type {type(layer_kv).__name__}") + return None + # ensure key/value blocks are at least 2D for block indexing + if key_blocks.ndim < 2 or value_blocks.ndim < 2: + logger.warning( + f"Layer {layer_idx} for request {req_id} has invalid KV block shape: " + f"got key={tuple(key_blocks.shape)} value={tuple(value_blocks.shape)}" + ) + return None + + return key_blocks, value_blocks + def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None: """Transfer KV cache data to downstream stage via OmniConnector.