diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index 024bd9f22..ba4270017 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/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): diff --git a/tests/kernels/test_rdna_gemm.py b/tests/kernels/test_rdna_gemm.py index 77aab2beb..138193d1d 100644 --- a/tests/kernels/test_rdna_gemm.py +++ b/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.