fix(vt): reshape_and_cache indexes paged KV cache by tensor strides

Fixes the stride-handling defect in e231196 (M1.6 Task 2). ReshapeAndCache
derived block/page strides from k_cache.shape and required whole-cache
IsContiguous(), which is wrong for our committed layout: get_kv_cache_shape
returns one (num_blocks, 2, block_size, H, D) allocation and K/V are its two
dim-1 unbind slices — rank-4 STRIDED views with block stride 2*bs*H*D (not
bs*H*D), never contiguous. Feeding the real slices either threw on the guard
or (with a shape-derived block_stride = half the real stride) silently wrote
block b into interleaved K/V memory, clobbering the other slice.

Mirror pinned csrc/libtorch_stable/cache_kernels.cu @ e24d1b24: source
block/page strides from key_cache.stride(0/1) and the token stride from
key.stride(0) (host ~L797-801; kernel ~L337-347). Each cache slice is indexed
with ITS OWN strides.

- cpu_cache.cpp / cuda_cache.cu: dst = block*stride[0] + offset*stride[1],
  src = token*stride[0]; per-token page stays one dense memcpy/run given the
  head-contiguous NHD slice (pinned is_contiguous_heads fast path).
- ops.cpp: relax the guard — no cache contiguity; require only elem stride 1,
  head-contiguous page (stride[2]==head_size), and contiguous k/v/slot_mapping.
- tests: add strided-unbind-slice CPU tests (drive the exact missed case;
  failed/threw before, pass after) + a build-guarded CUDA strided-parity test
  (dgx-pending). Existing contiguous tests still pass (contiguous cache is the
  stride[0]==bs*H*D special case).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: mudler

docs/superpowers/plans/2026-07-03-m1.6-paged-attention.md

@@ -21,10 +21,13 @@
 Read `/home/mudler/_git/vllm/vllm/v1/attention/backend.py` (VERIFY current API). Port `CommonAttentionMetadata` (the exact T0 field set), `AttentionBackend` (ABC: get_impl_cls/get_builder_cls/get_kv_cache_shape), `AttentionImpl` (`forward(layer, q, k, v, kv_cache, attn_metadata, output, ...)`), `AttentionMetadataBuilder` (`build(common_attn_metadata) -> backend metadata`). Behavioral interface + a metadata builder that turns the M1.5 step-inputs (query_start_loc/seq_lens/slot_mapping/block_table) into the backend metadata. Unit tests: build CommonAttentionMetadata from step-inputs, the max_query_len/max_seq_len derivation, get_kv_cache_shape.
 
 ### Task 2: reshape_and_cache + paged KV cache layout
