[py_connector] restore pre-0.26 vLLM KV layouts via shape-based detection

PR #257's rewrite only accepted vLLM >= 0.26.0's packed 4-D KV layout
(num_blocks, H, block, 2D), regressing full-attention support on older
vLLM. Detect all three flash_attn layouts from the tensor shape itself
(never from version strings) and normalize each into token-major
per-pointer views:

* 4-D packed (vllm >= 0.26.0): one pointer per layer, as before.
* 5-D N-first (num_blocks, 2, block, H, D) (vllm 0.23.0 - 0.25.x): two
  pointers per layer, K/V interleaved per block -> kernel strided path.
* 5-D KV-first (2, num_blocks, block, H, D) (vllm <= 0.22.1): two flat
  pointers per layer.

The Triton gather/scatter kernel already addresses through a flat
pointer array ([K0, V0, K1, V1, ...] for non-MLA), so it needs zero
changes; TransferGroup grows num_kv_ptrs (pointer count, = layer_num
for packed, 2x for split layouts) and the staging buffer views in
data_transfer use it. per_block_bytes is identical across layouts, so
the manager, storage and transfer protocol are unaffected. Unrecognized
layouts still fail fast at startup.

Hybrid (mamba) models on vllm <= 0.22.x are rejected with a clear
NotImplementedError: those schedulers assert
num_external_computed_tokens == 0 in _mamba_block_aligned_split, so the
first external hit would crash mid-flight. The gate probes the
installed scheduler for that blocking assert (capability check, not a
version comparison).

Note the saved byte layout differs between the packed and split-K/V
eras, so KV cache is not portable across vLLM upgrades; instance_id
isolation already prevents such mixing in practice.

New unit tests cover view construction for all three layouts (shape /
stride / pointer math on stub tensors, no torch required), fail-fast on
unrecognized and ambiguous shapes, and the hybrid gate on old/new/
unprobeable schedulers. The e2e VerifyingConnector reuses attn_kv_views
so its captures follow the same normalization.

Author: lpdink

integration_test/vllm_e2e/README.md

@@ -6,7 +6,8 @@ End-to-end integration tests for the KVCM vLLM connector
 through the OpenAI API and verifies that the KV cache data saved to / loaded
 from KVCM is correct.
 
-Requires 1-2 GPUs and vLLM >= 0.26.0.
+Requires 1-2 GPUs and vLLM (0.22.1, 0.23.0 and 0.26.0 are e2e-verified; the
+connector detects the KV cache layout of each era from the tensor shape).
 
 ## What is verified
 
@@ -98,7 +99,7 @@ All environment variables used by the e2e harness:
 | Variable | Required | Meaning |
 |---|---|---|
 | `KVCM_E2E_MODEL` | yes | Path to a local HF model directory (`config.json` + weights). Full-attention coverage needs a plain attention model (e.g. Qwen2.5-7B-Instruct); hybrid coverage needs a mamba/linear + attention model (e.g. Qwen3.5-4B). Hybrid models are auto-detected from `config.json`. |
