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

test: separate TMEM regions for A and C in TS MMA

Browse Source
This commit is contained in:
biondizzle
2026-05-28 13:32:22 +00:00
parent 37a502e476
commit c05cc1ac93
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/unit/test_mma_ts.cu
View File

@@ -42,17 +42,20 @@ test_mma_ts(float* o_out)
}
__syncthreads();
// TMEM alloc — 32 columns (16 for A, 16 for C)
if (wid == 0) tmem_alloc(__cvta_generic_to_shared(sTmemBase), 32);
// TMEM alloc — 64 columns (16 for A at offset 0, 16 for C at offset 32)
// TMEM alloc requires power of 2, minimum 32
if (wid == 0) tmem_alloc(__cvta_generic_to_shared(sTmemBase), 64);
__syncthreads();
uint32_t tb = *sTmemBase;
uint32_t tb_a = tb; // A starts at column 0
uint32_t tb_c = tb + 32; // C starts at column 32
// Write A = all 1.0 into TMEM columns 0-15 (128 rows × 16 columns)
if (wid == 0) {
for (int col = 0; col < 16; col++) {
// Each column: 128 FP32. Lane i writes positions i*4..i*4+3
float v0 = 1.0f, v1 = 1.0f, v2 = 1.0f, v3 = 1.0f;
tmem_store(tb + col, f32_to_u32(v0), f32_to_u32(v1), f32_to_u32(v2), f32_to_u32(v3));
tmem_store(tb_a + col, f32_to_u32(v0), f32_to_u32(v1), f32_to_u32(v2), f32_to_u32(v3));
}
tmem_fence_store();
}
@@ -63,7 +66,7 @@ test_mma_ts(float* o_out)
float check = 0.0f;
for (int col = 0; col < 16; col++) {
uint32_t u0, u1, u2, u3;
tmem_load(tb + col, u0, u1, u2, u3);
tmem_load(tb_a + col, u0, u1, u2, u3);
tmem_fence_load();
check += u32_to_f32(u0);
}
@@ -82,7 +85,7 @@ test_mma_ts(float* o_out)
printf("Before MMA: tb=%u, dv=%lu, idesc=%u, tid=%d\n", tb, dv, idesc, tid);
if (tid == 0) {
umma_ts_f16(tb, tb, dv, idesc, false);
umma_ts_f16(tb_c, tb_a, dv, idesc, false);
}
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
@@ -94,7 +97,7 @@ test_mma_ts(float* o_out)
float c_vals[16];
for (int n = 0; n < 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.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_c + n*8));
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) c_vals[n*8+c] = tmp[c];
}
@@ -109,7 +112,7 @@ test_mma_ts(float* o_out)
}
}
if (wid == 0) tmem_dealloc(tb, 32);
if (wid == 0) tmem_dealloc(tb, 64);
}
int main() {