refactor: stride-based attn interface with layout and causal mask

- Replace is_causal + causal_offset with unified causal_offset (-1 = off, >=0 = first Q pos)
- Causal and mask can now coexist (was mutually exclusive)
- Add stride-based addressing for Q/KV/O (layout-agnostic, zero-copy)
- Add layout param ("bhld"/"blhd") parsed in Python, passed as int to C++
- Support 2D [batch, kv_len] and 3D [batch, q_len, kv_len] mask
- Vectorize paged KV gather in Python fallback (was per-token Python loop)
- Extract shared helpers: compute_num_splits, alloc_split_partials, dispatch_head_dim
- Unify paged_decode entry via attn_pack_paged_params
- Update mma_softmax_tile for 3D mask with per-row qrow indexing

astrai/extension/__init__.py

@@ -5,6 +5,14 @@ Public API:
     - ``attn_prefill`` — multi-query prefill attention
     - ``attn_paged_decode`` — paged decode attention (direct page-table access)
 
+Interface (shared by all wrappers):
+    causal_offset: -1 = non-causal; >=0 = absolute position of first Q token
+    mask:       2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True = keep)
+    scale:      0.0 = auto (1/sqrt(head_dim)); >0 = explicit
+    layout:     "bhld" (default) or "blhd"
+
+Causal and mask can coexist — both are applied simultaneously.
+
 Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``)
 when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
 """

astrai/extension/ops.py

@@ -3,15 +3,43 @@
 Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when
 available, otherwise falls back to ``torch`` SDPA.
 
+Interface (all functions):
+    causal_offset: -1 = non-causal; >=0 = absolute position of first Q token
+    mask:       2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool)
+    scale:      0.0 = auto (1/sqrt(head_dim)); >0 = explicit
+    layout:     "bhld" (default) or "blhd"
+
 Add new kernel wrappers here; split into per-variant files only if this file
 grows large.
 """
 
+import math
+
 import torch
 import torch.nn.functional as F
 
 from astrai.extension.loader import _available, _modules
 
