feat: support NHD paged KV cache layout via kv_layout="NHD"

python/fmha_sm100/api.py

@@ -1035,6 +1035,7 @@ def fmha_sm100(
     q_offset_override: Optional[Union[int, torch.Tensor]] = None,
     out: Optional[torch.Tensor] = None,
     max_score: Optional[torch.Tensor] = None,
+    kv_layout: str = "HND",
     **kwargs
 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
     """Run dense, paged, or sparse SM100 FMHA using a precomputed plan.
@@ -1047,7 +1048,10 @@ def fmha_sm100(
         128.
     k : torch.Tensor
         Dense layout ``[total_kv_len, num_kv_heads, head_dim]`` or paged layout
-        ``[total_pages, num_kv_heads, page_size, head_dim]``.
+        ``[total_pages, num_kv_heads, page_size, head_dim]`` (``kv_layout="HND"``,
+        default).  Paged NHD caches ``[total_pages, page_size, num_kv_heads,
+        head_dim]`` are supported with ``kv_layout="NHD"`` (contiguous NHD
+        caches are consumed zero-copy; requires ``kv_indices``).
     v : torch.Tensor
         Same layout as ``k``.  The output head dimension follows ``v.shape[-1]``.
     plan_info : tuple
@@ -1072,6 +1076,9 @@ def fmha_sm100(
     max_score : torch.Tensor, optional
         Preallocated per-KV-tile score buffer with shape
         ``[num_qo_heads, max_k_tiles, total_qo_len]`` and dtype float32.
+    kv_layout : str, optional
+        Paged K/V cache layout, ``"HND"`` (default) or ``"NHD"``.  Only valid
+        for paged (4D) ``k``/``v``.
     **kwargs
         Runtime options forwarded to the kernel runner.  Common options are
         ``sm_scale``, ``q_scale``, ``k_scale``, ``v_scale``, ``o_scale``,
@@ -1084,6 +1091,16 @@ def fmha_sm100(
         output was disabled.  When both decode and prefill sub-plans are used,
         outputs are concatenated back into the original batch order.
     """
+    if kv_layout != "HND":
+        if kv_layout != "NHD":
+            raise ValueError(f"kv_layout must be 'HND' or 'NHD', got {kv_layout!r}")
+        if kv_indices is None or k.dim() != 4 or v.dim() != 4:
+            raise ValueError(
+                "kv_layout='NHD' requires paged k/v with shape "
+                "[total_pages, page_size, num_kv_heads, head_dim] and kv_indices"
+            )
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
     has_mixed_prefill, split, batch_size, decode, prefill = plan_info
     if not has_mixed_prefill:
         return _fmha_sm100(q, k, v, decode, out=out, max_score=max_score, kv_indices=kv_indices,kv_block_indexes=kv_block_indexes, q_offset_override=q_offset_override, **kwargs)

python/fmha_sm100/csrc/fmha_sm100_inst.jinja

@@ -66,5 +66,6 @@ cudaError_t {{ func_name }}(const FMHACutlassSM100Params& p) {
                                       , p.gmem_bounds
 #endif
                                       },
-      p.num_ctas);
+      p.num_ctas,
+      p.k_stride_t, p.v_stride_t);
 }

python/fmha_sm100/csrc/fmha_sm100_params.h

