From 5544d3a0a483836f0bc6a0d053797666e9a54d92 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 04:48:26 +0000 Subject: [PATCH] fix: TMEM reads must be outside my_row_active (warp-collective) --- .../fmha_6warp_tma_multirow_multitile.cuh | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh b/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh index ff30b6bc..efcb7090 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh @@ -277,16 +277,16 @@ fmha_6warp_tma_multirow_multitile_kernel(FmhaTmaMultiRowMultiTileParams params) float tile_max = my_row_active ? sTileRowMax[my_row] : 0.0f; float tile_sum = my_row_active ? sTileRowSum[my_row] : 0.0f; - // Read O_tile from TMEM (same pattern as multirow epilogue) + // Read O_tile from TMEM — MUST be outside my_row_active guard (warp-collective!) float o_tile_buf[HD]; - if (my_row_active) { - for (int n = 0; n < N_NSUB * 2; n++) { - float tmp[8]; - asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];" - : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]), - "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7]) - : "r"(tb + n * 8)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); + for (int n = 0; n < N_NSUB * 2; n++) { + float tmp[8]; + asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];" + : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]), + "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7]) + : "r"(tb + n * 8)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + if (my_row_active) { for (int c = 0; c < 8; c++) { int d = n * 8 + c; if (d < HD) o_tile_buf[d] = tmp[c];