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

test: write O to tb (overwriting P), same as isolated test

Browse Source
This commit is contained in:
biondizzle
2026-05-28 13:46:04 +00:00
parent a944f90040
commit 1c14ada386
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/unit/test_fmha_ts_hd16.cu
View File

@@ -119,11 +119,11 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
__syncthreads();
*/
// Zero O TMEM region
// Zero O TMEM region at tb (overwriting P — same as isolated test)
if (wid == 0) {
for (int n = 0; n < TMEM_O / 8; n++) {
for (int n = 0; n < 2; n++) { // 16 cols / 8 = 2 iterations
float z0=0,z1=0,z2=0,z3=0,z4=0,z5=0,z6=0,z7=0;
asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0],{%1,%2,%3,%4,%5,%6,%7,%8};" :: "r"(tb_o+n*8),"f"(z0),"f"(z1),"f"(z2),"f"(z3),"f"(z4),"f"(z5),"f"(z6),"f"(z7));
asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0],{%1,%2,%3,%4,%5,%6,%7,%8};" :: "r"(tb+n*8),"f"(z0),"f"(z1),"f"(z2),"f"(z3),"f"(z4),"f"(z5),"f"(z6),"f"(z7));
}
tmem_fence_store();
}
@@ -139,7 +139,7 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV_base), 16);
uint32_t tmem_a = tb; // P's first 16 columns
if (tid == 0) umma_ts_f16(tb_o, tmem_a, dv, idesc_pv, false); // no accumulate for first tile
if (tid == 0) umma_ts_f16(tb, tmem_a, dv, idesc_pv, false); // Write C to tb (overwriting P)
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
}
@@ -151,7 +151,7 @@ test_fmha_ts(const bf16_t* q, const bf16_t* k, const bf16_t* v,
float o_vals[HD];
for (int n = 0; n < HD / 8; 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_o + n*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 (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c] * 2.0f; // Undo MMA 0.5 scale
}