ck_tile: fix 2:4-sparse SWMMAC correctness on gfx1201/RDNA4 (3 bugs) + fail→pass repro

[The-Monk]

0. TL;DR for the maintainer (schung-amd)

Three correctness bugs in ck_tile's SPARSE MmaOpFamily path produce wrong results on gfx1201/RDNA4 (v_swmmac_*_iu4). Root-caused, fixed (334-line patch, two files), and reproduced with a standalone test that fails on current CK and passes with the fix. We found these while building a production-grade 2:4-sparse fp8/int4 GEMM for LLM inference on RDNA4 — §3-4 give that context, which doubles as validation that the fixed path works end-to-end at scale.

1. The three bugs (files: sparse_mma_pipeline.hpp, sparse_transforms.hpp)

Bug 1 — compress_a_impl writes phantom values on <2-nonzero 2:4 groups.
The fallback ADataType nonzero_elems[2] = {a_vec[i*4+2], a_vec[i*4+3]} seeds the compressed pair from fixed input positions rather than true zero. A 2:4 group with 0 or 1 real nonzeros then reconstructs with a phantom value from a wrong lane (observed as register aliasing / UB). Fix: default nonzero_elems to true {0,0} and fill only real survivors — 0/1/2-nonzero groups all reconstruct exactly.

Bug 2 — packed sub-byte (pk_int4_t) group-of-4 scan treats a byte as one element.
The compaction's "nonzero per group of 4" scan operates on bytes, but a pk_int4_t byte holds two 4-bit values; a byte is "nonzero" if either nibble is, so the group-of-4 assumption over-counts and writes out of bounds into nonzero_elems[2]/[3]. Fix: packed-sub-byte-aware nibble counting (guarded with a static_assert that the packed path currently implements only pk_int4_t).

Bug 3 — compress_a_impl emits the two per-group idx metadata fields in an order that mismatches what the hardware reads, requiring a SWAP + XOR-1 to correct.
Within CK's own pk_int4_t packing, the metadata CK writes has the two idx fields swapped relative to which compressed survivor they govern, plus a XOR-1, versus what v_swmmac_i32_iu4 consumes on gfx1201 — so the sparse result comes out wrong (pos(HIGH-nibble survivor, found first) = idx1_field XOR 1, pos(LOW-nibble survivor, found second) = idx0_field XOR 1). Dense iu4 was confirmed correct first (11 seeds) to isolate this to the sparse metadata path. Fix: emit the idx fields with the swap+XOR so the hardware reconstructs the correct positions (Idx0 < Idx1, the documented 2:4 contract, is naturally satisfied).
Precision note (so this isn't over-stated): we verified this is specific to CK's compress/packing convention, not a universal "the hardware swaps." An independently-derived encoder in our own driver produces correct metadata without the swap — and applying the swap to it breaks a working encoding. So the fix corrects CK's own ordering to match the hardware; it is not a claim about v_swmmac_iu4 behavior for all encoders.

Patch: gfx1201-sparse-fixes-b675945.patch (against b675945; survives rebase).

2. The test AMD asked for (fails on current CK, passes with the fix)

sparse_ck_warptile_probe.cpp drives CK's own machinery for the SPARSE MmaOpFamily — it builds Pipeline::AWarpDstrEncoding internally (no manual compression math), fills arbitrary-position 2:4 data (REAL_A_16x128, a genuine int8 2:4-pruned weight tile exported from a Quark sparse-llama checkpoint, density=0.494 exact2frac=0.977 — i.e. ~2.3% of 4-groups have <2 real nonzeros, which is exactly what triggers Bug 1's fallback), runs the SWMMAC GEMM (K=32/64/128, three FragsK tile sizes), and checks against a CPU int64-accumulate reference.

Re-run 2026-07-23, fresh scratch clone (~/ck-3753-scratch, exact commit b67594561, detached HEAD, untouched otherwise) — GPU0 only (HIP_VISIBLE_DEVICES=0, Radeon AI PRO R9700 gfx1201):

