[gilby]

What this changes

Flash Next's sparse attention on HIP gets a second kernel, tf_fn_sparse_lanes{6,12}_{f16,bf16} (zig/kernels/hip/flashnext/sparse_lanes.hpp), that does the same fp32 operations in the same order as tf_fn_sparse_parts{6,12} and so gives the same bits: the launcher takes it for 16-bit caches whose head dim is a multiple of 8, and keeps the scalar kernel for fp32 caches. Prompt rows, decode rows and MTP verify rows all go through it, as they did through the scalar kernel.

The scalar kernel spends most of its time on LDS reads (two per multiply-add) and a cache-kind branch per element. The new kernel keeps each row's 32-key tiles, each (head, key) score as a sequential dim-by-dim sum, the per-tile xor-shuffle max and sum trees, the running max/sum update tile by tile and P·V key by key, and changes only where the operands come from:

• a lane holds one key and computes several heads' scores (each staged K value is read once per dim, not once per head), with q read four dims at a time;
• K and V are staged in the cache type with 16-byte loads (the widening to fp32 is exact) and V is read from LDS instead of global memory for every key;
• the cache kind is a template parameter, and heads past the group compute on unused rows instead of branching;
• one tile at a time in 31 KB of LDS, so two blocks fit a CU.

Refs #1597.

Receipt

• Environment: TensorFold main 9e8c368b plus this change, zig build native-hip -Dhip-arch=gfx1151, Ubuntu archive HIP 7.1 (clang 21), AMD Ryzen AI MAX+ 395 with Radeon 8060S gfx1151 (128 GB unified memory). Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6, MTP 3 drafts, --parallel 1 --context 131072 --prompt-cache-gib 0, base and candidate alternating on fresh servers on the same machine.
• Bits: a new device test (hip/kernel_tests/flashnext_sparse.zig) runs both kernels on the same rows and compares the parts, the maxima and sums and the merged rows word for word: bf16 and fp16 caches, 12, 6 and 3 query heads a kv head, head dims 256 and 128, 1, 3, 5 and 16 parts, key counts 0, 1, 7, 31, 32, 33, 63, 64, 65, 97, 300, 650, 2,000 and 2,048, listed and dense rows. It needs the new launcher entry, so it fails to build on main. zig build hip-kernel-test 110 passed, 1 skipped, 0 failed (run directly from the test binary); with two of the kernel's additions swapped, the new test fails.
• Contract: a record for tf_fn_sparse_lanes12_{f16,bf16} and a runner (zig/tests/hip/contract_sparse_lanes.zig) against an fp64 reference: SUMMARY gfx1151 53 kernels: 53 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 130 limit problems (main: 52 kernels, the same 2 known findings and 130 limit problems). The limit problems are the same register differences in existing kernels as on main; the new records pass their live limits. All four chips' records come from zig build hip-kernel-contract-limits -Dhip-relax-branches=true (every new record current on gfx1030, gfx1100, gfx1151 and gfx1201); the live check passes on gfx1151.
• Tokens: token_sha256 of tools/bench_openai.py, tools/bench_concurrent.py and both tools/bench_cache.py --require-cold prompts identical to main; bench_concurrent --alone --serial every cell equal, 0 failed.
• Prompt speed, tools/bench_cache.py --require-cold --tokens 16 on public 7,377- and 29,364-token arrays, cold requests after the first (fresh servers in the order main, this, ..., main): 7,377 tokens 19.51 / 19.52 s -> 12.24 / 12.28 s (-37%), 29,364 tokens 87.22 / 86.96 s -> 53.38 / 53.35 s (-39%). The first request after start: 45.4 / 45.0 -> 38.4 s and 97.7 / 97.2 -> 62.9 s.
• Decode speed, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, median tok/s code / chat: main 49.87 / 43.85 and 50.06 / 44.02, this change 50.57 / 44.50 (unchanged); MTP acceptance unchanged (2.753 / 2.37 tokens a round), as the bits are.
• Kernel time on a cold 7,869-token prompt (rocprofv3 kernel trace): tf_fn_sparse_parts12 9.82 s → tf_fn_sparse_lanes12_bf16 1.96 s, at the same 1,614 launches; the whole GPU window 21.7 -> 13.7 s (profiled on main d377fcc6 before the rebase).
• Tests run: zig build test -Dhip=true --summary all 485/485 steps, 0 failed; python3 tools/zig/lean_check.py 424 on main and change; zig fmt --check and git diff --check clean.
• Not run: gfx1030, gfx1100 and gfx1201 (no card; their fp16 path is the same source, and the device test covers fp16 caches on gfx1151); Metal and CUDA (HIP kernels only).

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
• No precision traded for speed. If the bits change, the description says which sums change and why. (Bits unchanged.)
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[gilby]

Rebased onto current main (cfcbecb). The conflicts were the CHANGELOG, the record list and the kernel and contract registries, all resolved by keeping both sides. The merge took hip_records.zig to 609 lines, so the two sparse-lanes records now live in their own hip_records_sparse.zig and hip_records.zig is back to main's 593.

Receipts on the rebased head (14f04fe), one Strix Halo iGPU (gfx1151), ROCm 7.1 hipcc, zig 0.17.0:

• zig fmt clean, lean_check 424 problems, the same as main.
• hip-kernel-contract -Dhip-arch=gfx1151: SUMMARY gfx1151 58 kernels: 58 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed. The four sparse-lanes checks (nonfinite, repeat, shapes, sums) pass for the bf16 kernel.
• Limit records for tf_fn_sparse_lanes12_f16 and _bf16 are current on all four chips. Other kernels' limit records read stale on this ROCm 7.1 toolchain because main's were re-derived with ROCm 10.1 (test(hip): re-derive every gate-5 record with ROCm 10.1 #1979); I did not touch them.
• Device test: the lanes kernel gives the scalar kernel's bits at every count, part and cache kind. The kernel test binary reports 119 passed, 1 skipped, 0 failed.

Not run: gfx1030, gfx1100 and gfx1201 hardware, and any other family. Still waiting on a review.

[ashhart]

Thanks @gilby. A lane kernel that does the scalar kernel's fp32 operations in the same order, for 16-bit caches with head dims that are multiples of 8, is the right shape, and so is moving the two records to their own file to keep hip_records.zig under 600. The branch conflicts with main again (#1978, #2046 and #2034 touched the HIP contract lists). Before it lands:

1. The kernel contract run on your GPU with replaces against tf_fn_sparse_parts (bytes equal at prompt, decode and MTP-verify row counts), its SUMMARY line, and a planted control that goes red.
2. token_sha256 equal to main on served prompts and drafted decode.
3. The A/B: cold prompt at 8K and 32K and decode at 1 and 4 streams, main against this head, alternating runs.

We have no AMD GPU, so your receipts are what we go on. Then rebase, take it out of draft, and we'll land it.