From 793cf1112d93e67cab8d5ae57c9779bcc6dc8303 Mon Sep 17 00:00:00 2001 From: AlumKal Date: Wed, 12 Aug 2026 20:14:25 +0800 Subject: [PATCH] feat: support NHD paged KV layout via kv_layout="NHD" Both stacks previously accepted only HND paged caches [num_pages, Hkv, page_size, D]. Add kv_layout="HND"|"NHD" to fmha_sm100, sparse_atten_func, sparse_decode_atten_func, and SparseDecodePagedAttentionWrapper.run; NHD caches [num_pages, page_size, Hkv, D] are consumed zero-copy as strided views. - cute: normalize layout at the interface; replace unconditional k/v.contiguous() with _kv_kernel_view, which forwards tensors that satisfy the TMA contract (contiguous last dim, 16B-aligned base and outer strides) and falls back to the historical compacting copy otherwise. - csrc: plumb the within-page token stride (k_stride_t/v_stride_t) from torch strides through params/jinja into stride_K/stride_V, replacing the hardcoded head_dim token stride. Default 0 keeps stale JIT caches HND-correct; clear the JIT cache dir to pick up NHD support. - docs: fix cute README paged-layout description (claimed NHD while the code required HND) and document kv_layout, including the decode wrapper table row. Verified bitwise HND==NHD on both stacks (cute prefill bf16/fp8, cute dense decode fp8, csrc sparse prefill fp8/bf16, csrc sparse decode fp8); full cute test suite green (1169 passed, 141 pre-existing skips). Perf delta within noise: <=0.07% at 32k bs1 sparse prefill, +0.00% at bs256 32k sparse decode with 16k shared prefix. --- python/fmha_sm100/api.py | 19 ++++- python/fmha_sm100/csrc/fmha_sm100_inst.jinja | 3 +- python/fmha_sm100/csrc/fmha_sm100_params.h | 4 + .../csrc/fmha_sm100_variant_run.cu.jinja | 2 + .../csrc/include/fmha_cutlass_sm100.cuh | 18 +++-- python/fmha_sm100/cute/README.md | 7 +- python/fmha_sm100/cute/interface.py | 79 ++++++++++++++++++- 7 files changed, 120 insertions(+), 12 deletions(-) diff --git a/python/fmha_sm100/api.py b/python/fmha_sm100/api.py index b02f747..cd75d75 100644 --- a/python/fmha_sm100/api.py +++ b/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) diff --git a/python/fmha_sm100/csrc/fmha_sm100_inst.jinja b/python/fmha_sm100/csrc/fmha_sm100_inst.jinja index 6949035..ff211b5 100644 --- a/python/fmha_sm100/csrc/fmha_sm100_inst.jinja +++ b/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); } diff --git a/python/fmha_sm100/csrc/fmha_sm100_params.h b/python/fmha_sm100/csrc/fmha_sm100_params.h index 878a90c..c19dea3 100644 --- a/python/fmha_sm100/csrc/fmha_sm100_params.h +++ b/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; diff --git a/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja b/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja index 7e21585..55b2313 100644 --- a/python/fmha_sm100/csrc/fmha_sm100_variant_run.cu.jinja +++ b/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(k.stride(2)) : 0; + params.v_stride_t = is_paged ? static_cast(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); diff --git a/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh b/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh index b11abe6..b772a4a 100644 --- a/python/fmha_sm100/csrc/include/fmha_cutlass_sm100.cuh +++ b/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::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 diff --git a/python/fmha_sm100/cute/README.md b/python/fmha_sm100/cute/README.md index 3e2b846..c1af0e5 100644 --- a/python/fmha_sm100/cute/README.md +++ b/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 diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py index d72b17a..60290fa 100644 --- a/python/fmha_sm100/cute/interface.py +++ b/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: