[Bugfix] Fix tuple/list KV cache extraction crash (#1405)

Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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"]

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.