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

test: Layout D TMEM read using 32x32b.x8 format, 4 warps

Browse Source
This commit is contained in:
biondizzle
2026-05-28 10:07:15 +00:00
parent 68d1a7920c
commit 9e98c067ab
1 changed files with 45 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
tests/unit/test_umma_qk.cu
View File

@@ -1,8 +1,7 @@
/**
* UMMA QK GEMM Test (HD=16, SK=128) — debugging 4× output factor.
* MMA produces S[0,0] = 4× scalar reference. Investigating why.
*
* Working: 1 tmem_load at column 0. Multi-column load crashes.
* UMMA QK GEMM Test (HD=16, SK=128)
* Uses Layout D TMEM read (32x32b.x8, 4 warps) to correctly read MMA output.
* Debugging the 4× scaling factor observed with 16x256b read.
*/
#include <cuda_runtime.h>
@@ -29,73 +28,57 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k,
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
bf16_t* sK = sQ + 128 * 16 + 4096; // 8KB padding after Q
bf16_t* sK = sQ + 128 * 16 + 4096;
float* sQ_row = (float*)(sK + 128 * 16);
for (int d = tid; d < 16; d += NTHREADS) sQ_row[d] = bf16_to_f32(q[d]);
// TMEM alloc (128 cols for MMA output N=128)
// TMEM alloc (128 cols)
if (wid == 0) {
uint32_t sp = __cvta_generic_to_shared(sTmemBase);
tmem_alloc(sp, 128);
tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128);
}
__syncthreads();
uint32_t tmem_base = *sTmemBase;
uint32_t tb = *sTmemBase;
// Load Q and K into SMEM
// Load Q and K
write_q_to_smem<16>(sQ, q);
write_k_to_smem<128, 16>(sK, k);
__syncthreads();
// Verify SMEM data
if (tid == 0) {
for (int d = 0; d < 16; d++) {
int ck = d / 8, lc = d % 8;
s_out[160 + d] = bf16_to_f32(sQ[ck * 16 * 64 + lc]);
s_out[176 + d] = bf16_to_f32(sK[ck * 16 * 64 + lc]);
}
float dot = 0.0f;
for (int d = 0; d < 16; d++) {
int ck = d / 8, lc = d % 8;
dot += bf16_to_f32(sQ[ck * 16 * 64 + lc]) * bf16_to_f32(sK[ck * 16 * 64 + lc]);
}
s_out[192] = dot * scale;
}
__syncthreads();
// Descriptors
uint32_t sQ_smem = __cvta_generic_to_shared(sQ);
uint32_t sK_smem = __cvta_generic_to_shared(sK);
uint64_t desc_q = make_umma_desc_kmajor_none(sQ_smem, 64); // M=64
uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 64); // M=64
// Test with different N values to understand scaling
// N=8 → n_dim=1
uint32_t idesc = make_idesc(64, 128); // Try M=64
if (tid == 0) {
memcpy(&s_out[128], &desc_q, 8);
memcpy(&s_out[130], &desc_k, 8);
memcpy(&s_out[132], &idesc, 4);
s_out[133] = (float)sQ_smem;
s_out[134] = (float)sK_smem;
s_out[135] = (float)tmem_base;
}
__syncthreads();
uint64_t desc_q = make_umma_desc_kmajor_none(sQ_smem, 128);
uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 128);
uint32_t idesc = make_idesc(128, 128);
// MMA
if (tid == 0) {
umma_ss_f16(tmem_base, desc_q, desc_k, idesc, false);
umma_ss_f16(tb, desc_q, desc_k, idesc, false);
}
__syncwarp();
if (wid == 0 && lane == 0) tmem_fence_store();
__syncthreads();
// Read S[0,0] from TMEM column 0
if (wid == 0) {
uint32_t u0, u1, u2, u3;
tmem_load(tmem_base, u0, u1, u2, u3);
if (lane == 0) s_out[0] = u32_to_f32(u0);
// Read from TMEM using Layout D: 32x32b.x8 format
// Each warp reads 32 rows × 8 columns
// Warp 0: rows 0-31, warp 1: rows 32-63, warp 2: rows 64-95, warp 3: rows 96-127
if (wid < 4) {
const int row = wid * 32;
const int col = 0;
const int addr = tb + (row << 16) + col;
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"(addr));
asm volatile("tcgen05.wait::ld.sync.aligned;");
// Lane 0 gets S[wid*32 + 0, 0..7] in tmp[0..7]
if (lane == 0) {
for (int c = 0; c < 8; c++) {
s_out[wid * 8 + c] = tmp[c];
}
}
}
__syncthreads();
@@ -110,11 +93,11 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k,
}
__syncthreads();
if (wid == 0) tmem_dealloc(tmem_base, 128);
if (wid == 0) tmem_dealloc(tb, 128);
}
int main() {
printf("=== UMMA QK GEMM Test (HD=16) ===\n");
printf("=== UMMA QK GEMM Test (Layout D read) ===\n");
const int HD = 16, SK = 128;
const float SCALE = 1.0f / sqrtf((float)HD);
@@ -143,15 +126,18 @@ int main() {
cudaMemcpy(h_s_out, d_s_out, 256*sizeof(float), cudaMemcpyDeviceToHost);
cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost);
printf("S[0,0] MMA: %.6f\n", h_s_out[0]);
printf("S[0,0] SMEM: %.6f\n", h_s_out[192]);
printf("S[0,0] scalar: %.6f\n", h_s_scalar[0]);
printf("Ratio: %.4f\n", h_s_scalar[0] != 0 ? h_s_out[0]/h_s_scalar[0] : 0);
printf("Q SMEM: "); for(int d=0;d<4;d++) printf("%.4f ",h_s_out[160+d]); printf("...\n");
printf("Q orig: "); for(int d=0;d<4;d++) printf("%.4f ",bf16_to_f32_host(h_q[d])); printf("...\n");
printf("K[0] SMEM: "); for(int d=0;d<4;d++) printf("%.4f ",h_s_out[176+d]); printf("...\n");
printf("K[0] orig: "); for(int d=0;d<4;d++) printf("%.4f ",bf16_to_f32_host(h_k[d])); printf("...\n");
// Print S[0,0] from each warp's perspective
printf("S from TMEM (Layout D, 32x32b.x8):\n");
for (int w = 0; w < 4; w++) {
printf(" Warp %d (rows %d-%d): ", w, w*32, w*32+31);
for (int c = 0; c < 8; c++) printf("%.4f ", h_s_out[w*8+c]);
printf("\n");
}
printf("S[0,0..7] scalar: ");
for (int c = 0; c < 8; c++) printf("%.4f ", h_s_scalar[c]);
printf("\n");
printf("Ratio warp0/col0 vs scalar[0]: %.4f\n",
h_s_scalar[0] != 0 ? h_s_out[0] / h_s_scalar[0] : 0);
cudaFree(d_q); cudaFree(d_k); cudaFree(d_s_out); cudaFree(d_s_scalar);
free(h_q); free(h_k); free(h_s_out); free(h_s_scalar);