192 lines
7.3 KiB
Plaintext
192 lines
7.3 KiB
Plaintext
/**
|
||
* UMMA QK GEMM Test — HD=64 (4 K-tiles, separate SMEM per K-tile)
|
||
*
|
||
* Each K-tile gets its own (128, 16) SMEM region — no offset descriptors.
|
||
* Source data stride handled correctly (SRC_HD=64, SMEM_HD=16).
|
||
*/
|
||
|
||
#include <cuda_runtime.h>
|
||
#include <cstdio>
|
||
#include <cmath>
|
||
#include <cstdlib>
|
||
#include <cstring>
|
||
|
||
#include "dsv4/kernels/attention/fmha_common.cuh"
|
||
#include "dsv4/kernels/attention/fmha_umma_desc.cuh"
|
||
|
||
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); }
|
||
|
||
constexpr int HD = 64;
|
||
constexpr int SK = 128;
|
||
constexpr int NKT = HD / MMA_K_BF16; // 4
|
||
constexpr int BLOCK_MN = 128;
|
||
constexpr int TILE_SZ = BLOCK_MN * MMA_K_BF16; // 128*16 = 2048 BF16 per K-tile
|
||
|
||
__global__ void __launch_bounds__(128)
|
||
test_umma_hd64(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
|
||
float* __restrict__ s_out, float* __restrict__ s_scalar, float scale)
|
||
{
|
||
const int tid = threadIdx.x;
|
||
const int wid = tid / WARP, lane = tid % WARP;
|
||
|
||
// SMEM: tmem_base(4) + pad(12) + Q tiles (4 × 2048 BF16) + K tiles (4 × 2048 BF16)
|
||
extern __shared__ char sbuf[];
|
||
uint32_t* sTmemBase = (uint32_t*)sbuf;
|
||
bf16_t* sQ0 = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
|
||
bf16_t* sQ1 = sQ0 + TILE_SZ;
|
||
bf16_t* sQ2 = sQ1 + TILE_SZ;
|
||
bf16_t* sQ3 = sQ2 + TILE_SZ;
|
||
bf16_t* sK0 = sQ3 + TILE_SZ;
|
||
bf16_t* sK1 = sK0 + TILE_SZ;
|
||
bf16_t* sK2 = sK1 + TILE_SZ;
|
||
bf16_t* sK3 = sK2 + TILE_SZ;
|
||
|
||
constexpr int CORES_MN = 16; // 128/8
|
||
|
||
// Load Q K-tiles: Q is (1, 64), each K-tile takes 16 dims
|
||
// Zero all tiles
|
||
for (int i = tid; i < NKT * TILE_SZ; i += 128) { sQ0[i] = 0; sK0[i] = 0; }
|
||
__syncthreads();
|
||
|
||
// Write Q row 0 to each K-tile's SMEM
|
||
for (int kt = 0; kt < NKT; kt++) {
|
||
bf16_t* sq = sQ0 + kt * TILE_SZ;
|
||
for (int d = tid; d < MMA_K_BF16; d += 128) {
|
||
int ck = d / 8, lc = d % 8;
|
||
sq[ck * CORES_MN * 64 + lc] = q[kt * MMA_K_BF16 + d];
|
||
}
|
||
// Write K for this K-tile: K[r, 16*kt + d] for r=0..127, d=0..15
|
||
bf16_t* sk = sK0 + kt * TILE_SZ;
|
||
for (int r = 0; r < SK; r++) {
|
||
for (int d = tid; d < MMA_K_BF16; d += 128) {
|
||
int ck = d / 8, lc = d % 8;
|
||
int tmn = r / 8, lr = r % 8;
|
||
sk[ck * CORES_MN * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * MMA_K_BF16 + d];
|
||
}
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
// TMEM alloc
|
||
if (wid == 1) {
|
||
tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128);
|
||
}
|
||
__syncthreads();
|
||
uint32_t tb = *sTmemBase;
|
||
|
||
// Multi-K-tile QK GEMM with separate SMEM per K-tile
|
||
bf16_t* sQ_arr[NKT] = {sQ0, sQ1, sQ2, sQ3};
|
||
bf16_t* sK_arr[NKT] = {sK0, sK1, sK2, sK3};
|
||
uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN);
|
||
|
||
for (int kt = 0; kt < NKT; kt++) {
|
||
uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ_arr[kt]), BLOCK_MN);
|
||
uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK_arr[kt]), BLOCK_MN);
|
||
|
||
if (tid == 0) {
|
||
umma_ss_f16(tb, dq, dk, idesc, kt > 0);
|
||
}
|
||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||
__syncthreads();
|
||
}
|
||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||
__syncthreads();
|
||
|
||
// Read S from 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;
|
||
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;");
|
||
|
||
int out_row = wid * 32 + lane;
|
||
if (out_row < SK) {
|
||
for (int c = 0; c < 8; c++) {
|
||
int out_col = n * 8 + c;
|
||
if (out_col < SK) {
|
||
s_out[out_row * SK + out_col] = tmp[c] * scale;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
__syncthreads();
|
||
|
||
// Scalar reference
|
||
if (tid == 0) {
|
||
for (int j = 0; j < SK; j++) {
|
||
float dot = 0.0f;
|
||
for (int d = 0; d < HD; d++)
|
||
dot += bf16_to_f32(q[d]) * bf16_to_f32(k[j * HD + d]);
|
||
s_scalar[j] = dot * scale;
|
||
}
|
||
}
|
||
|
||
if (wid == 0) tmem_dealloc(tb, 128);
|
||
}
|
||
|
||
int main() {
|
||
printf("=== UMMA QK GEMM HD=64 (separate SMEM per K-tile) ===\n");
|
||
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(SK * SK, 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);
|
||
|
||
bf16_t *d_q, *d_k; float *d_s_out, *d_s_scalar;
|
||
cudaMalloc(&d_q, HD*sizeof(bf16_t)); cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));
|
||
cudaMalloc(&d_s_out, SK*SK*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);
|
||
|
||
// SMEM: 4 + 12(pad) + 8 * 2048*2
|
||
int smem = (4 + 16 + NKT * 2 * TILE_SZ * sizeof(bf16_t) + 256 + 127) & ~127;
|
||
printf("SMEM: %d bytes (%d KB)\n", smem, smem / 1024);
|
||
|
||
test_umma_hd64<<<1, 128, smem>>>(d_q, d_k, d_s_out, d_s_scalar, SCALE);
|
||
|
||
cudaError_t err = cudaDeviceSynchronize();
|
||
if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }
|
||
|
||
cudaMemcpy(h_s_out, d_s_out, SK*SK*sizeof(float), cudaMemcpyDeviceToHost);
|
||
cudaMemcpy(h_s_scalar, d_s_scalar, SK*sizeof(float), cudaMemcpyDeviceToHost);
|
||
|
||
printf("S[0,0..7] MMA: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n");
|
||
printf("S[0,0..7] ref: "); for(int c=0;c<8;c++) printf("%.6f ",h_s_scalar[c]); printf("\n");
|
||
printf("S[0,64..71] MMA: "); for(int c=64;c<72;c++) printf("%.6f ",h_s_out[0*SK+c]); printf("\n");
|
||
printf("S[0,64..71] ref: "); for(int c=64;c<72;c++) printf("%.6f ",h_s_scalar[c]); printf("\n");
|
||
|
||
float max_diff = 0.0f, max_val = 0.0f;
|
||
for (int c = 0; c < SK; c++) {
|
||
float diff = fabsf(h_s_out[0*SK+c] - h_s_scalar[c]);
|
||
max_diff = fmaxf(max_diff, diff);
|
||
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 (128 cols): %.8f\n", rel_err);
|
||
|
||
float max_nonzero = 0.0f;
|
||
for (int r = 1; r < SK; r++)
|
||
for (int c = 0; c < SK; c++)
|
||
max_nonzero = fmaxf(max_nonzero, fabsf(h_s_out[r*SK+c]));
|
||
printf("Rows 1-127 max abs: %.8f\n", max_nonzero);
|
||
|
||
bool row0_ok = rel_err < 0.01f;
|
||
bool rows_zero = max_nonzero < 1e-4f;
|
||
printf("Row 0: %s | Rows 1-127 zero: %s\n",
|
||
row0_ok ? "PASS" : "FAIL", rows_zero ? "PASS" : "FAIL");
|
||
printf("Overall: %s\n", (row0_ok && rows_zero) ? "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 (row0_ok && rows_zero) ? 0 : 1;
|
||
}
|