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

test: clean HD=64, 1 K-tile only, verify SMEM writes + compare vs scalar

Browse Source
This commit is contained in:
biondizzle
2026-05-28 12:04:54 +00:00
parent e587e26b06
commit afb18caf2d
1 changed files with 55 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

200
tests/unit/test_umma_qk_hd64.cu
View File

@@ -1,7 +1,7 @@
/**
* UMMA QK GEMM Test — HD=64 (4 K-tiles), SK=128
* Multi-K-tile accumulate: call MMA 4× with accumulate=true
* Each K-tile: 16 BF16 columns, separate descriptor
* UMMA QK GEMM Test — HD=64, SK=128, 1 K-tile (columns 0-15 only)
* Verifies the K-tile descriptor offset is correct by comparing
* against the HD=16 result (which is proven correct).
*/
#include <cuda_runtime.h>
@@ -18,172 +18,88 @@ using namespace dsv4::kernels::attention;
static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; }
// MMA K-tile = 16 BF16
constexpr int MMA_K = 16;
constexpr int N_WARPS = 4;
constexpr int BLOCK_M = 128;
__global__ void __launch_bounds__(N_WARPS * 32)
test_umma_qk_hd64(const bf16_t* q, const bf16_t* k,
__global__ void __launch_bounds__(128)
test_umma_qk_hd64_1ktile(const bf16_t* q, const bf16_t* k,
float* s_out, float* s_scalar, float scale, int hd, int sk)
{
const int tid = threadIdx.x;
const int wid = tid / 32, lane = tid % 32;
const int n_ktiles = hd / MMA_K; // 4 for hd=64
// SMEM: sQ (128, HD) canonical + sK (128, HD) canonical
// Each K-tile of (128, 16) = 4096 bytes
// Full (128, HD) = n_ktiles * 4096 bytes
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 * hd;
// TMEM alloc (128 cols for N=128)
// TMEM alloc
if (wid == 1) {
tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128);
}
__syncthreads();
uint32_t tb = *sTmemBase;
// Zero all first
for (int i = tid; i < 128 * hd; i += N_WARPS * 32) {
sQ[i] = 0;
sK[i] = 0;
}
// Zero SMEM
for (int i = tid; i < 128 * hd; i += 128) { sQ[i] = 0; sK[i] = 0; }
__syncthreads();
// Write Q (1, hd) to row 0 of sQ in canonical layout
for (int d = tid; d < hd; d += N_WARPS * 32) {
int core_k = d / 8, local_c = d % 8;
int idx = core_k * 16 * 64 + local_c; // tile_mn=0, local_r=0
sQ[idx] = q[d];
// Also write to flat offset for verification
if (d < 8) s_out[220 + d] = (float)idx; // Log the indices we write to
// Write Q (1, hd) to sQ row 0 in canonical layout
for (int d = tid; d < hd; d += 128) {
int ck = d / 8, lc = d % 8;
sQ[ck * 16 * 64 + lc] = q[d];
}
// Write K (sk, hd) to sK in canonical layout
for (int i = tid; i < sk * hd; i += N_WARPS * 32) {
for (int i = tid; i < sk * hd; i += 128) {
int r = i / hd, c = i % hd;
int tile_mn = r / 8, core_k = c / 8;
int local_r = r % 8, local_c = c % 8;
int idx = core_k * 16 * 64 + tile_mn * 64 + local_r * 8 + local_c;
sK[idx] = k[i];
int tmn = r / 8, ck = c / 8, lr = r % 8, lc = c % 8;
sK[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[i];
}
__syncthreads();
// Verify SMEM data for first K-tile (columns 0-15)
// In canonical layout, Q[d] for row 0 is at core_k * 16 * 64 + local_c
// Verify Q[0..7] in SMEM
if (tid == 0) {
// Q row 0, d=0..7: core_k=0, local_c=d → sQ[d]
for (int d = 0; d < 8; d++)
s_out[200+d] = bf16_to_f32(sQ[d]);
// Q row 0, d=8..15: core_k=1, local_c=d-8 → sQ[1024 + d-8]
for (int d = 0; d < 8; d++)
s_out[208+d] = bf16_to_f32(sQ[1024 + d]);
for (int d = 0; d < 8; d++) s_out[200+d] = bf16_to_f32(sQ[d]);
for (int d = 0; d < 8; d++) s_out[208+d] = bf16_to_f32(sQ[1024+d]); // Q[8..15]
}
__syncthreads();
// MMA with K-tile 0 (columns 0-15)
// Descriptor: start = sQ_smem (base of the matrix)
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, 128);
uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 128);
uint32_t idesc = make_idesc(128, 128);
// For each K-tile, construct a descriptor pointing to that 16-column slice
// K-tile k: columns [16k, 16k+16) of the (128, hd) matrix
// In canonical layout, the k-th 16-column slice starts at:
// k * CORES_MN * 64 = k * 16 * 64 = k * 1024 BF16 = k * 2048 bytes
// Each K-tile has BLOCK_M=128 rows and 16 columns.
// The descriptor for K-tile k: start_addr = sQ_smem + k * 2048
// But wait — gau-nernst uses A_smem + k * BLOCK_M * 32 for the start address
// BLOCK_M * 32 = 128 * 32 = 4096 bytes. But our K-tile is only 2048 bytes.
// Actually, gau-nernst's offset is for the SMEM start of the K-tile.
// In his layout, each (BLOCK_M, 32B) slice is BLOCK_M * 32 = 4096 bytes apart.
// But 32B = 16 BF16 = one K-tile. And his start_address = A_smem + k * BLOCK_M * 32.
// Wait — that's 4096 bytes per K-tile, but our (128, 16) matrix is only 4096 bytes.
// In canonical layout, the K-tile at columns [16k, 16k+16) starts at:
// The 2 core-matrix columns (c=2k and c=2k+1) are at offsets 2k*2048 and (2k+1)*2048.
// The full K-tile spans both columns: core(0..15, 2k) and core(0..15, 2k+1).
// The first column starts at 2k * 2048 and the second at (2k+1) * 2048.
// Total span: from 2k*2048 to (2k+1)*2048 + 16*128 = (2k+1)*2048 + 2048 = (2k+2)*2048.
// Hmm, this is getting complicated. The descriptor for a K-tile should describe
// a (128, 16) matrix starting at the right offset in SMEM.
// For K-tile 0 (columns 0-15): start at sQ_smem, LBO=2048, SBO=128
// For K-tile 1 (columns 16-31): start at sQ_smem + 2*2048, LBO=2048, SBO=128
// Wait, but columns 16-23 are core_k=2 and columns 24-31 are core_k=3.
// The K-tile at columns [16k, 16k+16) has core_k = 2k and 2k+1.
// core_k=2 starts at 2 * 1024 = 2048 BF16 = 4096 bytes from sQ.
// core_k=3 starts at 3 * 1024 = 3072 BF16 = 6144 bytes from sQ.
// But the descriptor's start_address is for the BEGINNING of the (128, 16) tile.
// The descriptor with LBO=2048 walks: column 0 at start, column 1 at start+2048.
// So for K-tile k, we need:
// start = sQ_smem + (2k) * 1024 * 2 = sQ_smem + 2k * 2048 bytes
// Column 0 = core_k 2k, column 1 = core_k 2k+1
// LBO = 2048 (same as before)
// SBO = 128 (same as before)
// Actually, this IS the same descriptor but with a different start_address.
// The descriptor for K-tile k: start = sQ + 2k * 2048
uint32_t idesc = make_idesc(BLOCK_M, 128);
// K-tile loop with accumulate
for (int kt = 0; kt < n_ktiles; kt++) {
// Descriptor for Q's k-th K-tile
uint32_t q_ktile_addr = sQ_smem + kt * 2 * 2048; // 2 core-matrix columns per K-tile
uint32_t k_ktile_addr = sK_smem + kt * 2 * 2048;
uint64_t desc_q = make_umma_desc_kmajor_none(q_ktile_addr, BLOCK_M);
uint64_t desc_k = make_umma_desc_kmajor_none(k_ktile_addr, BLOCK_M);
bool accumulate = (kt > 0);
// Debug: write descriptor info
if (tid == 0 && kt == 0) {
s_out[240] = (float)q_ktile_addr;
s_out[241] = (float)k_ktile_addr;
s_out[242] = (float)desc_q;
s_out[243] = (float)desc_k;
}
// 4 warp leaders call MMA
if (lane == 0) {
umma_ss_f16(tb, desc_q, desc_k, idesc, accumulate);
}
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
if (lane == 0) {
umma_ss_f16(tb, desc_q, desc_k, idesc, false);
}
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
// Read from TMEM: 32x32b.x8, each warp reads 32 rows × 8 columns
// Read TMEM
for (int n = 0; n < 128 / 8; n++) {
const int row = wid * 32;
const int col = n * 8;
const int addr = tb + (row << 16) + col;
const int addr = tb + (row << 16) + n * 8;
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.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;");
int out_row = wid * 32 + lane;
if (n < 2 && out_row < 128) { // First 16 cols
for (int c = 0; c < 8; c++) {
if (n < 2 && out_row < 128) {
for (int c = 0; c < 8; c++)
s_out[out_row * 16 + n * 8 + c] = tmp[c] * scale;
}
}
}
__syncthreads();
// Scalar reference
// Scalar: S[0,j] = sum(Q[0,d]*K[j,d], d=0..15) * scale (first K-tile only)
if (tid == 0) {
float* q_row = new float[hd];
for (int d = 0; d < hd; d++) q_row[d] = bf16_to_f32(q[d]);
for (int c = 0; c < sk; c++) {
for (int j = 0; j < sk; j++) {
float dot = 0.0f;
for (int d = 0; d < hd; d++)
dot += q_row[d] * bf16_to_f32(k[c * hd + d]);
s_scalar[c] = dot * scale;
for (int d = 0; d < 16; d++) // Only first K-tile
dot += bf16_to_f32(q[d]) * bf16_to_f32(k[j * hd + d]);
s_scalar[j] = dot * scale;
}
delete[] q_row;
}
__syncthreads();
@@ -191,15 +107,16 @@ test_umma_qk_hd64(const bf16_t* q, const bf16_t* k,
}
int main() {
printf("=== UMMA QK GEMM Test (HD=64, 4 K-tiles) ===\n");
printf("=== UMMA QK HD=64, 1 K-tile ===\n");
const int HD = 64, SK = 128;
const float SCALE = 1.0f / sqrtf((float)HD);
bf16_t* h_q = (bf16_t*)malloc(HD * sizeof(bf16_t));
bf16_t* h_k = (bf16_t*)malloc(SK * HD * sizeof(bf16_t));
float* h_s_out = (float*)calloc(128 * 16, sizeof(float));
float* h_s_out = (float*)calloc(128*16, sizeof(float));
float* h_s_scalar = (float*)calloc(SK, sizeof(float));
srand(42);
for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
for (int i = 0; i < SK*HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
@@ -208,13 +125,9 @@ int main() {
cudaMalloc(&d_s_out, 128*16*sizeof(float)); cudaMalloc(&d_s_scalar, SK*sizeof(float));
cudaMemcpy(d_q, h_q, HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemset(d_s_out, 0, 128*16*sizeof(float)); cudaMemset(d_s_scalar, 0, SK*sizeof(float));
// SMEM: sTmemBase(4) + pad(16) + sQ(128*64*2=16384) + sK(128*64*2=16384) + pad
int smem = (4 + 16 + 128*HD*2 + 128*HD*2 + 256 + 127) & ~127;
printf("SMEM: %d bytes\n", smem);
test_umma_qk_hd64<<<1, N_WARPS*32, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE, HD, SK);
int smem = (4+16+128*HD*2+128*HD*2+256+127)&~127;
test_umma_qk_hd64_1ktile<<<1, 128, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE, HD, SK);
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }
@@ -222,25 +135,22 @@ int main() {
cudaMemcpy(h_s_out, d_s_out, 128*16*sizeof(float), cudaMemcpyDeviceToHost);
cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost);
printf("Q[d=0..7] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ", h_s_out[200+d]); printf("\n");
printf("Q[d=16..23] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ", h_s_out[208+d]); printf("\n");
printf("Q[d=0..7] orig: "); for(int d=0;d<8;d++) printf("%.4f ", bf16_to_f32_host(h_q[d])); printf("\n");
printf("Q[d=16..23] orig: "); for(int d=0;d<8;d++) printf("%.4f ", bf16_to_f32_host(h_q[16+d])); printf("\n");
// Compare row 0
printf("S[0,0..7] MMA: ");
for (int c = 0; c < 8; c++) printf("%.6f ", h_s_out[0*16+c]);
printf("\nS[0,0..7] ref: ");
for (int c = 0; c < 8; c++) printf("%.6f ", h_s_scalar[c]);
printf("\n");
// SMEM verify
printf("Q[0..7] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ",h_s_out[200+d]); printf("\n");
printf("Q[0..7] orig: "); for(int d=0;d<8;d++) printf("%.4f ",bf16_to_f32_host(h_q[d])); printf("\n");
printf("Q[8..15] SMEM: "); for(int d=0;d<8;d++) printf("%.4f ",h_s_out[208+d]); printf("\n");
printf("Q[8..15] orig: "); for(int d=0;d<8;d++) printf("%.4f ",bf16_to_f32_host(h_q[8+d])); printf("\n");
// Compare
printf("S[0,0..7] MMA: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_out[0*16+c]); printf("\n");
printf("S[0,0..7] ref: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_scalar[c]); printf("\n");
float max_diff = 0.0f, max_val = 0.0f;
for (int c = 0; c < 16; c++) {
for (int c = 0; c < 8; c++) {
max_diff = fmaxf(max_diff, fabsf(h_s_out[0*16+c] - h_s_scalar[c]));
max_val = fmaxf(max_val, fabsf(h_s_scalar[c]));
}
float rel_err = max_val > 0 ? max_diff / max_val : max_diff;
printf("Row 0 rel err (16 cols): %.6f\n", rel_err);
printf("Row 0 rel err: %.6f\n", rel_err);
printf("Test %s\n", rel_err < 0.01f ? "PASSED" : "FAILED");
cudaFree(d_q); cudaFree(d_k); cudaFree(d_s_out); cudaFree(d_s_scalar);