-| `KVCM_E2E_PYTHON` | yes | Python interpreter of a venv with vLLM >= 0.26.0 and both KVCM wheels (`kvcm_py_client`, `kvcm_vllm_connector`) installed. |
+| `KVCM_E2E_PYTHON` | yes | Python interpreter of a venv with vLLM (any supported version, see above) and both KVCM wheels (`kvcm_py_client`, `kvcm_vllm_connector`) installed. |
 | `KVCM_E2E_CAPTURE_DIR` | internal | Set by the driver for the vLLM subprocess; tells `VerifyingConnector` where to write `.pt` captures. Do not set manually. |
 
 The driver also sets vLLM knobs for the spawned server (`VLLM_KV_CACHE_LAYOUT=NHD`,

integration_test/vllm_e2e/test_connector.py

@@ -64,7 +64,8 @@
 
 from kv_cache_manager.py_connector.common.logger import logger
 from kv_cache_manager.py_connector.vllm.metadata import TairKvCacheConnectorMetadata
-from kv_cache_manager.py_connector.vllm.v1_connector import TairKvCacheConnector
+from kv_cache_manager.py_connector.vllm.v1_connector import (
+    TairKvCacheConnector, attn_kv_views)
 
 CAPTURE_DIR_ENV = "KVCM_E2E_CAPTURE_DIR"
 
@@ -89,7 +90,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         for meta in self._group_metas:
             if meta.is_attention:
                 ref = kv_caches[meta.layer_names[0]]
-                kernel_bs = ref.shape[2]
+                kernel_bs = attn_kv_views(ref)[0].shape[1]
             else:
                 kernel_bs = 0
             self._cap_groups.append(
@@ -259,13 +260,20 @@ def _capture_block(self, kind, token_ids, block_ids_per_group, manager_block_idx
                 slot_tensor = torch.tensor(slots, dtype=torch.long, device=self._device)
                 for layer_name in layer_names:
                     kv_cache = self._kv_caches[layer_name]
-                    # vLLM >= 0.26.0: (num_blocks, num_kv_heads, kernel_bs,
-                    # 2*head_size) packed, NHD memory order is token-major. Flatten
-                    # the (block, token) dims and gather the whole per-token vector
-                    # (K and V packed) -- the packing is opaque to verification.
-                    per_token = kv_cache.shape[1] * kv_cache.shape[3]
-                    flat = kv_cache.permute(0, 2, 1, 3).reshape(-1, per_token)
-                    gathered = flat[slot_tensor, :].contiguous()  # [n_tok, per_token]
+                    # Normalize the layout (packed 4-D or split K/V 5-D) into
+                    # token-major views via the production helper and gather the
+                    # whole per-token vector by (block, token) advanced indexing
+                    # -- split K/V views are non-contiguous, so flattening them
+                    # first would copy the entire cache tensor. Split views are
+                    # concatenated on the content dim, so a capture is
+                    # comparable across save/load within one run.
+                    parts = []
+                    for v in attn_kv_views(kv_cache):
+                        blk = slot_tensor // kernel_bs
+                        tok = slot_tensor % kernel_bs
+                        parts.append(v[blk, tok].reshape(len(slots), -1))
+                    gathered = (parts[0] if len(parts) == 1
+                                else torch.cat(parts, dim=-1)).contiguous()
                     kv_by_layer[layer_name] = gathered.cpu()
             else:
                 # State stored once per group block; the manager block's last

kv_cache_manager/py_connector/common/types.py

@@ -30,12 +30,19 @@ class TransferGroup:
     layer_num: int = 0
 
     # --- Attention-only fields (is_attention == True) ---
-    # int64 tensor of [K0, V0, K1, V1, ...] data ptrs on the compute device.
+    # int64 tensor of transfer-pointer bases on the compute device. For packed
+    # K/V layouts (vllm >= 0.26.0) one pointer per layer [L0, L1, ...]; for
+    # split K/V layouts (vllm <= 0.25.x) two per layer [K0, V0, K1, V1, ...].
     kvcache_ptr_tensor_gpu: Optional[torch.Tensor] = None
-    per_token_dim: int = 0          # num_kv_heads * head_size
-    kernel_block_size: int = 0      # tensor.shape[2]
-    kv_stride: int = 0              # tensor.stride(0), 0 => contiguous flat layout
-    block_stride: int = 0           # tensor.stride(1), 0 => contiguous flat layout
+    # Number of transfer pointers (rows of the staging buffer view per block).
+    # layer_num for packed layouts, 2 * layer_num for split K/V layouts.
+    num_kv_ptrs: int = 0
+    per_token_dim: int = 0          # heads * content dim per pointer
+    kernel_block_size: int = 0      # tokens per kernel page
+    kv_stride: int = 0              # K->V element offset added by the kernel; 0
+                                    # because every pointer is its own base
+    block_stride: int = 0           # element stride between kernel pages of one
+                                    # pointer, 0 => contiguous flat layout
 
     # --- State-only fields (is_attention == False) ---
     # Per layer (num_blocks, page_size_bytes) uint8 views into the state storage.

kv_cache_manager/py_connector/test/BUILD

@@ -28,6 +28,13 @@ py_test(
     deps = [":vllm_stubs"],
 )
 
+py_test(
+    name = "test_kv_layouts",
+    srcs = ["test_kv_layouts.py"],
+    tags = ["no-remote-exec"],
+    deps = [":vllm_stubs"],
+)
+
 py_test(
     name = "test_scheduler_state",
     srcs = ["test_scheduler_state.py"],

kv_cache_manager/py_connector/test/test_kv_layouts.py

@@ -0,0 +1,263 @@
+"""Unit tests for the multi-version KV cache layout detection.
+
+``attn_kv_views`` must recognize the three flash_attn layouts vLLM has shipped
+(detected from the tensor shape, never from version strings) and reject
+anything else:
+
+* 4-D packed   ``(num_blocks, H, block, 2D)``   -- vLLM >= 0.26.0
+* 5-D N-first  ``(num_blocks, 2, block, H, D)`` -- vLLM 0.23.0 - 0.25.x
+* 5-D KV-first ``(2, num_blocks, block, H, D)`` -- vLLM <= 0.22.1
+
+``_build_transfer_group`` must derive the transfer pointers / strides from the
+normalized views, and ``ensure_hybrid_supported`` must fail fast when the
+installed vLLM's scheduler rejects external KV loads for hybrid models
+(vLLM <= 0.22.x).
+
+Runs without torch: a minimal FakeTensor models the strided-view semantics
+(shape / stride / offset / data_ptr) that the code under test reads.
+"""
+
+import sys
+import types
+import unittest
+
+from kv_cache_manager.py_connector.test.vllm_stubs import make_connector
+from kv_cache_manager.py_connector.vllm.v1_connector import (
+    attn_kv_views, ensure_hybrid_supported, GroupMeta)
+
+ITEMSIZE = 2  # bf16/fp16
+BASE_PTR = 1 << 20
+
+
+class FakeTensor:
+    """Minimal strided tensor: only what attn_kv_views / _build_transfer_group
+    read (dim/shape/stride/permute/indexing/data_ptr)."""
+
+    def __init__(self, shape, strides, offset=0, base=BASE_PTR):
+        self.shape = tuple(shape)
+        self._strides = tuple(strides)
+        self._offset = offset
+        self._base = base
+
+    @classmethod
+    def contiguous(cls, shape, base=BASE_PTR):
+        strides, acc = [], 1
+        for s in reversed(shape):
+            strides.append(acc)
+            acc *= s
+        return cls(shape, tuple(reversed(strides)), base=base)
+
+    def dim(self):
+        return len(self.shape)
+
+    def stride(self, i=None):
+        return self._strides if i is None else self._strides[i]
+
+    def data_ptr(self):
+        return self._base + self._offset * ITEMSIZE
+
+    def permute(self, *dims):
+        return FakeTensor([self.shape[d] for d in dims],
+                          [self._strides[d] for d in dims],
+                          self._offset, self._base)
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):  # t[i]: drop dim 0
+            return FakeTensor(self.shape[1:], self._strides[1:],
+                              self._offset + idx * self._strides[0], self._base)
+        if isinstance(idx, tuple) and idx[0] == slice(None) and isinstance(idx[1], int):
+            # t[:, i]: drop dim 1
+            return FakeTensor(self.shape[:1] + self.shape[2:],
+                              self._strides[:1] + self._strides[2:],
+                              self._offset + idx[1] * self._strides[1], self._base)
+        raise TypeError(f"unsupported index {idx!r}")
+
+
+def packed_4d(n=10, h=4, b=16, d2=256, base=BASE_PTR):
+    """vLLM >= 0.26.0: NHD memory is (n, b, h, d2) contiguous; the registered
+    tensor is its (n, h, b, d2) permuted view."""
+    return FakeTensor.contiguous([n, b, h, d2], base=base).permute(0, 2, 1, 3)
+
+
+def kv_first_5d(n=10, b=16, h=4, d=128, base=BASE_PTR):
+    """vLLM <= 0.22.1: (2, n, b, h, d) contiguous."""
+    return FakeTensor.contiguous([2, n, b, h, d], base=base)
+
+
+def n_first_5d(n=10, b=16, h=4, d=128, base=BASE_PTR):
+    """vLLM 0.23.0 - 0.25.x: (n, 2, b, h, d) contiguous."""
+    return FakeTensor.contiguous([n, 2, b, h, d], base=base)
+
+
+class TestAttnKvViews(unittest.TestCase):
+    def test_packed_4d(self):
+        views = attn_kv_views(packed_4d())
+        self.assertEqual(len(views), 1)
+        v = views[0]
+        self.assertEqual(v.shape, (10, 16, 4, 256))          # (n, b, h, 2d)
+        self.assertEqual(v.stride(), (16 * 4 * 256, 4 * 256, 256, 1))
+        self.assertEqual(v.data_ptr(), BASE_PTR)             # storage base
+
+    def test_kv_first_5d(self):
+        views = attn_kv_views(kv_first_5d())
+        self.assertEqual(len(views), 2)
+        k, v = views
+        for view in (k, v):
+            self.assertEqual(view.shape, (10, 16, 4, 128))
+            self.assertEqual(view.stride(), (16 * 4 * 128, 4 * 128, 128, 1))
+        self.assertEqual(k.data_ptr(), BASE_PTR)
+        # V base = K base + num_blocks * block * h * d elements.
+        self.assertEqual(v.data_ptr() - k.data_ptr(),
+                         10 * 16 * 4 * 128 * ITEMSIZE)
+
+    def test_n_first_5d(self):
+        views = attn_kv_views(n_first_5d())
+        self.assertEqual(len(views), 2)
+        k, v = views
+        for view in (k, v):
+            self.assertEqual(view.shape, (10, 16, 4, 128))
+            # K and V of one block are interleaved: the block stride covers
+            # both halves while the inner page stays token-major.
+            self.assertEqual(view.stride(), (2 * 16 * 4 * 128, 4 * 128, 128, 1))
+        self.assertEqual(v.data_ptr() - k.data_ptr(),
+                         16 * 4 * 128 * ITEMSIZE)
+
+    def test_unrecognized_layouts_fail_fast(self):
+        bad = [
+            FakeTensor.contiguous([10, 16, 4]),           # 3-D
+            FakeTensor.contiguous([10, 2, 16, 4, 128, 2]),  # 6-D
+            FakeTensor.contiguous([10, 16, 2, 4, 128]),   # 5-D, K/V dim misplaced
+        ]
+        for t in bad:
+            with self.subTest(shape=t.shape):
+                with self.assertRaises(NotImplementedError):
+                    attn_kv_views(t)
+
+    def test_ambiguous_layout_fails_fast(self):
+        # num_blocks == 2 in a KV-first shape is indistinguishable from a
+        # two-block N-first shape; refusing beats guessing.
+        with self.assertRaises(NotImplementedError):
+            attn_kv_views(FakeTensor.contiguous([2, 2, 16, 4, 128]))
+
+
+def _make_group_conn():
+    conn = make_connector(manager_block_size=16)
+    conn._self_spec_names = ["tp0_g0"]
+    conn._device = "cpu"
+    return conn
+
+
+def _attn_meta(layer_names, block_size=16):
+    return GroupMeta(group_idx=0, is_attention=True, layer_names=layer_names,
+                     block_size=block_size, per_block_bytes=0)
+
+
+class TestBuildTransferGroup(unittest.TestCase):
+    """Pointer construction per layout. Layer tensors get distinct bases so the
+    interleaving [K0, V0, K1, V1, ...] is observable. The pointer list is
+    captured by patching ``torch.tensor`` (works with both the stubbed and a
+    real torch: no tensor math happens on the captured value)."""
+
+    def _build(self, kv_caches):
+        import unittest.mock as mock
+        import kv_cache_manager.py_connector.vllm.v1_connector as v1c
+        conn = _make_group_conn()
+        captured = []
+
+        def fake_tensor(data, **kw):
+            captured[:] = list(data)
+            t = mock.MagicMock()
+            t.to.return_value = t
+            return t
+
+        with mock.patch.object(v1c.torch, "tensor", side_effect=fake_tensor):
+            g = conn._build_transfer_group(
+                _attn_meta(list(kv_caches.keys())), kv_caches)
+        return g, captured
+
+    def test_packed_one_ptr_per_layer(self):
+        kv = {"l0": packed_4d(base=BASE_PTR), "l1": packed_4d(base=2 * BASE_PTR)}
+        g, ptrs = self._build(kv)
+        self.assertEqual(g.num_kv_ptrs, 2)
+        self.assertEqual(g.layer_num, 2)
+        self.assertEqual(g.per_token_dim, 4 * 256)
+        self.assertEqual(g.kernel_block_size, 16)
+        self.assertEqual(g.block_stride, 0)  # flat
+        self.assertEqual(ptrs, [BASE_PTR, 2 * BASE_PTR])
+
+    def test_kv_first_two_ptrs_per_layer(self):
+        kv = {"l0": kv_first_5d(base=BASE_PTR), "l1": kv_first_5d(base=2 * BASE_PTR)}
+        g, ptrs = self._build(kv)
+        self.assertEqual(g.num_kv_ptrs, 4)
+        self.assertEqual(g.layer_num, 2)
+        self.assertEqual(g.per_token_dim, 4 * 128)
+        self.assertEqual(g.block_stride, 0)  # each half is flat token-major
+        v_off = 10 * 16 * 4 * 128 * ITEMSIZE
+        self.assertEqual(ptrs, [BASE_PTR, BASE_PTR + v_off,
+                                2 * BASE_PTR, 2 * BASE_PTR + v_off])
+
+    def test_n_first_strided_blocks(self):
+        kv = {"l0": n_first_5d(base=BASE_PTR)}
+        g, ptrs = self._build(kv)
+        self.assertEqual(g.num_kv_ptrs, 2)
+        self.assertEqual(g.per_token_dim, 4 * 128)
+        # K/V interleaved per block -> kernel must walk the strided path.
+        self.assertEqual(g.block_stride, 2 * 16 * 4 * 128)
+        v_off = 16 * 4 * 128 * ITEMSIZE
+        self.assertEqual(ptrs, [BASE_PTR, BASE_PTR + v_off])
+
+    def test_unrecognized_layout_fails_fast(self):
+        with self.assertRaises(NotImplementedError):
+            self._build({"l0": FakeTensor.contiguous([10, 16, 4])})
+
+
+class _BlockedScheduler:
+    """Mimics vLLM <= 0.22.x: external loads are asserted away."""
+
+    def _mamba_block_aligned_split(self, request, num_new_tokens,
+                                   num_new_local_computed_tokens=0,
+                                   num_external_computed_tokens=0):
+        assert num_external_computed_tokens == 0, (
+            "External KV connector is not verified yet"
+        )
+
+
+class _OpenScheduler:
+    """Mimics vLLM >= 0.23.0: the split handles external tokens."""
+
+    def _mamba_block_aligned_split(self, request, num_new_tokens,
+                                   num_new_local_computed_tokens=0,
+                                   num_external_computed_tokens=0):
+        return num_new_tokens
+
+
+class TestHybridGate(unittest.TestCase):
+    MOD = "vllm.v1.core.sched.scheduler"
+
+    def _with_scheduler(self, cls):
+        mod = types.ModuleType(self.MOD)
+        mod.Scheduler = cls
+        old = sys.modules.get(self.MOD)
+        sys.modules[self.MOD] = mod
+        self.addCleanup(lambda: (sys.modules.pop(self.MOD, None),
+                                 old and sys.modules.__setitem__(self.MOD, old)))
+
+    def test_old_vllm_hybrid_raises_gracefully(self):
+        self._with_scheduler(_BlockedScheduler)
+        with self.assertRaises(NotImplementedError) as ctx:
+            ensure_hybrid_supported()
+        # The message must tell the operator what to do.
+        self.assertIn("vLLM >= 0.23.0", str(ctx.exception))
+        self.assertIn("hybrid", str(ctx.exception))
+
+    def test_new_vllm_hybrid_passes(self):
+        self._with_scheduler(_OpenScheduler)
+        ensure_hybrid_supported()  # must not raise
+
+    def test_unprobeable_scheduler_does_not_block(self):
+        self._with_scheduler(object)  # no _mamba_block_aligned_split at all
+        ensure_hybrid_supported()  # must not raise
+
+
+if __name__ == "__main__":
+    unittest.main()

kv_cache_manager/py_connector/vllm/data_transfer.py

@@ -147,7 +147,7 @@ def save_task(self, multi_result: MultiResult, task_idx, group: TransferGroup,
                                          dtype=torch.uint8, device=self._device)
                 if group.is_attention:
                     view = gpu_buffer.view(self._info.dtype).view(
-                        len(valid), group.layer_num,
+                        len(valid), group.num_kv_ptrs,
                         self._manager_block_size, group.per_token_dim)
                     batch_gather_scatter_helper.batch_gather_kv_caches(
                         group.kvcache_ptr_tensor_gpu, view, block_token_indices,
@@ -226,7 +226,7 @@ def load_task(self, multi_result: MultiResult, task_idx, group: TransferGroup,
                 gpu_buffer = cpu_buffer.to(self._device, non_blocking=True)
                 if group.is_attention:
                     view = gpu_buffer.view(self._info.dtype).view(
-                        n, group.layer_num, self._manager_block_size, group.per_token_dim)
+                        n, group.num_kv_ptrs, self._manager_block_size, group.per_token_dim)
                     batch_gather_scatter_helper.batch_scatter_kv_caches(
                         group.kvcache_ptr_tensor_gpu, view, block_token_indices,
                         list(range(n)), self._manager_block_size, group.per_token_dim,