-Read the pinned reshape_and_cache (csrc/cache_kernels or the CPU backend's cache write) + the KV cache tensor layout (`get_kv_cache_shape` for the full-attn spec). Port `vt::reshape_and_cache(k, v, k_cache, v_cache, slot_mapping)` — write new K/V into the paged cache at the slot ids. CPU + CUDA. Match the exact paged layout (block-major [num_blocks, block_size, num_kv_heads, head_size] or the upstream shape). Pinned-oracle golden. Unit tests: writing tokens to slots, reading them back at the right block/offset; CUDA vs CPU.
+Port `vt::reshape_and_cache(k, v, k_cache, v_cache, slot_mapping)` — write new K/V into the paged cache at the slot ids. CPU + CUDA.
+**LAYOUT (BINDING, from Task 1 + review): Task 1 committed the flash NHD shape `get_kv_cache_shape = (num_blocks, 2, block_size, num_kv_heads, head_size)`.** So `k_cache`/`v_cache` are the two dim-1 slices, each `[num_blocks, block_size, num_kv_heads, head_size]`; a slot id maps to `block = slot / block_size, offset = slot % block_size`, and the write goes to `[block, offset, kv_head, :]`. **TRAP TO AVOID:** the pinned `cpu_attn.py` uses a DIFFERENT (HND) internal layout `(num_blocks, num_kv_heads, block_size, 2*head_size)` — DO NOT port cpu_attn's cache-indexing/view arithmetic. Take cpu_attn (or flash_attn) only for the write SEMANTICS; index against the NHD shape Task 1 allocates. Crossing the two layouts silently corrupts every output.
+**Golden strategy (review): compose the reference from math, not backend cache bytes** — a reshape_and_cache "golden" is just: after writing, reading slot s back from the NHD cache yields the input k/v for that token. Unit-test WRITE→READ round-trip directly (host, layout-consistent); no external oracle needed. CUDA vs CPU parity for the kernel.
 
 ### Task 3: Paged attention op (varlen prefill + paged decode)
-Port the correctness-grade paged attention: `vt::paged_attention(out, q, k_cache, v_cache, block_table, seq_lens, query_start_loc, scale, ...)` — for each query token, causal GQA softmax over the K/V read from the paged blocks (block_table → block ids → cache slots) up to seq_len. This generalizes M0.9's dense `vt::Attention` to the paged/varlen/batched case. CPU reference + CUDA (correctness-grade block-per-(query,head), the FlashInfer perf kernel is M2.4). Pinned-oracle golden (dump from the pinned CPU attention backend for a small batched case). Validate against M0.9's dense attention on the single-sequence case (must agree). Unit tests + parity golden CPU+CUDA on dgx.
+Port the correctness-grade paged attention: `vt::paged_attention(out, q, k_cache, v_cache, block_table, seq_lens, query_start_loc, scale, ...)` — for each query token, causal GQA softmax over the K/V read from the paged blocks (`block_table` → block ids; a token at absolute position p reads block `block_table[req, p/block_size]`, offset `p%block_size` in the **NHD** cache from Task 2) up to seq_len. This generalizes M0.9's dense `vt::Attention` to the paged/varlen/batched case. CPU reference + CUDA (correctness-grade block-per-(query,head); FlashInfer perf kernel is M2.4). **READ against the NHD layout Task 2 writes (same trap: not cpu_attn's HND arithmetic).**
+**Golden strategy (review): COMPOSE the reference math (M0.9-style), do NOT dump backend cache bytes.** The attention OUTPUT is layout-agnostic: build the golden as per-token causal GQA softmax over gathered K/V (the same reference approach as M0.9's `dense_attention` golden). **Anchor: on the single-sequence contiguous case, paged_attention MUST agree bit-for-tolerance with M0.9's dense `vt::Attention`** — assert this directly (it's the strongest correctness check + needs no new oracle). Then a small batched varlen case (2 reqs: prefill + decode) via composed reference. Unit tests + CUDA-vs-CPU parity on dgx.
 
 ### Task 4: GDN attention metadata (prefill/decode segmentation)
 Read `/home/mudler/_git/vllm/vllm/v1/attention/backends/gdn_attn.py::GDNAttentionMetadata`. Port the metadata that segments a batched step into GDN prefill (chunked-scan) vs decode (recurrence) vs spec segments — num_prefills/num_decodes, the has_initial_state masks, the segment offsets — so the M0.7 GDN ops (GdnPrefill/GdnDecode) can be driven by a batched SchedulerOutput. This is the hybrid-model glue: the metadata builder splits the batch's GDN-layer work. Behavioral (no new GDN kernels — reuse M0.7). Unit tests: a batched step with 1 prefill + 1 decode request produces the right GDN segmentation.

src/vt/cpu/cpu_cache.cpp

@@ -15,35 +15,59 @@ namespace {
 // Write each new per-token K/V into the paged NHD cache at its slot id. The
 // "auto" cache path is a raw element copy (cache dtype == k/v dtype); we copy
 // bytes so f32/f16/bf16 are all bit-exact (upstream KV_T == CACHE_T in auto).
+//
+// CRITICAL (M1.6 Task-2): the destination is indexed from the TENSOR STRIDES,
+// not from k_cache.shape. get_kv_cache_shape hands us ONE (num_blocks, 2,
+// block_size, H, D) allocation and K/V are its two dim-1 unbind slices, so the
+// block stride is 2*bs*H*D (NOT bs*H*D) and K/V are NON-contiguous rank-4 views.
+// This mirrors pinned csrc/libtorch_stable/cache_kernels.cu @ e24d1b24:
+//   host reshape_and_cache_flash (~L797-801): key_stride = key.stride(0);
+//     block_stride = key_cache.stride(0); page_stride = key_cache.stride(1);
+//     head_stride = key_cache.stride(2);
+//   kernel reshape_and_cache_flash_kernel (~L337-347): key_src = key +
+//     token_idx*key_stride; key_dst = key_cache + block_idx*block_stride +
+//     block_offset*page_stride; is_contiguous_heads = (head_stride == head_size).
+// The wrapper guarantees head_stride == head_size && elem stride == 1 (the NHD
+// unbind slice), so the per-token page is one dense run of n_elems inside the
+// block — the is_contiguous_heads fast path, i.e. a single memcpy per token.
 void ReshapeAndCacheKernel(Queue&, const Tensor& k, const Tensor& v, Tensor& k_cache,
                            Tensor& v_cache, const Tensor& slot_mapping) {
   const int64_t num_slots = slot_mapping.shape[0];
   const int64_t block_size = k_cache.shape[1];
   const int64_t num_kv_heads = k_cache.shape[2];
   const int64_t head_size = k_cache.shape[3];
   const int64_t n_elems = num_kv_heads * head_size;  // one token's page (NHD)
-  const int64_t page_stride = n_elems;               // stride over block_size (dim 1)
-  const int64_t block_stride = block_size * n_elems; // stride over num_blocks  (dim 0)
+  // Destination strides come from the tensors (unbind-slice aware), each cache
+  // with ITS OWN strides. Source token stride comes from k/v.stride(0); the
+  // per-token [H, D] payload is packed (input k/v are contiguous rows).
+  const int64_t k_block_stride = k_cache.stride[0];
+  const int64_t k_page_stride = k_cache.stride[1];
+  const int64_t v_block_stride = v_cache.stride[0];
+  const int64_t v_page_stride = v_cache.stride[1];
+  const int64_t k_tok_stride = k.stride[0];
+  const int64_t v_tok_stride = v.stride[0];
   const size_t elem = SizeOf(k.dtype);
 
   const int64_t* slots = slot_mapping.Ptr<int64_t>();
   const auto* ksrc = static_cast<const uint8_t*>(k.data);
   const auto* vsrc = static_cast<const uint8_t*>(v.data);
   auto* kdst = static_cast<uint8_t*>(k_cache.data);
   auto* vdst = static_cast<uint8_t*>(v_cache.data);
+  const size_t bytes = static_cast<size_t>(n_elems) * elem;
 
   for (int64_t t = 0; t < num_slots; ++t) {
     const int64_t slot = slots[t];
     if (slot < 0) continue;  // padded token → skip (upstream NOTE: slot can be -1)
     const int64_t block = slot / block_size;
     const int64_t offset = slot % block_size;
-    const int64_t dst = block * block_stride + offset * page_stride;  // element offset
-    const int64_t src = t * n_elems;
-    const size_t bytes = static_cast<size_t>(n_elems) * elem;
-    std::memcpy(kdst + static_cast<size_t>(dst) * elem, ksrc + static_cast<size_t>(src) * elem,
-                bytes);
-    std::memcpy(vdst + static_cast<size_t>(dst) * elem, vsrc + static_cast<size_t>(src) * elem,
-                bytes);
+    const int64_t kdst_off = block * k_block_stride + offset * k_page_stride;  // elements
+    const int64_t vdst_off = block * v_block_stride + offset * v_page_stride;
+    const int64_t ksrc_off = t * k_tok_stride;
+    const int64_t vsrc_off = t * v_tok_stride;
+    std::memcpy(kdst + static_cast<size_t>(kdst_off) * elem,
+                ksrc + static_cast<size_t>(ksrc_off) * elem, bytes);
+    std::memcpy(vdst + static_cast<size_t>(vdst_off) * elem,
+                vsrc + static_cast<size_t>(vsrc_off) * elem, bytes);
   }
 }
 

src/vt/cuda/cuda_cache.cu

@@ -27,22 +27,36 @@ cudaStream_t AsStream(const Queue& q) { return static_cast<cudaStream_t>(q.handl
 
 // Word is the raw storage type (uint32_t for f32, uint16_t for f16/bf16): the
 // auto cache path is a bit-exact copy, so no dtype-aware conversion is needed.
+//
+// CRITICAL (M1.6 Task-2): destination is indexed from TENSOR STRIDES, not from
+// k_cache.shape. K/V are the two dim-1 unbind slices of one (num_blocks, 2,
+// block_size, H, D) allocation, so their block stride is 2*bs*H*D (NOT bs*H*D)
+// and they are NON-contiguous rank-4 views. Mirrors pinned cache_kernels.cu @
+// e24d1b24 (host ~L797-801 sources key_stride/block_stride/page_stride/
+// head_stride from the tensor strides; kernel ~L337-347 does key_dst =
+// key_cache + block_idx*block_stride + block_offset*page_stride). The wrapper
+// guarantees head_stride == head_size && elem stride == 1 (the NHD unbind
+// slice), so the per-token page is one dense run of n_elems inside the block —
+// pinned's is_contiguous_heads fast path; here the threads stride that run.
 template <typename Word>
-__global__ void ReshapeAndCacheKernel(const Word* __restrict__ key,
-                                      const Word* __restrict__ value, Word* __restrict__ key_cache,
-                                      Word* __restrict__ value_cache,
-                                      const int64_t* __restrict__ slot_mapping, int64_t block_size,
-                                      int64_t n_elems) {
+__global__ void ReshapeAndCacheKernel(
+    const Word* __restrict__ key, const Word* __restrict__ value,
+    Word* __restrict__ key_cache, Word* __restrict__ value_cache,
+    const int64_t* __restrict__ slot_mapping, int64_t block_size, int64_t n_elems,
+    int64_t k_block_stride, int64_t k_page_stride, int64_t v_block_stride,
+    int64_t v_page_stride, int64_t k_tok_stride, int64_t v_tok_stride) {
   const int64_t token = blockIdx.x;
   const int64_t slot = slot_mapping[token];
   if (slot < 0) return;  // padded token → skip
   const int64_t block = slot / block_size;
   const int64_t offset = slot % block_size;
-  const int64_t dst = (block * block_size + offset) * n_elems;  // NHD element offset
-  const int64_t src = token * n_elems;
+  const int64_t kdst = block * k_block_stride + offset * k_page_stride;  // element offset
+  const int64_t vdst = block * v_block_stride + offset * v_page_stride;
+  const int64_t ksrc = token * k_tok_stride;
+  const int64_t vsrc = token * v_tok_stride;
   for (int64_t e = threadIdx.x; e < n_elems; e += blockDim.x) {
-    key_cache[dst + e] = key[src + e];
-    value_cache[dst + e] = value[src + e];
+    key_cache[kdst + e] = key[ksrc + e];
+    value_cache[vdst + e] = value[vsrc + e];
   }
 }
 
@@ -52,6 +66,12 @@ void ReshapeAndCacheKernelCuda(Queue& q, const Tensor& k, const Tensor& v, Tenso
   const int64_t block_size = k_cache.shape[1];
   const int64_t n_elems = k_cache.shape[2] * k_cache.shape[3];
   if (num_slots == 0 || n_elems == 0) return;
+  const int64_t k_block_stride = k_cache.stride[0];
+  const int64_t k_page_stride = k_cache.stride[1];
+  const int64_t v_block_stride = v_cache.stride[0];
+  const int64_t v_page_stride = v_cache.stride[1];
+  const int64_t k_tok_stride = k.stride[0];
+  const int64_t v_tok_stride = v.stride[0];
   const unsigned grid = static_cast<unsigned>(num_slots);
   const unsigned block = static_cast<unsigned>(n_elems < 512 ? n_elems : 512);
   const cudaStream_t s = AsStream(q);
@@ -60,12 +80,14 @@ void ReshapeAndCacheKernelCuda(Queue& q, const Tensor& k, const Tensor& v, Tenso
     case 4:
       ReshapeAndCacheKernel<uint32_t><<<grid, block, 0, s>>>(
           k.Ptr<uint32_t>(), v.Ptr<uint32_t>(), k_cache.Ptr<uint32_t>(), v_cache.Ptr<uint32_t>(),
-          slots, block_size, n_elems);
+          slots, block_size, n_elems, k_block_stride, k_page_stride, v_block_stride,
+          v_page_stride, k_tok_stride, v_tok_stride);
       break;
     case 2:
       ReshapeAndCacheKernel<uint16_t><<<grid, block, 0, s>>>(
           k.Ptr<uint16_t>(), v.Ptr<uint16_t>(), k_cache.Ptr<uint16_t>(), v_cache.Ptr<uint16_t>(),
-          slots, block_size, n_elems);
+          slots, block_size, n_elems, k_block_stride, k_page_stride, v_block_stride,
+          v_page_stride, k_tok_stride, v_tok_stride);
       break;
     default: VT_CHECK(false, "cuda reshape_and_cache: unsupported dtype element size");
   }

src/vt/ops.cpp

@@ -367,9 +367,24 @@ void ReshapeAndCache(Queue& q, const Tensor& k, const Tensor& v, Tensor& k_cache
                v_cache.dtype == k.dtype,
            "reshape_and_cache: k/v/k_cache/v_cache must share one float dtype (auto cache path)");
   VT_CHECK(slot_mapping.dtype == DType::kI64, "reshape_and_cache: slot_mapping must be i64");
-  VT_CHECK(k.IsContiguous() && v.IsContiguous() && k_cache.IsContiguous() &&
-               v_cache.IsContiguous() && slot_mapping.IsContiguous(),
-           "reshape_and_cache: contiguous tensors required");
+  // The paged KV cache is ONE (num_blocks, 2, block_size, H, D) allocation;
+  // k_cache/v_cache are the two dim-1 unbind slices, i.e. rank-4 STRIDED views
+  // (block stride 2*bs*H*D, not bs*H*D). We therefore must NOT require the cache
+  // to be contiguous — indexing is driven by k_cache/v_cache strides (mirroring
+  // pinned cache_kernels.cu::reshape_and_cache_flash, which reads block/page/
+  // head strides from key_cache.stride(0/1/2)). We only require what the copy
+  // actually needs: the innermost element access is well-defined (elem stride 1)
+  // and the per-token page is dense (head stride == head_size, i.e. dim-2/3
+  // packed), which holds for the NHD unbind slice. The input k/v rows and
+  // slot_mapping must be contiguous (upstream reads k/v inner packed, applying
+  // only key.stride(0) for the token, and indexes slot_mapping directly).
+  VT_CHECK(k.IsContiguous() && v.IsContiguous() && slot_mapping.IsContiguous(),
+           "reshape_and_cache: k/v inputs and slot_mapping must be contiguous");
+  VT_CHECK(k_cache.stride[3] == 1 && v_cache.stride[3] == 1,
+           "reshape_and_cache: k_cache/v_cache innermost (head_size) stride must be 1");
+  VT_CHECK(k_cache.stride[2] == head_size && v_cache.stride[2] == head_size,
+           "reshape_and_cache: k_cache/v_cache page must be head-contiguous "
+           "(stride[2] == head_size) — the NHD unbind-slice layout");
   VT_CHECK(k.device == q.device && v.device == q.device && k_cache.device == q.device &&
                v_cache.device == q.device && slot_mapping.device == q.device,
            "reshape_and_cache: device mismatch (k/v/k_cache/v_cache/slot_mapping/queue)");