test: clean rewrite, 32 TMEM cols, MMA N=32, tmem_load loop
This commit is contained in:
@@ -1,11 +1,6 @@
|
||||
/**
|
||||
* Standalone CUDA test for UMMA QK GEMM (tcgen05.mma SS, BF16).
|
||||
*
|
||||
* Tests that tcgen05.mma produces correct QK attention scores.
|
||||
* Uses K-major, SWIZZLE_NONE, BF16 descriptors with canonical SMEM layout.
|
||||
* Based on gau-nernst's tcgen05 tutorial.
|
||||
*
|
||||
* Test: HD=16, SK=128 (single K-tile, single MMA call)
|
||||
* Test: HD=16, SK=128, single K-tile, MMA with N=32 (32 TMEM columns)
|
||||
*/
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
@@ -19,128 +14,63 @@
|
||||
|
||||
using namespace dsv4::kernels::attention;
|
||||
|
||||
// Host-side BF16 conversion
|
||||
static bf16_t f32_to_bf16_host(float f) {
|
||||
uint32_t u; memcpy(&u, &f, 4);
|
||||
return (uint16_t)(u >> 16);
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// Test kernel: UMMA QK GEMM for HD=16, SK=128
|
||||
// ==================================================================
|
||||
// A = Q: (128, 16) padded, K-major core-matrix layout
|
||||
// B = K: (128, 16) in K-major core-matrix layout
|
||||
// MMA computes S = A × B^T = Q × K^T → (128, 128) in TMEM
|
||||
|
||||
__global__ void __launch_bounds__(NTHREADS)
|
||||
test_umma_qk_hd16(
|
||||
const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
|
||||
float* __restrict__ s_out, // output: S[0, 0..127] from TMEM
|
||||
float* __restrict__ s_scalar, // scalar reference
|
||||
float* __restrict__ s_out, float* __restrict__ s_scalar,
|
||||
float scale
|
||||
) {
|
||||
const int tid = threadIdx.x;
|
||||
const int wid = tid / WARP, lane = tid % WARP;
|
||||
|
||||
// ================================================================
|
||||
// SMEM layout
|
||||
// ================================================================
|
||||
// sQ: (128, 16) in canonical layout = 4096 bytes
|
||||
// sK: (128, 16) in canonical layout = 4096 bytes
|
||||
// sTmemBase: 4 bytes
|
||||
// sQ_row: 16 floats for scalar ref
|
||||
extern __shared__ char sbuf[];
|
||||
uint32_t* sTmemBase = (uint32_t*)sbuf;
|
||||
// Align to 16 bytes for UMMA
|
||||
bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
|
||||
bf16_t* sK = sQ + 128 * 16; // 4096 bytes after Q
|
||||
bf16_t* sK = sQ + 128 * 16;
|
||||
float* sQ_row = (float*)(sK + 128 * 16);
|
||||
|
||||
// Load Q to sQ_row for scalar reference
|
||||
for (int d = tid; d < 16; d += NTHREADS) {
|
||||
// Load Q for scalar reference
|
||||
for (int d = tid; d < 16; d += NTHREADS)
|
||||
sQ_row[d] = bf16_to_f32(q[d]);
|
||||
}
|
||||
|
||||
// ================================================================
|
||||
// TMEM allocation
|
||||
// ================================================================
|
||||
// TMEM allocation — 32 columns
|
||||
if (wid == 0) {
|
||||
uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase);
|
||||
tmem_alloc(smem_ptr, 32); // 32 columns like minimal test
|
||||
tmem_alloc(smem_ptr, 32);
|
||||
}
|
||||
__syncthreads();
|
||||
uint32_t tmem_base = *sTmemBase;
|
||||
|
||||
// Debug: write tmem_base to output BEFORE tmem_store
|
||||
if (tid == 0) {
|
||||
s_out[138] = (float)tmem_base; // should be 0.0f if first alloc
|
||||
s_out[139] = 1.0f; // sentinel: reached this point
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Zero TMEM — skip (minimal test didn't zero)
|
||||
// if (wid == 0) {
|
||||
// tmem_store(tmem_base, 0, 0, 0, 0);
|
||||
// __syncwarp();
|
||||
// tmem_store(tmem_base + 1, 0, 0, 0, 0);
|
||||
// __syncwarp();
|
||||
// }
|
||||
// __syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
s_out[140] = 2.0f; // sentinel: survived tmem_store
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// Load Q and K into SMEM in canonical layout
|
||||
// ================================================================
|
||||
write_q_to_smem<16>(sQ, q);
|
||||
write_k_to_smem<128, 16>(sK, k);
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// Construct descriptors and instruction descriptor
|
||||
// ================================================================
|
||||
// Construct descriptors
|
||||
uint32_t sQ_smem = __cvta_generic_to_shared(sQ);
|
||||
uint32_t sK_smem = __cvta_generic_to_shared(sK);
|
||||
|
||||
// K-major NONE: LBO = BLOCK_MN * 16, SBO = 128
|
||||
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);
|
||||
uint32_t idesc = make_idesc(128, 32); // M=128, N=32
|
||||
|
||||
// Debug: write descriptor values to output (positions 128-135)
|
||||
// Debug output
|
||||
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)__cvta_generic_to_shared(sQ);
|
||||
s_out[136] = (float)(sQ_smem >> 4); // start_address for desc
|
||||
s_out[137] = (float)(sQ_smem & 0xF); // alignment check
|
||||
s_out[135] = (float)tmem_base;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// SKIP MMA for now — just test descriptor construction
|
||||
// ================================================================
|
||||
// if (tid == 0) {
|
||||
// umma_ss_f16(tmem_base, desc_q, desc_k, idesc, /*accumulate=*/false);
|
||||
// }
|
||||
// __syncwarp();
|
||||
// if (wid == 0 && lane == 0) {
|
||||
// tmem_fence_store();
|
||||
// }
|
||||
// __syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// Call tcgen05.mma SS — test with 32 columns (N=32)
|
||||
// ================================================================
|
||||
// Rebuild idesc for N=32
|
||||
uint32_t idesc = make_idesc(128, 32);
|
||||
|
||||
// Call tcgen05.mma SS
|
||||
if (tid == 0) {
|
||||
umma_ss_f16(tmem_base, desc_q, desc_k, idesc, /*accumulate=*/false);
|
||||
}
|
||||
@@ -150,29 +80,24 @@ test_umma_qk_hd16(
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// Read S from TMEM
|
||||
// ================================================================
|
||||
if (wid == 0) {
|
||||
for (int col = 0; col < 32; col++) {
|
||||
uint32_t u0, u1, u2, u3;
|
||||
tmem_load(tmem_base + col, u0, u1, u2, u3);
|
||||
if (lane == 0) {
|
||||
s_out[col] = u32_to_f32(u0); // S[0, col]
|
||||
s_out[col] = u32_to_f32(u0);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// Scalar reference
|
||||
// ================================================================
|
||||
if (tid == 0) {
|
||||
for (int c = 0; c < 128; c++) {
|
||||
float dot = 0.0f;
|
||||
for (int d = 0; d < 16; d++) {
|
||||
for (int d = 0; d < 16; d++)
|
||||
dot += sQ_row[d] * bf16_to_f32(k[c * 16 + d]);
|
||||
}
|
||||
s_scalar[c] = dot * scale;
|
||||
}
|
||||
}
|
||||
@@ -184,12 +109,8 @@ test_umma_qk_hd16(
|
||||
}
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// Host
|
||||
// ==================================================================
|
||||
|
||||
int main() {
|
||||
printf("=== UMMA QK GEMM Test (HD=16, SK=128) ===\n");
|
||||
printf("=== UMMA QK GEMM Test (HD=16, SK=128, N=32) ===\n");
|
||||
|
||||
const int HD = 16, SK = 128;
|
||||
const float SCALE = 1.0f / sqrtf((float)HD);
|
||||
@@ -240,31 +161,23 @@ int main() {
|
||||
printf("\n");
|
||||
|
||||
float max_diff = 0.0f, max_val = 0.0f;
|
||||
for (int c = 0; c < SK; c++) {
|
||||
for (int c = 0; c < 32; c++) {
|
||||
max_diff = fmaxf(max_diff, fabsf(h_s_out[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("Max abs diff: %.6f, Max rel err: %.6f\n", max_diff, rel_err);
|
||||
|
||||
// Debug: print descriptor values
|
||||
uint64_t dbg_desc_q, dbg_desc_k;
|
||||
uint32_t dbg_idesc;
|
||||
memcpy(&dbg_desc_q, &h_s_out[128], 8);
|
||||
memcpy(&dbg_desc_k, &h_s_out[130], 8);
|
||||
memcpy(&dbg_idesc, &h_s_out[132], 4);
|
||||
printf("desc_q = 0x%016lx\n", dbg_desc_q);
|
||||
printf("desc_k = 0x%016lx\n", dbg_desc_k);
|
||||
printf("idesc = 0x%08x\n", dbg_idesc);
|
||||
printf("sQ_smem = %.0f\n", h_s_out[133]);
|
||||
printf("sK_smem = %.0f\n", h_s_out[134]);
|
||||
printf(" desc_q start_addr = %lu, LBO = %lu, SBO = %lu\n",
|
||||
dbg_desc_q & 0x3FFF, (dbg_desc_q >> 16) & 0x3FFF, (dbg_desc_q >> 32) & 0x3FFF);
|
||||
printf(" desc_k start_addr = %lu, LBO = %lu, SBO = %lu\n",
|
||||
dbg_desc_k & 0x3FFF, (dbg_desc_k >> 16) & 0x3FFF, (dbg_desc_k >> 32) & 0x3FFF);
|
||||
// Debug
|
||||
uint64_t dq, dk; uint32_t idi;
|
||||
memcpy(&dq, &h_s_out[128], 8);
|
||||
memcpy(&dk, &h_s_out[130], 8);
|
||||
memcpy(&idi, &h_s_out[132], 4);
|
||||
printf("desc_q=0x%016lx (addr=%lu,LBO=%lu,SBO=%lu)\n", dq, dq&0x3FFF, (dq>>16)&0x3FFF, (dq>>32)&0x3FFF);
|
||||
printf("desc_k=0x%016lx (addr=%lu,LBO=%lu,SBO=%lu)\n", dk, dk&0x3FFF, (dk>>16)&0x3FFF, (dk>>32)&0x3FFF);
|
||||
printf("idesc=0x%08x, tmem_base=%.0f\n", idi, h_s_out[135]);
|
||||
|
||||
printf("Test %s\n", rel_err < 0.01f ? "PASSED" : "FAILED");
|
||||
|
||||
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);
|
||||
return (rel_err < 0.01f) ? 0 : 1;
|
||||
|
||||
Reference in New Issue
Block a user