@@ -34,6 +34,10 @@ struct FMHACutlassSM100Params {
   int k_stride_h;
   int v_stride_n;
   int v_stride_h;
+  // Within-page token stride for paged K/V (elements). 0 selects the
+  // HND-contiguous default head_dim; NHD caches pass Hkv * head_dim.
+  int k_stride_t = 0;
+  int v_stride_t = 0;
   int batch_size;
   int total_qo_len;
   int total_kv_len;

python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja

@@ -84,6 +84,8 @@ void FMHAVariantRun_{{ variant_name }}(ffi::TensorView workspace_buffer, ffi::Te
   params.k_stride_h = k.stride(1);
   params.v_stride_n = v.stride(0);
   params.v_stride_h = v.stride(1);
+  params.k_stride_t = is_paged ? static_cast<int>(k.stride(2)) : 0;
+  params.v_stride_t = is_paged ? static_cast<int>(v.stride(2)) : 0;
   params.batch_size = qo_segment_lens.size(0);
   params.total_qo_len = q.size(0);
   params.total_kv_len = is_paged ? 0 : k.size(0);

python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh

@@ -131,7 +131,9 @@ struct FwdRunner {
                          int q_stride_h_original = 0,
                          int h_r_original = 0,
                          PackGQAUnpackParams pack_gqa = {},
-                         int num_ctas = 0) {
+                         int num_ctas = 0,
+                         int k_stride_t = 0,
+                         int v_stride_t = 0) {
     cutlass::KernelHardwareInfo hw_info;
     hw_info.device_id = 0;
     hw_info.sm_count = (num_ctas > 0)
@@ -169,8 +171,12 @@ struct FwdRunner {
       int k_stride_head = k_stride_h;
       int v_stride_page = v_stride_n;
       int v_stride_head = v_stride_h;
-      stride_K = make_stride(head_dim_qk, _1{}, k_stride_head, k_stride_page);
-      stride_V = make_stride(_1{}, head_dim_vo, v_stride_head, v_stride_page);
+      // Within-page token stride: head_dim for HND-contiguous caches; NHD
+      // caches ([page, token, head, dim]) pass num_kv_heads * head_dim.
+      int k_stride_token = k_stride_t > 0 ? k_stride_t : head_dim_qk;
+      int v_stride_token = v_stride_t > 0 ? v_stride_t : head_dim_vo;
+      stride_K = make_stride(k_stride_token, _1{}, k_stride_head, k_stride_page);
+      stride_V = make_stride(_1{}, v_stride_token, v_stride_head, v_stride_page);
       auto shape_K = make_shape(KVPageSize, head_dim_qk, num_kv_heads, total_page_num);
       auto shape_V = make_shape(head_dim_vo, KVPageSize, num_kv_heads, total_page_num);
 
@@ -441,7 +447,9 @@ cudaError_t run_fmha_fwd(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn
                          int q_stride_h_original = 0,
                          int h_r_original = 0,
                          PackGQAUnpackParams pack_gqa = {},
-                         int num_ctas = 0) {
+                         int num_ctas = 0,
+                         int k_stride_t = 0,
+                         int v_stride_t = 0) {
   return FwdRunner<DTypeIn, DTypeOut, IdType, TileShapeQK, TileShapePV, ActiveMask,
                    ThreadShape, IsSplitKV, SingleSoftmaxWarpGroup, KVPageSize, kSparseAttnMode>::run(
       workspace_buffer, q, k, v, qo_segment_lens, kv_segment_lens,
@@ -456,7 +464,7 @@ cudaError_t run_fmha_fwd(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn
       maybe_max_score, max_k_tiles,
       kv_block_indexes, kv_block_num,
       pack_factor, q_stride_n_original, q_stride_h_original, h_r_original,
-      pack_gqa, num_ctas);
+      pack_gqa, num_ctas, k_stride_t, v_stride_t);
 }
 
 };  // namespace flashinfer

python/fmha_sm100/cute/README.md

@@ -163,11 +163,15 @@ objects raise an error instead of silently falling back.
 ### Sparse Page Attention
 
 - `q`: `[total_q, Hq, D]`
-- `k`, `v`: `[num_pages, page_size, Hkv, D]`
+- `k`, `v`: `[num_pages, Hkv, page_size, D]` (`kv_layout="HND"`, default) or
+  `[num_pages, page_size, Hkv, D]` (`kv_layout="NHD"`)
 - `page_table`: `[B, max_num_pages_per_seq]`
 - `seqused_k`: optional `[B]`, logical valid KV length per batch
 - `cu_seqlens_q`: `[B + 1]`
 - do not pass `cu_seqlens_k` together with `page_table`
+- `kv_layout`: `"HND"` (default) or `"NHD"`; a standard contiguous NHD cache is
+  consumed zero-copy through strided TMA — inputs violating the alignment
+  contract (last dim strided, unaligned base/strides) fall back to a copy
 
 Use `seqused_k` whenever logical KV length is smaller than the physical page
 capacity. This is the normal way to represent partially used tail pages.
@@ -587,6 +591,7 @@ The decode path is intentionally narrower than the dense sparse path:
 | `page_size` | `128` (must equal `blk_kv`) |
 | Causal | `True` |
 | Batch | `1 ≤ B ≤ 1024` |
+| KV cache layout | `[num_pages, Hkv, 128, 128]` HND (default); `run(..., kv_layout="NHD")` accepts `[num_pages, 128, Hkv, 128]` (contiguous NHD caches are zero-copy) |
 
 `seqused_k` may vary across batch (variable-length decode is the design
 target).  The schedule includes a load-balance heuristic that triggers

python/fmha_sm100/cute/interface.py

@@ -136,6 +136,53 @@ def _prepare_paged_kv_for_tma(k, v, blk_kv: int):
     return k, v
 
 
+def _apply_kv_layout(k, v, kv_layout: str, *, paged: bool):
+    """Normalize K/V to the canonical HND view.
+
+    ``kv_layout="HND"`` expects paged K/V as ``[num_pages, Hkv, page_size, D]``
+    and is a no-op.  ``kv_layout="NHD"`` expects ``[num_pages, page_size, Hkv,
+    D]`` and returns permuted views in the canonical HND shape.  The kernels
+    consume K/V through dynamic strides, so a standard (contiguous) NHD cache
+    is used zero-copy; inputs that violate the kernel alignment contract fall
+    back to a compacting copy in :func:`_kv_kernel_view`.
+    """
+    if kv_layout == "HND":
+        return k, v
+    if kv_layout != "NHD":
+        raise ValueError(f"kv_layout must be 'HND' or 'NHD', got {kv_layout!r}")
+    if not paged:
+        raise ValueError(
+            "kv_layout='NHD' requires paged K/V (pass page_table); dense flat "
+            "K/V [total_k, Hkv, D] is already token-major"
+        )
+    if k.ndim != 4 or v.ndim != 4:
+        raise ValueError(
+            "kv_layout='NHD' requires k/v with shape [num_pages, page_size, Hkv, D]"
+        )
+    return k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
+
+
+_KV_ALIGN_BYTES = 16  # from_dlpack(assumed_align=16) + assume_tensor_aligned (128-bit)
+
+
+def _kv_kernel_view(t: torch.Tensor) -> torch.Tensor:
+    """Pass K/V through when it satisfies the kernel's TMA contract; else copy.
+
+    Contract: contiguous last dim, 16-byte-aligned base pointer, and every
+    non-last stride a 16-byte multiple.  Layout-permuted paged caches (HND
+    views of NHD storage) satisfy this and stay zero-copy; exotic strided or
+    offset views take the compacting copy that the previous unconditional
+    ``.contiguous()`` provided.
+    """
+    if t.stride(-1) != 1 or t.data_ptr() % _KV_ALIGN_BYTES:
+        return t.contiguous()
+    esize = t.element_size()
+    for size, stride in zip(t.shape[:-1], t.stride()[:-1]):
+        if size > 1 and (stride * esize) % _KV_ALIGN_BYTES:
+            return t.contiguous()
+    return t
+
+
 def _validate_cu_seqlens(
     cu_seqlens: torch.Tensor,
     *,
@@ -618,6 +665,7 @@ def sparse_atten_func(
     return_softmax_lse: bool = False,
     page_table: Optional[torch.Tensor] = None,
     seqused_k: Optional[torch.Tensor] = None,
+    kv_layout: str = "HND",
     schedule: Optional[SparseAttentionSchedule] = None,
     usable_SM_count: int = -1,
     qk_dtype: Optional[torch.dtype] = None,
@@ -680,6 +728,13 @@ def sparse_atten_func(
     seqused_k : torch.Tensor, optional
         Shape ``[batch_size]``, dtype int32.  Effective KV length per request
         for paged causal attention.
+    kv_layout : str, optional
+        Paged K/V cache layout: ``"HND"`` (default) for
+        ``[num_pages, Hkv, blk_kv, D]`` or ``"NHD"`` for
+        ``[num_pages, blk_kv, Hkv, D]``.  A standard contiguous NHD cache is
+        consumed zero-copy via strided views; inputs violating the kernel
+        alignment contract fall back to a compacting copy.  Only valid
+        together with ``page_table``.
     schedule : SparseAttentionSchedule, optional
         Prebuilt sparse forward schedule.  If omitted, the schedule is built
         during the call.
@@ -715,6 +770,7 @@ def sparse_atten_func(
         raise ValueError("return_temperature_lse=True requires return_softmax_lse=True")
     partial_dtype = _normalize_partial_dtype(partial_dtype)
     qk_dtype, pv_dtype = _resolve_forward_mma_dtypes(q, k, v, qk_dtype, pv_dtype)
+    k, v = _apply_kv_layout(k, v, kv_layout, paged=page_table is not None)
 
     if cu_seqlens_q is None or cu_seqlens_k is None:
         raise ValueError(
@@ -738,8 +794,8 @@ def sparse_atten_func(
 
     return _sparse_atten_csr_varlen_forward(
         q.contiguous(),
-        k.contiguous(),
-        v.contiguous(),
+        _kv_kernel_view(k),
+        _kv_kernel_view(v),
         k2q_row_ptr.contiguous(),
         k2q_q_indices.contiguous(),
         int(topK),
@@ -1006,6 +1062,7 @@ def sparse_decode_atten_func(
     softmax_scale: Optional[float] = None,
     return_softmax_lse: bool = False,
     schedule: Optional[DecodeAttentionSchedule] = None,
+    kv_layout: str = "HND",
     O_partial: Optional[torch.Tensor] = None,
     LSE_partial: Optional[torch.Tensor] = None,
 ):
@@ -1046,6 +1103,11 @@ def sparse_decode_atten_func(
     O_partial, LSE_partial : torch.Tensor, optional
         Optional split-KV partial workspaces.  Normally owned by
         ``SparseDecodePagedAttentionWrapper``.
+    kv_layout : str, optional
+        Paged K/V cache layout: ``"HND"`` (default) for
+        ``[num_pages, Hkv, blk_kv, 128]`` or ``"NHD"`` for
+        ``[num_pages, blk_kv, Hkv, 128]`` (contiguous NHD caches are consumed
+        zero-copy).
 
     Returns
     -------
@@ -1055,6 +1117,7 @@ def sparse_decode_atten_func(
     """
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** -0.5
+    k, v = _apply_kv_layout(k, v, kv_layout, paged=True)
     batch, head_kv = _validate_sparse_decode_inputs(
         q,
         k,
@@ -1100,8 +1163,8 @@ def sparse_decode_atten_func(
     )
     _call_sparse_decode_forward_sm100_paged_fp8(
         q.contiguous(),
-        k.contiguous(),
-        v.contiguous(),
+        _kv_kernel_view(k),
+        _kv_kernel_view(v),
         None if q2k_indices is None else q2k_indices.contiguous(),
         page_table.contiguous(),
         seqused_k.contiguous(),
@@ -1343,6 +1406,7 @@ def run(
         return_softmax_lse: bool = False,
         out: Optional[torch.Tensor] = None,
         lse: Optional[torch.Tensor] = None,
+        kv_layout: str = "HND",
     ):
         """Launch decode using metadata cached by ``plan``.
 
@@ -1362,6 +1426,10 @@ def run(
             Preallocated BF16 output buffer with shape ``q.shape``.
         lse : torch.Tensor, optional
             Preallocated float32 LSE buffer with shape ``[total_q, Hq]``.
+        kv_layout : str, optional
+            Paged K/V cache layout: ``"HND"`` (default) or ``"NHD"`` for
+            ``[num_pages, blk_kv, Hkv, 128]`` caches (contiguous NHD caches
+            are consumed zero-copy).
 
         Returns
         -------
@@ -1381,8 +1449,11 @@ def run(
                 softmax_scale=softmax_scale, return_softmax_lse=return_softmax_lse,
                 schedule=self.decode_schedule,
                 O_partial=self.O_partial, LSE_partial=self.LSE_partial,
+                kv_layout=kv_layout,
             )
 
+        k, v = _apply_kv_layout(k, v, kv_layout, paged=True)
+        k, v = _kv_kernel_view(k), _kv_kernel_view(v)
         if softmax_scale is None:
             softmax_scale = q.shape[-1] ** -0.5
         if out is None: