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

[alumkal]

Motivation

Serving engines differ on paged-KV cache layout: HND [num_pages, Hkv, page_size, D] vs NHD [num_pages, page_size, Hkv, D] (FlashInfer naming). MSA previously accepted only HND, forcing NHD-based engines to repack or transpose. Both MSA stacks already consume K/V through runtime-stride TMA, so NHD can be served zero-copy — the layout assumption was only skin-deep.

What changed

New keyword kv_layout="HND" | "NHD" (default "HND", fully backward compatible) on all runtime entry points:

Entry point	Stack
fmha_sm100(...)	csrc JIT
sparse_atten_func(...)	CuTe-DSL sparse prefill (BF16/FP8)
sparse_decode_atten_func(...) / SparseDecodePagedAttentionWrapper.run(...)	CuTe-DSL paged FP8 decode

fmha_sm100_plan is unchanged — plans never see KV tensors.

CuTe-DSL (cute/interface.py)

• _apply_kv_layout normalizes NHD input to a permuted HND-shaped view; kernels consume it through their existing dynamic-stride layouts — no kernel changes.
• Replaced the unconditional k/v.contiguous() (which silently repacked any strided view) with _kv_kernel_view: forwards tensors satisfying the TMA contract (contiguous last dim, 16 B-aligned base pointer and outer strides), otherwise falls back to the historical compacting copy.

csrc JIT (params.h, both jinjas, fmha_cutlass_sm100.cuh)

• The paged branch hardcoded the within-page token stride to head_dim. Now k_stride_t/v_stride_t are taken from k.stride(2)/v.stride(2) and land in stride_K/stride_V; the stride modes were already runtime ints, so this is pure plumbing. Default 0 selects the old head_dim behavior.

Docs (cute/README.md)

• Fixed a pre-existing error: the Sparse Page Attention section documented NHD while the code required HND. Both layouts and kv_layout are now documented, including the decode-wrapper table.

Out of scope

• NVFP4 path: the packed FP4 data is stride-driven and would work the same way, but the 128×4 scale-row order bakes in HND flattening (page, head, token) — a separate contract, left unchanged.
• The CuTe-DSL paged FP8 decode kernel's sparse mode (q2k_indices) is unimplemented upstream (sparse_decode_atten_func raises NotImplementedError for it; sparse decode is served by the csrc sparse path instead). kv_layout is still forwarded on that branch so NHD works as soon as the q2k gather kernel lands.

Verification

• Bitwise HND == NHD on an SM100-class GPU (cc10.3) across six cases: cute sparse prefill (BF16, FP8), cute dense decode (FP8), csrc sparse prefill (FP8, BF16), csrc sparse decode (FP8, shared-prefix page table).
• Full cute test suite: 1169 passed, 141 skipped (pre-existing skips). csrc proxy-KV smoke test passes.
• Guard behavior: NHD views verified zero-copy (pointer/stride preserved); non-contiguous misaligned and odd-stride views are repaired by copy exactly as before; NHD without kv_indices raises a clear ValueError.

Performance (HND vs NHD, CUDA events, cold L2, median)

Case (h=64/4, topk=16, page=128)	HND µs	NHD µs	delta
cute sparse prefill 32k bs1 bf16	4720.8	4724.2	+0.07%
cute sparse prefill 32k bs1 fp8	3740.6	3742.1	+0.04%
fmha_sm100 sparse prefill 32k bs1 bf16	4812.8	4815.3	+0.05%
fmha_sm100 sparse prefill 32k bs1 fp8	4210.2	4211.8	+0.04%
csrc sparse decode 32k (16k shared prefix) bs256 fp8	706.5	706.6	+0.00%

NHD is effectively free: TMA box rows are D-contiguous in either layout; only the inter-row stride changes.