• Unpatched ck_tile (commit b67594561 as-is): FAIL — max_abs_err=352 on all three tile shapes (K32/K64/K128), deterministic/reproducible run-to-run (fixed real-weight A tile, seeded B):

  [K32_single_frag] M=16 N=16 K=32 -> max_abs_err=352 FAIL   (C_expected[0]=-813 C_actual[0]=-1165)
  [K64_2frag]       M=16 N=16 K=64 -> max_abs_err=352 FAIL   (C_expected[0]=-446 C_actual[0]=-798)
  [K128_4frag]      M=16 N=16 K=128 -> max_abs_err=352 FAIL  (C_expected[0]=429 C_actual[0]=77)
  === SOME FAILED ===
  

• Patched (git apply gfx1201-sparse-fixes-b675945.patch on the same clone, same commit, recompiled, rerun): PASS — max_abs_err=0 on all three shapes, twice in a row:

  [K32_single_frag] M=16 N=16 K=32 -> max_abs_err=0 PASS
  [K64_2frag]       M=16 N=16 K=64 -> max_abs_err=0 PASS
  [K128_4frag]      M=16 N=16 K=128 -> max_abs_err=0 PASS
  === ALL PASS ===
  

Reproduction (from scratch, no dependency on our fork):

git clone https://github.com/ROCm/composable_kernel ~/ck-3753-scratch && cd ~/ck-3753-scratch && git checkout b675945
# unpatched build/run:
clang++ -x hip --offload-arch=gfx1201 -D__HIP_PLATFORM_AMD__=1 \
  -isystem <rocm-devel>/include -I ~/ck-3753-scratch/include \
  -std=c++17 -O3 -DUSE_REAL_TILE \
  -I <path-to>/sparse24-ck sparse_ck_warptile_probe.cpp -o probe_unpatched \
  -L <rocm-devel>/lib -lamdhip64
HIP_VISIBLE_DEVICES=0 ./probe_unpatched     # -> FAIL, max_abs_err=352
# patched build/run:
git apply gfx1201-sparse-fixes-b675945.patch
clang++ ... (same flags) -o probe_patched
HIP_VISIBLE_DEVICES=0 ./probe_patched       # -> PASS, max_abs_err=0