+_LAYOUT_CODES: dict[str, int] = {"bhld": 0, "blhd": 1}
+
+
+def _parse_layout(layout: str | int) -> int:
+    if isinstance(layout, int):
+        return layout
+    code = _LAYOUT_CODES.get(layout.lower())
+    if code is None:
+        raise ValueError(
+            f"unknown layout '{layout}', expected one of {list(_LAYOUT_CODES)}"
+        )
+    return code
+
+
+def _to_bhld(t: torch.Tensor, layout: int) -> torch.Tensor:
+    """Normalize to b h l d view. Zero-copy transpose if layout==1 (b l h d)."""
+    if layout == 1:
+        return t.transpose(1, 2)
+    return t
+
 
 def _expand_kv_heads(
     k: torch.Tensor, v: torch.Tensor, q_head: int
@@ -26,20 +54,84 @@ def _expand_kv_heads(
     return k, v
 
 
+def _build_attn_mask(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    mask: torch.Tensor | None,
+    causal_offset: int,
+    scale: float,
+) -> tuple[torch.Tensor | None, float]:
+    """Build SDPA-compatible attn_mask + resolved scale.
+
+    q and k must already be in b h l d layout.
+    Causal and mask can coexist: causal sets -inf above the diagonal, mask
+    sets -inf for padded positions. Both are OR'd into a single bool mask.
+    """
+    q_len = q.size(2)
+    kv_len = k.size(2)
+    head_dim = q.size(3)
+    resolved_scale = scale if scale and scale > 0 else 1.0 / math.sqrt(head_dim)
+
+    attn_mask = None
+
+    if mask is not None:
+        if mask.dim() == 2:
+            # [batch, kv_len] → [batch, 1, 1, kv_len]
+            attn_mask = mask[:, None, None, :]
+        elif mask.dim() == 3:
+            # [batch, q_len, kv_len] → [batch, 1, q_len, kv_len]
+            attn_mask = mask[:, None, :, :]
+        else:
+            raise ValueError(f"mask must be 2D or 3D, got {mask.dim()}D")
+
+    if causal_offset >= 0:
+        batch = q.size(0)
+        # q row i attends to kv cols 0..(causal_offset + i)
+        q_idx = torch.arange(q_len, device=q.device).unsqueeze(1)  # [q_len, 1]
+        kv_idx = torch.arange(kv_len, device=q.device).unsqueeze(0)  # [1, kv_len]
+        causal_bool = kv_idx > (causal_offset + q_idx)  # True = masked out
+        causal_mask = causal_bool.unsqueeze(0).expand(
+            batch, -1, -1
+        )  # [batch, q_len, kv_len]
+        causal_mask = causal_mask[:, None, :, :]  # [batch, 1, q_len, kv_len]
+
+        if attn_mask is not None:
+            attn_mask = attn_mask | causal_mask
+        else:
+            attn_mask = causal_mask
+
+    return attn_mask, resolved_scale
+
+
 def _torch_fallback(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
     mask: torch.Tensor | None,
-    is_causal: bool,
+    causal_offset: int,
-    scale: float | None,
+    scale: float,
+    q_layout: int,
+    kv_layout: int | None = None,
 ) -> torch.Tensor:
-    """Reference attention via ``scaled_dot_product_attention``."""
+    """Reference attention via ``scaled_dot_product_attention``.
+
+    q_layout / kv_layout: 0 = b h l d, 1 = b l h d.
+    If kv_layout is None, uses q_layout (Q and K/V share the same layout).
+    """
+    if kv_layout is None:
+        kv_layout = q_layout
+    q = _to_bhld(q, q_layout)
+    k = _to_bhld(k, kv_layout)
+    v = _to_bhld(v, kv_layout)
     k, v = _expand_kv_heads(k, v, q.size(1))
-    attn_mask = mask[:, None, None, :] if mask is not None else None
+    attn_mask, resolved_scale = _build_attn_mask(q, k, mask, causal_offset, scale)
-    return F.scaled_dot_product_attention(
+    out = F.scaled_dot_product_attention(
-        q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale
+        q, k, v, attn_mask=attn_mask, is_causal=False, scale=resolved_scale
     )
+    # Restore Q's original layout
+    if q_layout == 1:
+        out = out.transpose(1, 2)
+    return out
 
 
 def _gather_kv_from_pages(
@@ -56,23 +148,22 @@ def _gather_kv_from_pages(
         k_cache    : [n_pages, page_size, n_kv_heads, head_dim]
         v_cache    : same as k_cache
     Returns:
-        k, v : [batch, n_kv_heads, kv_len, head_dim]
+        k, v : [batch, kv_len, n_kv_heads, head_dim]  (b l h d)
     """
     batch, max_pages = page_table.shape
-    n_pages, ps, n_kv_heads, head_dim = k_cache.shape
+    _, ps, n_kv_heads, head_dim = k_cache.shape
     if ps != page_size:
         raise ValueError(f"k_cache page_size mismatch: {ps} vs {page_size}")
 
-    k = k_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
+    # Vectorized gather: build physical page + offset indices, then advanced-index
-    v = v_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
+    positions = torch.arange(kv_len, device=page_table.device)
+    logical_pages = positions // page_size  # [kv_len]
+    page_offsets = positions % page_size  # [kv_len]
 
-    for b in range(batch):
+    phys_pages = page_table[:, logical_pages]  # [batch, kv_len]
-        for pos in range(kv_len):
+    # k_cache[phys_pages, page_offsets] → [batch, kv_len, n_kv_heads, head_dim] (b l h d)
-            log_pg = pos // page_size
+    k = k_cache[phys_pages, page_offsets]
-            pg_off = pos % page_size
+    v = v_cache[phys_pages, page_offsets]
-            phys = int(page_table[b, log_pg].item())
-            k[b, :, pos, :] = k_cache[phys, pg_off, :, :]
-            v[b, :, pos, :] = v_cache[phys, pg_off, :, :]
     return k, v
 
 
@@ -81,21 +172,22 @@ def attn_decode(
     k: torch.Tensor,
     v: torch.Tensor,
     mask: torch.Tensor | None = None,
-    is_causal: bool = False,
+    causal_offset: int = -1,
-    causal_offset: int = 0,
+    scale: float = 0.0,
-    scale: float | None = None,
+    layout: str = "bhld",
 ) -> torch.Tensor:
+    li = _parse_layout(layout)
     if _available["attn_decode"]:
         return _modules["attn_decode"].attn_decode(
             q,
             k,
             v,
             mask=mask,
-            is_causal=is_causal,
             causal_offset=causal_offset,
             scale=scale,
+            layout=li,
         )
-    return _torch_fallback(q, k, v, mask, is_causal, scale)
+    return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li)
 
 
 def attn_prefill(
@@ -103,21 +195,22 @@ def attn_prefill(
     k: torch.Tensor,
     v: torch.Tensor,
     mask: torch.Tensor | None = None,
-    is_causal: bool = False,
+    causal_offset: int = -1,
-    causal_offset: int = 0,
+    scale: float = 0.0,
-    scale: float | None = None,
+    layout: str = "bhld",
 ) -> torch.Tensor:
+    li = _parse_layout(layout)
     if _available["attn_prefill"]:
         return _modules["attn_prefill"].attn_prefill(
             q,
             k,
             v,
             mask=mask,
-            is_causal=is_causal,
             causal_offset=causal_offset,
             scale=scale,
+            layout=li,
         )
-    return _torch_fallback(q, k, v, mask, is_causal, scale)
+    return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li)
 
 
 def attn_paged_decode(
@@ -128,10 +221,11 @@ def attn_paged_decode(
     page_size: int,
     kv_len: int,
     mask: torch.Tensor | None = None,
-    is_causal: bool = False,
+    causal_offset: int = -1,
-    causal_offset: int = 0,
+    scale: float = 0.0,
-    scale: float | None = None,
+    layout: str = "bhld",
 ) -> torch.Tensor:
+    li = _parse_layout(layout)
     if _available["attn_paged_decode"]:
         return _modules["attn_paged_decode"].attn_paged_decode(
             q,
@@ -141,9 +235,12 @@ def attn_paged_decode(
             page_size,
             kv_len,
             mask=mask,
-            is_causal=is_causal,
             causal_offset=causal_offset,
             scale=scale,
+            layout=li,
         )
+    # Gathered K/V are always b l h d
     k, v = _gather_kv_from_pages(page_table, k_cache, v_cache, page_size, kv_len)
-    return _torch_fallback(q, k, v, mask, is_causal, scale)
+    return _torch_fallback(
+        q, k, v, mask, causal_offset, scale, q_layout=li, kv_layout=1
+    )

csrc/kernels/attn_common.h

@@ -10,11 +10,19 @@ struct AttentionParams {
     int kv_len;
     int head_dim;
     int use_mask;
-    int is_causal;
+    int causal_offset;   // -1 = non-causal; >=0 = absolute position of first Q token
-    int causal_offset;
     int num_splits;
     float scale;
 
+    // Q strides (element offsets for each dim — layout-agnostic)
+    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
+    // KV strides (K and V share the same layout — only base pointers differ)
+    int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
+
+    // Mask: 2D [batch, kv_len] (mask_q_stride=0) or 3D [batch, q_len, kv_len]
+    int mask_b_stride;   // = kv_len (both 2D and 3D)
+    int mask_q_stride;   // 2D: 0 (all q rows share); 3D: kv_len
+
     const T* __restrict__ q;
     const T* __restrict__ k;
     const T* __restrict__ v;
@@ -34,7 +42,6 @@ struct PagedAttentionParams {
     int kv_len;
     int head_dim;
     int use_mask;
-    int is_causal;
     int causal_offset;
     float scale;
 
@@ -42,6 +49,13 @@ struct PagedAttentionParams {
     int page_size;
     int max_pages;
 
+    // Q strides (layout-agnostic)
+    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
+
+    // Mask strides (2D or 3D)
+    int mask_b_stride;
+    int mask_q_stride;
+
     const T* __restrict__ q;
     const T* __restrict__ k_cache;
     const T* __restrict__ v_cache;

csrc/kernels/attn_decode.cu

@@ -5,24 +5,12 @@
 #include "attn_decode_split_kv_mma.cuh"
 #endif
 
-static int decode_num_splits(int base_blocks, int tiles_total) {
-    int sm_count = 0;
-    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
-    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
-    return std::max(1, std::min(n, std::min(tiles_total, 32)));
-}
-
 // Scalar fallback: one warp per query head, split-KV across grid.z.
 static void launch_scalar_decode(AttentionParams<bf16>& p) {
     int group_size = p.q_head / p.kv_head;
     int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
-    p.num_splits = decode_num_splits(p.batch * p.kv_head, chunks_total);
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
-
+    alloc_split_partials(p);
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
-    p.o_part = o_part.data_ptr<float>();
-    p.ml_part = ml_part.data_ptr<float>();
 
     size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
     attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, group_size), smem>>>(p);
@@ -38,13 +26,8 @@ static void launch_scalar_decode(AttentionParams<bf16>& p) {
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_mma_decode(AttentionParams<bf16>& p) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = decode_num_splits(p.batch * p.kv_head, tiles_total);
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
-
+    alloc_split_partials(p);
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
-    p.o_part = o_part.data_ptr<float>();
-    p.ml_part = ml_part.data_ptr<float>();
 
     attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
@@ -68,35 +51,21 @@ torch::Tensor attn_decode(
     torch::Tensor k,
     torch::Tensor v,
     c10::optional<torch::Tensor> mask,
-    bool is_causal = false,
+    int64_t causal_offset,
-    int64_t causal_offset = 0,
+    double scale,
-    c10::optional<double> scale = c10::nullopt
+    int64_t layout
 ) {
     AttentionParams<bf16> p;
-    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
+    attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
     TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
     TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
 
-    auto O = torch::empty_like(q);
+    // O matches Q's original layout
-    p.o = (bf16*)O.data_ptr();
+    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
+    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
+    p.o = (bf16*)O_view.data_ptr();
 
-    switch (p.head_dim) {
+    dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_decode<D>(p); });
-        case 32:
-            dispatch_decode<32>(p);
-            break;
-        case 64:
-            dispatch_decode<64>(p);
-            break;
-        case 128:
-            dispatch_decode<128>(p);
-            break;
-        case 256:
-            dispatch_decode<256>(p);
-            break;
-        default:
-            TORCH_CHECK(false, "decode: unsupported head_dim ", p.head_dim,
-                        " (supported: 32, 64, 128, 256)");
-    }
     return O;
 }
 
@@ -106,8 +75,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("k"),
         py::arg("v"),
         py::arg("mask") = py::none(),
-        py::arg("is_causal") = false,
+        py::arg("causal_offset") = -1,
-        py::arg("causal_offset") = 0,
+        py::arg("scale") = 0.0,
-        py::arg("scale") = py::none(),
+        py::arg("layout") = 0,
         "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
 }

csrc/kernels/attn_decode_split_kv.cuh

@@ -21,13 +21,16 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
     int lane = threadIdx.x;
     int hd_per_thread = p.head_dim / 32;
 
+    // Q: [batch, q_head, q_len=1, head_dim] — stride-based
     float q_reg[8];
-    int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
+    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
+              + lane * hd_per_thread * p.q_stride_d;
     for (int i = 0; i < hd_per_thread; i++)
-        q_reg[i] = __bfloat162float(p.q[q_off + i]);
+        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
 
-    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim;
+    // KV: [batch, kv_head, kv_len, head_dim] — stride-based base
-    int mask_base = batch * p.kv_len;
+    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
+    int mask_base = batch * p.mask_b_stride;
 
     float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
 
@@ -43,9 +46,15 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
         int chunk_start = ci * DC_CHUNK;
         int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start);
 
+        // Load K into shared memory (gather from strided global)
         int total = this_chunk * p.head_dim;
-        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y)
+        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) {
-            k_smem[i] = p.k[kv_base + chunk_start * p.head_dim + i];
+            int s = i / p.head_dim;
+            int d_dim = i % p.head_dim;
+            int kv_idx = chunk_start + s;
+            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
+            k_smem[i] = p.k[g_off];
+        }
         __syncthreads();
 
         for (int s = 0; s < this_chunk; s++) {
@@ -54,9 +63,10 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
                 partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
             partial = warp_reduce_sum(partial) * p.scale;
 
-            if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s])
+            int kv_idx = chunk_start + s;
+            if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx])
                 partial = -FLT_MAX;
-            if (p.is_causal && (chunk_start + s) > p.causal_offset)
+            if (p.causal_offset >= 0 && kv_idx > p.causal_offset)
                 partial = -FLT_MAX;
 
             float new_m = fmaxf(m, partial);
@@ -64,9 +74,10 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
             float beta  = expf(partial - new_m);
             d = d * alpha + beta;
 
-            int v_off = kv_base + (chunk_start + s) * p.head_dim + lane * hd_per_thread;
+            // V: stride-based read
+            int v_off = kv_base + kv_idx * p.kv_stride_l + lane * hd_per_thread * p.kv_stride_d;
             for (int i = 0; i < hd_per_thread; i++)
-                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i]) * beta;
+                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i * p.kv_stride_d]) * beta;
             m = new_m;
         }
         __syncthreads();
@@ -94,6 +105,9 @@ __global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
     int d = threadIdx.x;
     if (d >= p.head_dim) return;
 
+    int batch = bh / p.q_head;
+    int q_head = bh % p.q_head;
+
     size_t split_base = (size_t)bh * p.num_splits;
     const float* mlp = p.ml_part + split_base * 2;
     const float* op = p.o_part + split_base * p.head_dim;
@@ -112,5 +126,7 @@ __global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
     }
 
     float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
-    p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
+    // Stride-based output write (q_len=1 for decode, so stride_l not needed)
+    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
+    p.o[o_off] = __float2bfloat16(acc * inv);
 }

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -68,7 +68,8 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     // ---- Load Q directly from global into mma A-operand registers ----
     // Same layout as prefill: frag[0]/[2] = row gid, frag[1]/[3] = row gid+8
     // cols kt*16 + tid4*2 + {0,1} / +{8,9}. pau[0]=cols c,c+1; pau[4]=c+8,c+9.
-    const int q_base = (batch * p.q_head + q_head0) * HEAD_DIM;
+    // Stride-based: Q is [batch, q_head, q_len=1, head_dim]
+    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
     const int qra = gid;
     const int qrb = gid + 8;
     const bool va = qra < G, vb = qrb < G;
@@ -77,9 +78,9 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     for (int kt = 0; kt < KD; kt++) {
         int c = kt * 16 + tid4 * 2;
         const unsigned* pau = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qra * HEAD_DIM + c]);
+            &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]);
         const unsigned* pbu = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qrb * HEAD_DIM + c]);
+            &p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]);
         Qa[kt][0] = va ? pau[0] : 0u;
         Qa[kt][1] = vb ? pbu[0] : 0u;
         Qa[kt][2] = va ? pau[4] : 0u;
@@ -92,8 +93,9 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
     float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
 
-    const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM;
+    // KV: stride-based base — [batch, kv_head, kv_len, head_dim]
-    const int mask_base = batch * p.kv_len;
+    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
+    const int mask_batch_base = batch * p.mask_b_stride;
     const int tiles_total = (p.kv_len + BC - 1) / BC;
     const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
     const int ti_begin = split * tiles_per_split;
@@ -111,8 +113,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             int kc = kv0 + r;
             bool valid = kc < p.kv_len;
             int off = r * LD + swiz_col(d, r, SWIZ_MASK);
-            cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid);
+            // KV stride-based: contiguous within head_dim (stride_d == 1 typically)
-            cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + d], valid);
+            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
+            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
+            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
         }
         cp_async_commit();
     };
@@ -148,9 +152,13 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
             Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
 
-        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
+        // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant
+        int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
         mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
-                                    mask_base, p.mask, has_mask,
+                                    0, 0,
+                                    mask_batch_base, 0,
+                                    batch,
+                                    p.mask, has_mask,
                                     Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);

csrc/kernels/attn_entry_utils.cuh

@@ -1,43 +1,214 @@
 #pragma once
 #include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
 #include "attn_common.h"
 
+using bf16 = __nv_bfloat16;
+
+// ---------------------------------------------------------------------------
+// Shared dispatch helpers — eliminates duplication across .cu entry files.
+// ---------------------------------------------------------------------------
+
+inline int compute_num_splits(int base_blocks, int tiles_total) {
+    int sm_count = 0;
+    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
+    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
+    return std::max(1, std::min(n, std::min(tiles_total, 32)));
+}
+
+// Dispatch head_dim to a generic lambda FN (callable as FN.operator()<D>()).
+template <typename Fn>
+inline void dispatch_head_dim(int hd, Fn&& fn) {
+    switch (hd) {
+        case 32:  fn.template operator()<32>();  break;
+        case 64:  fn.template operator()<64>();  break;
+        case 128: fn.template operator()<128>(); break;
+        case 256: fn.template operator()<256>(); break;
+        default:
+            TORCH_CHECK(false, "unsupported head_dim ", hd,
+                         " (supported: 32, 64, 128, 256)");
+    }
+}
+
+// Allocate split-KV partial buffers and wire into params.
+template<typename P>
+inline void alloc_split_partials(P& p) {
+    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
+    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
+    p.o_part = (float*)o_part.data_ptr();
+    p.ml_part = (float*)ml_part.data_ptr();
+}
+
+// ---------------------------------------------------------------------------
+// Param packing — fills AttentionParams from torch tensors with validation.
+// ---------------------------------------------------------------------------
+
 template<typename T>
 inline void attn_pack_params(
     torch::Tensor q,
     torch::Tensor k,
     torch::Tensor v,
     c10::optional<torch::Tensor> mask,
-    bool is_causal,
     int64_t causal_offset,
-    c10::optional<double> scale,
+    double scale,
+    int64_t layout,   // 0 = b h l d, 1 = b l h d
     AttentionParams<T>& p
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+
     TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
     TORCH_CHECK(q.dtype() == torch::kBFloat16);
     TORCH_CHECK(k.dtype() == torch::kBFloat16);
     TORCH_CHECK(v.dtype() == torch::kBFloat16);
+    TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
+    TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
+
+    // Normalize to b h l d view (zero-copy transpose if user passed b l h d)
+    if (layout == 1) {
+        q = q.transpose(1, 2);
+        k = k.transpose(1, 2);
+        v = v.transpose(1, 2);
+    }
 
     p.batch = (int)q.size(0);
     p.q_head = (int)q.size(1);
-    p.kv_head = (int)k.size(1);
     p.q_len = (int)q.size(2);
-    p.kv_len = (int)k.size(2);
     p.head_dim = (int)q.size(3);
-    p.use_mask = mask.has_value() ? 1 : 0;
+    p.kv_head = (int)k.size(1);
-    p.is_causal = is_causal ? 1 : 0;
+    p.kv_len = (int)k.size(2);
+    TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
+
+    // Strides (layout-agnostic: works for b h l d and b l h d)
+    p.q_stride_b = (int)q.stride(0);
+    p.q_stride_h = (int)q.stride(1);
+    p.q_stride_l = (int)q.stride(2);
+    p.q_stride_d = (int)q.stride(3);
+    p.kv_stride_b = (int)k.stride(0);
+    p.kv_stride_h = (int)k.stride(1);
+    p.kv_stride_l = (int)k.stride(2);
+    p.kv_stride_d = (int)k.stride(3);
+
     p.causal_offset = (int)causal_offset;
-    p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
+    p.use_mask = mask.has_value() ? 1 : 0;
+    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
+
     p.q = (const T*)q.data_ptr();
     p.k = (const T*)k.data_ptr();
     p.v = (const T*)v.data_ptr();
+    p.o = nullptr;
+    p.o_part = nullptr;
+    p.ml_part = nullptr;
+
     if (p.use_mask) {
-        TORCH_CHECK(mask.value().dtype() == torch::kBool);
+        auto m = mask.value();
-        TORCH_CHECK(mask.value().dim() == 2);
+        TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
-        TORCH_CHECK(mask.value().size(0) == p.batch);
+        TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
-        TORCH_CHECK(mask.value().size(1) == p.kv_len);
+        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
-        p.mask = mask.value().data_ptr<bool>();
+        TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
+        if (m.dim() == 2) {
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_q_stride = 0;
+        } else if (m.dim() == 3) {
+            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_q_stride = (int)m.stride(1);
+        } else {
+            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
+        }
+        p.mask = m.data_ptr<bool>();
     } else {
         p.mask = nullptr;
+        p.mask_b_stride = 0;
+        p.mask_q_stride = 0;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Param packing for paged attention.
+// ---------------------------------------------------------------------------
+
+template<typename T>
+inline void attn_pack_paged_params(
+    torch::Tensor q,
+    torch::Tensor page_table,
+    torch::Tensor k_cache,
+    torch::Tensor v_cache,
+    int64_t page_size,
+    int64_t kv_len,
+    c10::optional<torch::Tensor> mask,
+    int64_t causal_offset,
+    double scale,
+    int64_t layout,   // 0 = b h l d, 1 = b l h d
+    PagedAttentionParams<T>& p
+) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+
+    TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16");
+    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16");
+    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16");
+    TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
+    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
+
+    // Normalize Q to b h l d view if user passed b l h d
+    if (layout == 1) {
+        q = q.transpose(1, 2);
+    }
+
+    p.batch = (int)q.size(0);
+    p.q_head = (int)q.size(1);
+    p.q_len = (int)q.size(2);
+    p.head_dim = (int)q.size(3);
+    p.kv_head = (int)k_cache.size(2);
+    p.kv_len = (int)kv_len;
+    p.page_size = (int)page_size;
+    p.max_pages = (int)page_table.size(1);
+
+    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
+    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
+    TORCH_CHECK(k_cache.size(1) == page_size,
+                "k_cache dim 1 must equal page_size, got ",
+                k_cache.size(1), " vs ", page_size);
+
+    // Q strides
+    p.q_stride_b = (int)q.stride(0);
+    p.q_stride_h = (int)q.stride(1);
+    p.q_stride_l = (int)q.stride(2);
+    p.q_stride_d = (int)q.stride(3);
+
+    p.causal_offset = (int)causal_offset;
+    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
+    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
+
+    p.page_table = page_table.data_ptr<int64_t>();
+    p.k_cache = (const T*)k_cache.data_ptr();
+    p.v_cache = (const T*)v_cache.data_ptr();
+    p.q = (const T*)q.data_ptr();
+    p.o = nullptr;
+    p.o_part = nullptr;
+    p.ml_part = nullptr;
+
+    if (p.use_mask) {
+        auto m = mask.value();
+        TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
+        TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
+        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
+        TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
+        if (m.dim() == 2) {
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_q_stride = 0;
+        } else if (m.dim() == 3) {
+            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
+            p.mask_b_stride = (int)m.stride(0);
+            p.mask_q_stride = (int)m.stride(1);
+        } else {
+            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
+        }
+        p.mask = m.data_ptr<bool>();
+    } else {
+        p.mask = nullptr;
+        p.mask_b_stride = 0;
+        p.mask_q_stride = 0;
     }
 }

csrc/kernels/attn_mma_utils.cuh

@@ -122,7 +122,9 @@ template <int KD, int NC8>
 __device__ inline void mma_compute_scores(
     const unsigned Qa[KD][4],
     const bf16* __restrict__ sK,
-    int LD, int SWIZ_MASK, int lane,
+    int LD, 
+    int SWIZ_MASK, 
+    int lane,
     float Sacc[NC8][4])
 {
     #pragma unroll
@@ -142,13 +144,20 @@ __device__ inline void mma_compute_scores(
 // Online softmax + Oacc rescale for one K/V tile.
 //   maxc0/maxc1: per-row KV column bounds (prefill: per-query-row causal limits;
 //     decode: same value for both rows since q_len==1).
+//   qrow0/qrow1: query row indices (for 3D mask indexing; decode passes 0).
+//   mask_b_stride/mask_q_stride: mask layout (2D: mask_q_stride=0; 3D: =kv_len).
 // Reads Sacc (Q@K^T scores), applies causal/mask, computes P = exp(S - nm),
 // rescales Oacc by exp(m_old - nm), and updates m/l — all in place.
 template <int NC8, int DN8>
 __device__ inline void mma_softmax_tile(
     int kv0,
-    int maxc0, int maxc1,
+    int maxc0, 
-    int mask_base,
+    int maxc1,
+    int qrow0, 
+    int qrow1,
+    int mask_b_stride, 
+    int mask_q_stride,
+    int mask_batch,
     const bool* __restrict__ mask,
     bool has_mask,
     float Sacc[NC8][4],
@@ -162,14 +171,16 @@ __device__ inline void mma_softmax_tile(
     // Mask out-of-bounds / masked columns: set -FLT_MAX so expf → 0 downstream
     // without per-element sentinel checks. Compute tile-local row maxima.
     float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
+    int mask_base0 = mask_batch * mask_b_stride + qrow0 * mask_q_stride;
+    int mask_base1 = mask_batch * mask_b_stride + qrow1 * mask_q_stride;
     #pragma unroll
     for (int n8 = 0; n8 < NC8; n8++) {
         int cc = kv0 + n8 * 8 + 2 * tid4;
         int c1 = cc + 1;
-        bool b0 = (cc >= maxc0) || (has_mask && !mask[mask_base + cc]);
+        bool b0 = (cc >= maxc0) || (has_mask && !mask[mask_base0 + cc]);
-        bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base + c1]);
+        bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base0 + c1]);
-        bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base + cc]);
+        bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base1 + cc]);
-        bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base + c1]);
+        bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base1 + c1]);
         float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
         float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
         float s2 = b2 ? -FLT_MAX : Sacc[n8][2];

csrc/kernels/attn_paged_decode.cu

@@ -3,26 +3,13 @@
 #include "attn_paged_decode_split_kv_mma.cuh"
 #endif
 
-#include <torch/extension.h>
+#include "attn_entry_utils.cuh"
-#include <c10/cuda/CUDAGuard.h>
-
-static int paged_decode_num_splits(int base_blocks, int tiles_total) {
-    int sm_count = 0;
-    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
-    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
-    return std::max(1, std::min(n, std::min(tiles_total, 32)));
-}
 
 static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
     int group_size = p.q_head / p.kv_head;
     int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
-    p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, chunks_total);
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
-
+    alloc_split_partials(p);
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
-    p.o_part = o_part.data_ptr<float>();
-    p.ml_part = ml_part.data_ptr<float>();
 
     size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
     dim3 grid = dim3(p.batch * p.kv_head, 1, p.num_splits);
@@ -35,13 +22,8 @@ static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, tiles_total);
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
-
+    alloc_split_partials(p);
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
-    p.o_part = o_part.data_ptr<float>();
-    p.ml_part = ml_part.data_ptr<float>();
 
     paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
@@ -68,68 +50,19 @@ torch::Tensor attn_paged_decode(
     int64_t page_size,
     int64_t kv_len,
     c10::optional<torch::Tensor> mask,
-    bool is_causal = false,
+    int64_t causal_offset,
-    int64_t causal_offset = 0,
+    double scale,
-    c10::optional<double> scale = c10::nullopt
+    int64_t layout
 ) {
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    PagedAttentionParams<bf16> p;
+    attn_pack_paged_params(q, page_table, k_cache, v_cache,
+                           page_size, kv_len, mask, causal_offset, scale, layout, p);
 
-    int batch = q.size(0);
+    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
-    int q_head = q.size(1);
+    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
-    int head_dim = q.size(3);
+    p.o = (bf16*)O_view.data_ptr();
-    int kv_head = k_cache.size(2);
-    int max_pages = page_table.size(1);
-
-    TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
-    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16");
-    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16");
-    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16");
-    TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
-    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
-    TORCH_CHECK(head_dim % 32 == 0, "head_dim must be multiple of 32");
-    TORCH_CHECK(k_cache.size(1) == page_size,
-                "k_cache dim 1 must equal page_size, got ",
-                k_cache.size(1), " vs ", page_size);
-    TORCH_CHECK(k_cache.size(0) >= 0, "k_cache must have at least 0 pages");
-
-    float scale_val = scale.has_value()
-        ? static_cast<float>(scale.value())
-        : 1.0f / std::sqrt(static_cast<float>(head_dim));
-
-    auto O = torch::empty_like(q);
-
-    PagedAttentionParams<bf16, float> p;
-    p.batch = batch;
-    p.q_head = q_head;
-    p.kv_head = kv_head;
-    p.q_len = static_cast<int>(q.size(2));
-    p.kv_len = static_cast<int>(kv_len);
-    p.head_dim = head_dim;
-    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
-    p.is_causal = is_causal ? 1 : 0;
-    p.causal_offset = static_cast<int>(causal_offset);
-    p.scale = scale_val;
-    p.page_size = static_cast<int>(page_size);
-    p.max_pages = max_pages;
-    p.page_table = page_table.data_ptr<int64_t>();
-    p.k_cache = reinterpret_cast<const bf16*>(k_cache.data_ptr());
-    p.v_cache = reinterpret_cast<const bf16*>(v_cache.data_ptr());
-    p.q = reinterpret_cast<const bf16*>(q.data_ptr());
-    p.mask = p.use_mask ? mask.value().data_ptr<bool>() : nullptr;
-    p.o = reinterpret_cast<bf16*>(O.data_ptr());
-    p.o_part = nullptr;
-    p.ml_part = nullptr;
-
-    switch (p.head_dim) {
-        case 32:  dispatch_paged_decode<32>(p);  break;
-        case 64:  dispatch_paged_decode<64>(p);  break;
-        case 128: dispatch_paged_decode<128>(p); break;
-        case 256: dispatch_paged_decode<256>(p); break;
-        default:
-            TORCH_CHECK(false, "paged_decode: unsupported head_dim ", p.head_dim,
-                        " (supported: 32, 64, 128, 256)");
-    }
 
+    dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_paged_decode<D>(p); });
     return O;
 }
 
@@ -142,8 +75,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("page_size"),
         py::arg("kv_len"),
         py::arg("mask") = py::none(),
-        py::arg("is_causal") = false,
+        py::arg("causal_offset") = -1,
-        py::arg("causal_offset") = 0,
+        py::arg("scale") = 0.0,
-        py::arg("scale") = py::none(),
+        py::arg("layout") = 0,
         "Paged GQA decode — split-KV with direct page-table access.");
 }

csrc/kernels/attn_paged_decode_split_kv.cuh

@@ -22,11 +22,13 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
     int lane = threadIdx.x;
     int hd_per_thread = p.head_dim / 32;
 
+    // Q: stride-based [batch, q_head, q_len=1, head_dim]
     float q_reg[8];
-    int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
+    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
+              + lane * hd_per_thread * p.q_stride_d;
     #pragma unroll
     for (int i = 0; i < hd_per_thread; i++)
-        q_reg[i] = __bfloat162float(p.q[q_off + i]);
+        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
 
     float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
 
@@ -37,7 +39,7 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
     int ch_begin = split * chunks_per_split;
     int ch_end = min(chunks_total, ch_begin + chunks_per_split);
 
-    const int mask_base = batch * p.kv_len;
+    const int mask_base = batch * p.mask_b_stride;
 
     for (int ci = ch_begin; ci < ch_end; ci++) {
         int chunk_start = ci * PDC_CHUNK;
@@ -70,9 +72,10 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
                 partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
             partial = paged_warp_reduce_sum(partial) * p.scale;
 
-            if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s])
+            int kv_idx = chunk_start + s;
+            if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx])
                 partial = -FLT_MAX;
-            if (p.is_causal && (chunk_start + s) > p.causal_offset)
+            if (p.causal_offset >= 0 && kv_idx > p.causal_offset)
                 partial = -FLT_MAX;
 
             float new_m = fmaxf(m, partial);
@@ -118,6 +121,9 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
     int d = threadIdx.x;
     if (d >= p.head_dim) return;
 
+    int batch = bh / p.q_head;
+    int q_head = bh % p.q_head;
+
     size_t split_base = (size_t)bh * p.num_splits;
     const float* mlp = p.ml_part + split_base * 2;
     const float* op = p.o_part + split_base * p.head_dim;
@@ -136,5 +142,6 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
     }
 
     float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
-    p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
+    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
+    p.o[o_off] = __float2bfloat16(acc * inv);
 }

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -42,7 +42,8 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
 
     // ---- Load Q directly from global into mma A-operand registers ----
-    const int q_base = (batch * p.q_head + q_head0) * HEAD_DIM;
+    // Q stride-based: [batch, q_head, q_len=1, head_dim]
+    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
     const int qra = gid;
     const int qrb = gid + 8;
     const bool va = qra < G, vb = qrb < G;
@@ -51,9 +52,9 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     for (int kt = 0; kt < KD; kt++) {
         int c = kt * 16 + tid4 * 2;
         const unsigned* pau = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qra * HEAD_DIM + c]);
+            &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]);
         const unsigned* pbu = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qrb * HEAD_DIM + c]);
+            &p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]);
         Qa[kt][0] = va ? pau[0] : 0u;
         Qa[kt][1] = vb ? pbu[0] : 0u;
         Qa[kt][2] = va ? pau[4] : 0u;
@@ -66,7 +67,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
     float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
 
-    const int mask_base = batch * p.kv_len;
+    const int mask_batch_base = batch * p.mask_b_stride;
     const int tiles_total = (p.kv_len + BC - 1) / BC;
     const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
     const int ti_begin = split * tiles_per_split;
@@ -130,9 +131,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
             Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
             Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
 
-        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
+        // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant
+        int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
         mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
-                                    mask_base, p.mask, has_mask,
+                                    0, 0,
+                                    mask_batch_base, 0,
+                                    batch,
+                                    p.mask, has_mask,
                                     Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);

csrc/kernels/attn_prefill.cu

@@ -35,34 +35,19 @@ torch::Tensor attn_prefill(
     torch::Tensor k,
     torch::Tensor v,
     c10::optional<torch::Tensor> mask,
-    bool is_causal = false,
+    int64_t causal_offset,
-    int64_t causal_offset = 0,
+    double scale,
-    c10::optional<double> scale = c10::nullopt
+    int64_t layout
 ) {
     AttentionParams<bf16> p;
-    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
+    attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
     TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
 
-    auto O = torch::empty_like(q);
+    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
-    p.o = (bf16*)O.data_ptr();
+    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
+    p.o = (bf16*)O_view.data_ptr();
 
-    switch (p.head_dim) {
+    dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_prefill<D>(p); });
-        case 32:
-            dispatch_prefill<32>(p);
-            break;
-        case 64:
-            dispatch_prefill<64>(p);
-            break;
-        case 128:
-            dispatch_prefill<128>(p);
-            break;
-        case 256:
-            dispatch_prefill<256>(p);
-            break;
-        default:
-            TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
-                        " (supported: 32,64,128,256)");
-    }
     return O;
 }
 
@@ -72,8 +57,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("k"),
         py::arg("v"),
         py::arg("mask") = py::none(),
-        py::arg("is_causal") = false,
+        py::arg("causal_offset") = -1,
-        py::arg("causal_offset") = 0,
+        py::arg("scale") = 0.0,
-        py::arg("scale") = py::none(),
+        py::arg("layout") = 0,
         "GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
 }

csrc/kernels/attn_prefill_split_q.cuh

@@ -50,12 +50,14 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM];
     __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM];
 
+    // Q: stride-based load [batch, q_head, q_len, head_dim]
     float qreg[DPT];
     if (q_row < p.q_len) {
-        int q_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
+        int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
+                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
 #pragma unroll
         for (int i = 0; i < DPT; i++)
-            qreg[i] = __bfloat162float(p.q[q_off + i]) * p.scale;
+            qreg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]) * p.scale;
     }
 
     float m = -FLT_MAX, l = 0.0f;
@@ -64,7 +66,9 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     for (int i = 0; i < DPT; i++)
         acc[i] = 0.0f;
 
-    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM;
+    // KV: stride-based base
+    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
+    int mask_batch_base = batch * p.mask_b_stride;
     int tiles   = (p.kv_len + P_BC - 1) / P_BC;
     int tt      = G * ROWS;
     int lid     = row * G + gpos;
@@ -79,15 +83,19 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
         int kv0  = ti * P_BC;
         int tlen = min(P_BC, p.kv_len - kv0);
 
+        // Load K/V into shared memory from strided global
         for (int i = lid; i < tlen * HEAD_DIM; i += tt) {
-            int gidx = kv_base + (kv0 + i / HEAD_DIM) * HEAD_DIM + (i % HEAD_DIM);
+            int s = i / HEAD_DIM;
-            sK[i] = p.k[gidx];
+            int d_dim = i % HEAD_DIM;
-            sV[i] = p.v[gidx];
+            int kv_idx = kv0 + s;
+            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
+            sK[i] = p.k[g_off];
+            sV[i] = p.v[g_off];
         }
         __syncthreads();
 
         int lim = tlen;
-        if (p.is_causal && q_row < p.q_len) {
+        if (p.causal_offset >= 0 && q_row < p.q_len) {
             int ep = q_row + p.causal_offset + 1;
             if (kv0 >= ep)
                 lim = 0;
@@ -95,6 +103,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
                 lim = ep - kv0;
         }
 
+        int mask_row_base = mask_batch_base + q_row * p.mask_q_stride;
         for (int s = 0; s < lim; s++) {
             const bf16* kr = sK + s * HEAD_DIM + gpos * DPT;
             float part = 0.0f;
@@ -108,7 +117,8 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
             }
             float dot = group_reduce_sum<G>(part, gmask);
 
-            if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s])
+            int kv_idx = kv0 + s;
+            if (p.use_mask && p.mask && !p.mask[mask_row_base + kv_idx])
                 dot = -FLT_MAX;
 
             float nm = fmaxf(m, dot);
@@ -131,10 +141,12 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     }
 
     if (q_row < p.q_len) {
-        int o_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
+        // O: stride-based write
+        int o_off = batch * p.q_stride_b + q_head * p.q_stride_h
+                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
         float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
 #pragma unroll
         for (int i = 0; i < DPT; i++)
-            p.o[o_off + i] = __float2bfloat16(acc[i] * rl);
+            p.o[o_off + i * p.q_stride_d] = __float2bfloat16(acc[i] * rl);
     }
 }

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -72,7 +72,8 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     // registers across the tile loop.
     //   frag[0]/[2]: row = qrow0 + gid ;  frag[1]/[3]: row = qrow0 + gid + 8
     //   frag[0]/[1]: cols kt*16 + tid4*2 + {0,1} ; frag[2]/[3]: + 8
-    const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
+    // Q stride-based: [batch, q_head, q_len, head_dim]
+    const int q_base = batch * p.q_stride_b + q_head * p.q_stride_h;
     const int qra = qrow0 + gid;
     const int qrb = qrow0 + gid + 8;
     const bool va = qra < p.q_len, vb = qrb < p.q_len;
@@ -81,9 +82,9 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     for (int kt = 0; kt < KD; kt++) {
         int c = kt * 16 + tid4 * 2;
         const unsigned* pau = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qra * HEAD_DIM + c]);
+            &p.q[q_base + qra * p.q_stride_l + c * p.q_stride_d]);
         const unsigned* pbu = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qrb * HEAD_DIM + c]);
+            &p.q[q_base + qrb * p.q_stride_l + c * p.q_stride_d]);
         Qa[kt][0] = va ? pau[0] : 0u;
         Qa[kt][1] = vb ? pbu[0] : 0u;
         Qa[kt][2] = va ? pau[4] : 0u;
@@ -96,18 +97,19 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
         Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
     float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
 
-    const int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM;
+    // KV: stride-based base
+    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
     const int tiles = (p.kv_len + BC - 1) / BC;
     const int qr0 = qrow0 + gid;      // row for c0/c1
     const int qr1 = qrow0 + gid + 8;  // row for c2/c3
 
-    // Causal tile-skip bounds (no-op when is_causal == 0)
+    // Causal tile-skip bounds (no-op when causal_offset < 0)
-    const int use_skip = p.is_causal;
+    const int use_skip = (p.causal_offset >= 0) ? 1 : 0;
     const int max_kv = qrow0 + BR - 1 + p.causal_offset;
     const int block_max_kv =
         blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset;
     const int has_mask = p.use_mask && p.mask;
-    const int mb = batch * p.kv_len;
+    const int mask_batch_base = batch * p.mask_b_stride;
 
     // Last active tile: block-level causal bound (all warps in the block share
     // the K/V load, so the prefetch range is the block max, not per-warp).
@@ -132,8 +134,9 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
             int kc = kv0 + r;
             bool valid = kc < p.kv_len;
             int off = r * LD + swiz_col(d, r, SWIZ_MASK);
-            cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid);
+            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
-            cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + d], valid);
+            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
+            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
         }
         cp_async_commit();
     };
@@ -170,12 +173,15 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
             Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
             Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
 
-        int maxc0 = p.is_causal ? min(p.kv_len, qr0 + p.causal_offset + 1)
+        int maxc0 = (p.causal_offset >= 0) ? min(p.kv_len, qr0 + p.causal_offset + 1)
                                         : p.kv_len;
-        int maxc1 = p.is_causal ? min(p.kv_len, qr1 + p.causal_offset + 1)
+        int maxc1 = (p.causal_offset >= 0) ? min(p.kv_len, qr1 + p.causal_offset + 1)
                                         : p.kv_len;
         mma_softmax_tile<NC8, DN8>(kv0, maxc0, maxc1,
-                                    mb, p.mask, has_mask,
+                                    qr0, qr1,
+                                    mask_batch_base, p.mask_q_stride,
+                                    batch,
+                                    p.mask, has_mask,
                                     Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
@@ -186,19 +192,20 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     // halves store count and removes the uncoalesced scalar-store penalty)
     float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
     float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
-    const int o_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
+    // O: stride-based write
+    const int o_base = batch * p.q_stride_b + q_head * p.q_stride_h;
 #pragma unroll
     for (int dn8 = 0; dn8 < DN8; dn8++) {
         int d = dn8 * 8 + 2 * tid4;
         if (qr0 < p.q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
                                                      Oacc[dn8][1] * rl0);
-            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr0 * HEAD_DIM + d]) = v;
+            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v;
         }
         if (qr1 < p.q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
                                                      Oacc[dn8][3] * rl1);
-            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr1 * HEAD_DIM + d]) = v;
+            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v;
         }
     }
 }

csrc/tests/attn_decode_test.cu

@@ -98,8 +98,9 @@ static void bench() {
 
         AttentionParams<bf16> p;
         p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
-        p.head_dim = D; p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
+        p.head_dim = D; p.use_mask = 0; p.causal_offset = -1;
         p.scale = 1.0f / sqrtf((float)D);
+        set_default_strides(p);
         p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO;
 
         DecodeScratch sc;
@@ -160,7 +161,7 @@ int main() {
 
         AttentionParams<bf16> p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
-        p.use_mask=0; p.is_causal=0; p.causal_offset=0;
+        p.use_mask=0; p.causal_offset=-1;
         p.scale=1.0f/sqrtf((float)D);
         p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
 
@@ -180,7 +181,7 @@ int main() {
         cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
 
         float* ref=new float[nQ];
-        cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, 0, 0);
+        cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, -1);
 
         float max_err=0;
         for (size_t i=0;i<nQ;i++){

csrc/tests/attn_paged_decode_test.cu

@@ -138,13 +138,14 @@ static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed)
     }
 
     float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
-    cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv, 1, kv_len, HEAD_DIM, 0, 0);
+    cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv, 1, kv_len, HEAD_DIM, -1);
 
     float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
     PagedAttentionParams<bf16, float> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1;
     p.kv_len = kv_len; p.head_dim = HEAD_DIM;
-    p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
+    p.use_mask = 0; p.causal_offset = -1;
+    set_default_strides(p);
     p.num_splits = 1; p.scale = scale_val;
     p.page_size = page_size; p.max_pages = max_pages;
     p.page_table = d_pt;
@@ -272,7 +273,8 @@ static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
     PagedAttentionParams<bf16, float> pa;
     pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1;
     pa.kv_len = kv_len; pa.head_dim = HEAD_DIM;
-    pa.use_mask = 0; pa.is_causal = 0; pa.causal_offset = 0;
+    pa.use_mask = 0; pa.causal_offset = -1;
+    set_default_paged_strides(pa);
     pa.num_splits = 1; pa.scale = scale_val;
     pa.page_size = page_size; pa.max_pages = max_pages;
     pa.page_table = d_pt;

csrc/tests/attn_prefill_test.cu

@@ -75,7 +75,8 @@ static void bench() {
 
         AttentionParams<bf16> p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
-        p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
+        p.use_mask=0; p.causal_offset=causal?0:-1;
+        set_default_strides(p);
         p.scale=1.0f/sqrtf((float)D);
         p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
 
@@ -143,7 +144,8 @@ int main() {
 
         AttentionParams<bf16> p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
-        p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
+        p.use_mask=0; p.causal_offset=causal?0:-1;
+        set_default_strides(p);
         p.scale=1.0f/sqrtf((float)D);
         p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
 
@@ -158,7 +160,7 @@ int main() {
         cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
 
         float* ref=new float[nQ];
-        cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal, 0);
+        cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal ? 0 : -1);
 
         float max_err=0;
         for (size_t i=0;i<nQ;i++) {

csrc/tests/test_utils.cuh

@@ -101,6 +101,32 @@ void dispatch_by_head_dim(int head_dim, Fn&& fn) {
     _HeadSwitch<32, 64, 128, 256>::call(head_dim, fn);
 }
 
+// Set default strides for contiguous b h l d layout on AttentionParams.
+template<typename P>
+inline void set_default_strides(P& p) {
+    p.q_stride_b  = p.q_head * p.q_len * p.head_dim;
+    p.q_stride_h  = p.q_len * p.head_dim;
+    p.q_stride_l  = p.head_dim;
+    p.q_stride_d  = 1;
+    p.kv_stride_b = p.kv_head * p.kv_len * p.head_dim;
+    p.kv_stride_h = p.kv_len * p.head_dim;
+    p.kv_stride_l = p.head_dim;
+    p.kv_stride_d = 1;
+    p.mask_b_stride = p.kv_len;
+    p.mask_q_stride = 0;
+}
+
+// Set default Q strides for contiguous b h l d layout on PagedAttentionParams.
+template<typename P>
+inline void set_default_paged_strides(P& p) {
+    p.q_stride_b  = p.q_head * p.q_len * p.head_dim;
+    p.q_stride_h  = p.q_len * p.head_dim;
+    p.q_stride_l  = p.head_dim;
+    p.q_stride_d  = 1;
+    p.mask_b_stride = p.kv_len;
+    p.mask_q_stride = 0;
+}
+
 // Generic CPU reference for multi-query / grouped-query attention.
 // Tensor shapes (all float*):
 //   Q : [B, Hq, q_len, D]
@@ -108,10 +134,11 @@ void dispatch_by_head_dim(int head_dim, Fn&& fn) {
 //   V : [B, Hk, kv_len, D]
 //   O : [B, Hq, q_len, D]
 // mask: if q_len == 1, shape is [B, kv_len]; otherwise mask is not supported.
+// causal_offset: -1 = non-causal; >=0 = absolute position of first Q token.
 static void cpu_attention_ref(
     const float* Q, const float* K, const float* V, const bool* mask,
     float* O, int B, int Hq, int Hk, int q_len, int kv_len, int D,
-    int is_causal, int causal_offset
+    int causal_offset
 ) {
     float scale = 1.0f / sqrtf((float)D);
     int n_rep = Hq / Hk;
@@ -122,7 +149,7 @@ static void cpu_attention_ref(
                 float mv = -INFINITY, sv = 0.0f;
                 float accum[256] = {0.0f};
                 int lim = kv_len;
-                if (is_causal) {
+                if (causal_offset >= 0) {
                     int c = qi + causal_offset + 1;
                     lim = (c < kv_len) ? c : kv_len;
                 }