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

Add softmax read from TMEM

Browse Source
This commit is contained in:
biondizzle
2026-05-28 14:40:35 +00:00
parent 38b0ff0bf8
commit a6f7ef7c45
1 changed files with 23 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
tests/unit/test_minimal_pv.cu
View File

@@ -35,8 +35,29 @@ test_minimal_pv()
for (int i = tid; i < TILE_SZ; i += 128) { sQ0[i] = 0; sK0[i] = 0; }
__syncthreads();
// Set s_p_vals = all 0.5
for (int j = tid; j < 16; j += 128) s_p_vals[j] = 0.5f;
// Softmax: read S from TMEM, compute P, write to s_p_vals
if (wid == 0) {
float s_vals[16], row_max = -INFINITY;
// Read first 16 columns of S (for K-tile 0)
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));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) { s_vals[c] = tmp[c]; row_max = fmaxf(row_max, tmp[c]); }
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+8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) { s_vals[8+c] = tmp[c]; row_max = fmaxf(row_max, tmp[c]); }
row_max = wmax(row_max);
float row_sum = 0.0f;
if (lane == 0) for (int j=0;j<16;j++) { s_vals[j] = expf(s_vals[j] - row_max); row_sum += s_vals[j]; }
row_sum = wsum(row_sum);
if (lane == 0) for (int j=0;j<16;j++) { s_vals[j] /= row_sum; s_p_vals[j] = s_vals[j]; }
}
__syncthreads();
// Fill sPk: (128, 16) canonical, row 0 = s_p_vals