(We compiled with the ROCm-devel clang toolchain directly — clang++ -x hip --offload-arch=gfx1201; the distro hipcc on this box is a stale 5.7-era wrapper unrelated to CK's target toolchain and will not build ck_tile headers. Any current ROCm/HIP clang works the same way.)

ASAN attempted, not obtained (honest note, not a blocker): we tried to get a device-sanitizer stack trace on the Bug-1/Bug-2 OOB write for extra evidence (-fsanitize=address -fgpu-sanitize). Device ASAN needs an xnack+ target-ID variant; RDNA4 (gfx1201) does not accept an xnack+/xnack- feature suffix at all (clang++: error: invalid target ID 'gfx1201:xnack+') — HIP device-side ASAN is a gfx9/CDNA-class feature (gfx90a/gfx94x with xnack+), not available on RDNA4 consumer/workstation targets in this ROCm toolchain. We did not chase this further since it's not required to reproduce the bug — the max_abs_err numbers above are a complete, deterministic fail→pass.

We'll port this into CK's gtest format for the PR; the standalone repro is included so it's runnable without the full CK test build.

3. Context — why these bugs mattered (the "whole thing")

We hit these building a library-grade 2:4-sparse GEMM for RDNA4 LLM inference (llama.cpp/ggml fork, gfx1201, 2× R9700). The journey, honestly:

• ISA ceiling is real: isolated v_swmmac microbench hits 765 TOP/s fp8-2:4 (2× dense) and 1531 TOP/s int4-2:4 (4×) at ILP≥4 — 88-100% of the R9700 spec.
• A hand-written MMQ-grade 2:4 kernel (cooperative tiles, LDS staging, ILP≥4) reaches ~93% of our native-fp8 WMMA MMQ (which we built — none exists upstream) on the whole model, and is at or ahead of the library on the three dominant GEMM shapes (+1.8% / +3.9% / +32% on gate-up / q-o / down-proj).
• The honest finding worth AMD seeing: the ISA 2×/4× sparsity ceiling does not translate to a compute-bound, well-tiled model GEMM — measured across three quality tiers, the sparsity factor is 2× (isolated ISA) → 1.25× (crude kernel) → ~1.01× (library-grade). Once the kernel is well-tiled, the GEMM is bound by the tiled system (LDS staging, occupancy), not the raw SWMMAC issue rate that sparsity accelerates. This is a real, reproducible microbench-vs-model boundary for RDNA4 structured sparsity.
• The mechanism, made precise (int4-2:4): the "4×" decomposes as 2×(int4 vs int8) × 2×(2:4 sparsity). The first 2× is real and already captured by dense wmma_i32_16x16x32_iu4 — gfx1201's dense int4 engine is already K=32-wide. The second 2× washes: swmmac_i32_16x16x32_iu4 (749,761 GOP/s) ≈ wmma_i32_16x16x32_iu4 (765,061) = 0.98×, instruction-for-instruction, because dense int4 already occupies the K=32 slot. The only sparse tensor edge is the K=64 form (1.77×) — the same "2×-K" mechanism fp8-2:4 already showed washing at library grade. So on high-arithmetic-intensity LLM GEMM shapes (compute-bound), structured 2:4 is a ~1.2× weight-bandwidth win (memory-bound regime only), not the tensor 2×/4×. (Correctness first regardless — the perf question is separate, and now answered.)
• The remaining gap to the library was pinned by direct profiling to two specific, non-magical things — a 2×-slow activation-quantize helper (LDS+sync vs shuffle) and one occupancy-starved small-N shape — both fixed [§4].

4. Closing the gap — result

The 8% was pinned (by direct rocprofv3 head-to-head) to two things, NOT the big GEMMs (we're at/ahead of our own fp8 MMQ there: +1.8% / +3.9% / +32% on gate-up / q-o / down-proj). Fixing them:

build	dense-MMQ (vs our fp8 MMQ)	2:4-MMQ (vs our fp8 MMQ)
capstone baseline	3786 (0.924×)	3823 (0.933×)
+ shuffle-based quantize (kept)	3900 (0.952–0.978×)	3944 (0.962–0.989×)

• Quantize helper rewrite (kept): replaced a 32-thread / LDS+__syncthreads() reduction with the upstream mmq.cuh quantize shape — 128 threads, float4 loads, warp-shuffle max-reduction, zero LDS/sync (halved it, 36µs→~18µs). Closed ~40% of the remaining gap; both kernels now high-90s% of our native-fp8 kernel (0.989× same-session).
• The N=1024 k/v small-shape (honest negative): both candidate levers — LDS bank-conflict padding and coarser K-per-sync — were implemented, correctness-verified, and measured worse (padding −2.5%; coarser-K a monotonic regression as it forced dropping double-buffering into an occupancy/LDS cliff). Both reverted. The outlier remains open, wants a different lever (shape-adaptive smaller tile without touching sync granularity).

Net: the hand-written 2:4 kernel now reaches ~96–99% of our native-fp8 WMMA MMQ (which we built — none exists upstream), is ahead on the FLOP-dominant GEMMs, and carries the measured ~1% sparsity edge over a comparable dense kernel — a library-competitive RDNA4 2:4 GEMM on top of the correctness fixes.

5. A related RDNA4 iu4 quirk (same class, different instruction — FYI, not part of the fix)

The Bug-3 element-ordering quirk on swmmac_iu4 is not isolated: an independent kernel-authoring effort on this box hit the same class on the dense v_dot8_i32_iu4/sudot8 path — "dots mismatched element pairs." RDNA4's iu4 instruction family appears to carry undocumented element-ordering conventions in both the sparse (SWMMAC metadata) and dense (dot8 operand pairing) paths. Documenting these in the ISA/CK would save the next implementer the multi-day reverse-engineering we did. Happy to write up the dense one too if useful.

6. Method (for reproducibility)

All numbers: gfx1201 (R9700), ROCm 7.14, GPU-isolated, warm, medians ≥3. Correctness gated by execution (CPU-reference compare), never inspection. The full capability-optimizer method (measure → grade vs the published peak → drive the lever) is what surfaced both the bugs and the microbench-vs-model boundary.