[Kernel] Fix single-accumulator RDNA3 GEMM tiles

[skyguan92]

Summary

• preserve the RDNA3 GEMM accumulator container when an scf.for has exactly one carried result
• add a native gfx11 regression for the resulting 16x16x32 one-atom tile

Problem

On current main (3df2d7c70059b53f352d83226f9ee8edeceecb33), the public tile parameters admit (reg_m, reg_n, reg_k, waves_m, waves_n)=(1,1,2,1,1). That produces one WMMA accumulator. FlyDSL returns the single post-loop result as the carried vector value, but the kernel always slices it as a sequence:

TypeError: unsupported index type: <class 'slice'>
  rdna3_f16_gemm.py:356: accs = list(results[:n_acc])

The fix wraps that one result before the epilogue and leaves the existing multi-result path unchanged.

Validation

Tested on a Radeon PRO W7900, native gfx1100, ROCm/HIP 7.2, PyTorch 2.9.1+gitff65f5b, with HSA_OVERRIDE_GFX_VERSION unset.

• new focused regression: 1 passed
• related RDNA F16/BF16 suite: 23 passed, 7 deselected
• Ruff 0.16.2 format check and lint: passed
• BF16 correctness against FP32 torch accumulation passed for the five M16 DeepSeek-V4-aligned (K,N) rows: (4096,1024), (1024,4096), (4096,512), (4096,2048), and (2048,4096)
• generated ISA for M16/N512/K4096 is wave32, contains 8 static v_wmma_f32_16x16x16_bf16 instructions, contains no scalar FMA fallback, and reports zero private segment/scratch

For bounded performance reconnaissance, the enabled 16x16x64 tile measured 17.735 us hot kernel p50 at M16/N512/K4096 versus 19.059 us for the previously buildable 16x32x64 control (200 captured launches, 4 replays/round, 31 rounds). It is not uniformly faster across all projection shapes, so this PR deliberately changes no default tile, autotune heuristic, or dispatch policy.

This only fixes a legal RDNA3 kernel configuration. It does not add M<16 tail handling or claim an end-to-end model speedup.

[skyguan92]

CI triage: the only failing check looks unrelated to this RDNA3 GEMM change and non-deterministic.

• The Navi job completed with 1,663 passed / 1 failed / 1,586 skipped. Its sole failure was the untouched tests/kernels/test_layernorm.py::test_layernorm_smoothquant case (M=64, N=2000, float32; max quant diff 4, while the scale diff remained within tolerance).
• This PR changes only kernels/gemm/rdna3_f16_gemm.py and tests/kernels/test_rdna_gemm.py.
• tests/kernels/test_layernorm.py has the same Git blob (fe7022cd9b08f9ce70bb2eaa2b1aa985781b963f) at this PR base (3df2d7c7) and later main (2d65dea1). On that later main, the Navi job passed the same SmoothQuant test and the full suite.

I attempted to rerun only the failed job, but GitHub requires repository admin rights for that action. Could a maintainer rerun the failed Navi job? I do not think a code change is justified unless the failure reproduces.

[skyguan92]

Current-main follow-up; no code change or rerun request in this update.

• The exact stable patch (79c2ecf9...ac01) applies cleanly to current main b8ed73fe and remains two files, +39/-1.
• On a physical W7900 with native gfx1100, HIP 7.2, and HSA_OVERRIDE_GFX_VERSION absent, using the verified current-main FlyDSL wheel, unmodified main reproduces the original TypeError: unsupported index type: <class 'slice'>. Main + this patch passes the 16x16x64 BF16 case (16x16x32 tile, max abs error 4.8749e-4) and the related non-large RDNA suite (19 passed / 5 skipped / 6 deselected).
• The old failing SmoothQuant row (M=64, N=2000, f32) passes 20/20 independent processes on the same card; every run reports max quant diff 1 and max scale diff 1.12e-08 (tolerance 0.001). test_layernorm.py remains exact blob fe7022cd...63f.
• Current-main ROCm 7.14 Navi job 93801446435 also passes test_layernorm_smoothquant. Its full pytest job is not green: the two failures are unrelated 1-GiB OOMs in RMSNorm-large and Softmax with zero device memory free (1642 passed / 2324 skipped / 2 failed).

Checksums, source identity, commands, and raw-result summary: W7900D evidence.

This strengthens the existing unrelated/non-deterministic classification for the old Navi failure and confirms the fix is still needed on current main. I am leaving the branch and LayerNorm tolerance unchanged and will wait for maintainer review.

[cschenjunlin]

Current-main follow-up; no code change or rerun request in this update.

• The exact stable patch (79c2ecf9...ac01) applies cleanly to current main b8ed73fe and remains two files, +39/-1.
• On a physical W7900 with native gfx1100, HIP 7.2, and HSA_OVERRIDE_GFX_VERSION absent, using the verified current-main FlyDSL wheel, unmodified main reproduces the original TypeError: unsupported index type: <class 'slice'>. Main + this patch passes the 16x16x64 BF16 case (16x16x32 tile, max abs error 4.8749e-4) and the related non-large RDNA suite (19 passed / 5 skipped / 6 deselected).
• The old failing SmoothQuant row (M=64, N=2000, f32) passes 20/20 independent processes on the same card; every run reports max quant diff 1 and max scale diff 1.12e-08 (tolerance 0.001). test_layernorm.py remains exact blob fe7022cd...63f.
• Current-main ROCm 7.14 Navi job 93801446435 also passes test_layernorm_smoothquant. Its full pytest job is not green: the two failures are unrelated 1-GiB OOMs in RMSNorm-large and Softmax with zero device memory free (1642 passed / 2324 skipped / 2 failed).

Checksums, source identity, commands, and raw-result summary: W7900D evidence.

This strengthens the existing unrelated/non-deterministic classification for the old Navi failure and confirms the fix is still needed on current main. I am leaving the branch and LayerNorm tolerance unchanged and will wait for maintainer review.

For the SmoothQuant failure, I’ve already identified the root cause and am working on a fix. I’ll create a separate PR once the fix is validated. Since the SmoothQuant test is currently skipped on the latest main, it should be possible to merge the latest main and rerun CI for this PR.