biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: TMEM reads must be outside my_row_active (warp-collective)

Browse Source
This commit is contained in:
biondizzle
2026-05-30 04:48:26 +00:00
parent 1dca8d8cfa
commit 5544d3a0a4
1 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh
View File

@@ -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];