[Kernel] Fix single-accumulator RDNA3 GEMM tiles

kernels/gemm/rdna3_f16_gemm.py

@@ -353,7 +353,14 @@ def _do_compute_rk(accs_in, rk, buf_offset):
 
             results = yield list(s_accs)
 
-        accs = list(results[:n_acc])
+        # A one-atom tile carries one scf.for result. FlyDSL intentionally
+        # unwraps a single result to its value, while multiple results remain a
+        # sliceable sequence. Keep the kernel-side accumulator container stable
+        # in both cases so 16x16 tiles can reach the epilogue.
+        if const_expr(n_acc == 1):
+            accs = [results]
+        else:
+            accs = list(results[:n_acc])
 
         last_read_off = ((num_k_tiles - 1) % 2) * c_lds_buf_stride
         for rk in range_constexpr(reg_k):

tests/kernels/test_rdna_gemm.py

@@ -109,6 +109,37 @@ def test_f16_gemm_correctness(M, N, K, in_dtype, out_dtype):
     assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05)
 
 
+def test_f16_gemm_single_accumulator_loop_result():
+    """A 16x16 tile keeps one WMMA accumulator across the pipelined K loop."""
+    _requires_rdna3()
+
+    M, N, K = 16, 16, 64
+    torch.manual_seed(42)
+    launch_fn, block_m, block_n, block_k = _create_wmma_gemm_module_gfx11(
+        M,
+        N,
+        K,
+        in_dtype="bf16",
+        out_dtype="bf16",
+        reg_m=1,
+        reg_n=1,
+        reg_k=2,
+        waves_m=1,
+        waves_n=1,
+    )
+    assert (block_m, block_n, block_k) == (16, 16, 32)
+
+    A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 0.1
+    B_T = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1
+    C = torch.zeros(M, N, dtype=torch.bfloat16, device="cuda")
+
+    launch_fn(C, A, B_T, torch.cuda.current_stream())
+    torch.cuda.synchronize()
+
+    C_ref = A.float() @ B_T.float().T
+    assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05)
+
+
 def test_f16_gemm_stochastic_rounding():
     """BF16 GEMM with the stochastic-rounding epilogue: bounded, seed-varying, reproducible.