HIP Flash Next: WMMA sparse attention for prompt chunks on gfx11 #1597

On AMD Ryzen AI MAX+ 395, Radeon 8060S gfx1151 (40 CUs, 128 GB unified memory, measured 225.6 GB/s read), Ubuntu 26.04, archive HIP 7.1 (clang 21), TensorFold 7c068895, TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP.

A cold prompt prefills at about 240 tok/s, flat from 8K to 262K tokens. A rocprofv3 kernel trace of one 7.9K-token prompt puts 33% of the window in tf_fn_sparse_parts12 (zig/kernels/hip/flashnext/sparse.hpp): 9.8 s over 1,669 calls, about 7 ms per 64-row slice once the indexer's selection is full. The kernel is scalar fp32: one block per (kv head, query row, part), each staging its own row's selected keys, so neighbouring query rows that select mostly the same blocks share nothing.

Proposal

A prompt-only kernel on gfx11 WMMA: 4 query rows x the kv head's 12 query heads as 48 WMMA columns, walking the union of those rows' selected blocks, with per-row block masks, the causal end per row, fp32 online softmax and the existing parts/merge layout. Decode keeps the scalar kernel. It changes the sums (bf16 Q.K and P.V fragments with fp32 accumulation instead of fp32 scalar), so it needs the agreed arithmetic gate.

A prototype behind an opt-in switch: the kernel matches an fp64 reference within 2.1e-3 absolute on served shapes (the scalar kernel 1.7e-7); prefill 231/260/241 -> 326/411/374 tok/s at 8K/32K/131K with needles found and decode unchanged; the sparse kernel's time per 8K prompt drops from 9.81 s to 0.73 s.

Gate to agree

• FP64 error bound and teacher-forced top-token comparison against the scalar kernel on served shapes.
• Kernel contract registration on the HIP runner (gates 1-5) and limits for every served RDNA chip.
• Drafted equals plain, concurrent equals solo, resumed equals fresh, on the WMMA path.