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

P6: Add minimal CUDA test for TMA store epilogue

Browse Source
This commit is contained in:
biondizzle
2026-05-30 16:59:45 +00:00
parent 6a7726e764
commit 2833eb56e7
1 changed files with 183 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

183
tests/unit/test_p6_tma_store.cu Normal file
View File

@@ -0,0 +1,183 @@
/**
* P6 Minimal TMA store epilogue test.
*
* Tests the TMA store path of the 6-warp multi-head kernel.
* Compares TMA store output vs direct GMEM write output.
* If they match, the TMA store pipeline is correct.
*/
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_runtime.h>
#include <cuda.h>
#include "dsv4/kernels/attention/fmha_common.cuh"
#include "dsv4/kernels/attention/fmha_umma_desc.cuh"
#include "dsv4/kernels/attention/fmha_tma.cuh"
#include "dsv4/kernels/attention/fmha_6warp_multihead.cuh"
using namespace dsv4::kernels::attention;
static float cosine_sim(const float* a, const float* b, int n) {
double dot = 0, na = 0, nb = 0;
for (int i = 0; i < n; i++) {
dot += a[i] * b[i];
na += a[i] * a[i];
nb += b[i] * b[i];
}
return (float)(dot / (sqrt(na) * sqrt(nb) + 1e-30));
}
int main() {
constexpr int HD = 64;
constexpr int N = 128;
constexpr int n_h = 4;
constexpr int batch = 1;
constexpr float scale = 1.0f / sqrtf((float)HD);
// Allocate tensors
bf16_t *d_q, *d_k, *d_v;
bf16_t *d_o_direct, *d_o_tma;
float *d_lse_direct, *d_lse_tma;
cudaMalloc(&d_q, batch * n_h * 1 * HD * sizeof(bf16_t));
cudaMalloc(&d_k, batch * n_h * N * HD * sizeof(bf16_t));
cudaMalloc(&d_v, batch * n_h * HD * N * sizeof(bf16_t));
cudaMalloc(&d_o_direct, batch * n_h * 1 * HD * sizeof(bf16_t));
cudaMalloc(&d_o_tma, batch * n_h * 1 * HD * sizeof(bf16_t));
cudaMalloc(&d_lse_direct, batch * n_h * 1 * sizeof(float));
cudaMalloc(&d_lse_tma, batch * n_h * 1 * sizeof(float));
// Initialize with random data
srand(42);
auto init_bf16 = [](bf16_t* d, int n) {
float* h = new float[n];
for (int i = 0; i < n; i++) h[i] = (float)rand() / RAND_MAX - 0.5f;
for (int i = 0; i < n; i++) {
unsigned short us;
asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(us) : "f"(h[i]));
d[i] = us;
}
delete[] h;
};
init_bf16(d_q, batch * n_h * HD);
init_bf16(d_k, batch * n_h * N * HD);
init_bf16(d_v, batch * n_h * HD * N);
// Launch with direct GMEM write (tma_o = nullptr)
{
FmhaParams params;
params.q = d_q;
params.k = d_k;
params.v = d_v;
params.o = d_o_direct;
params.lse = d_lse_direct;
params.s_k = N;
params.scale = scale;
params.head_dim = HD;
params.q_head_stride = HD;
params.q_batch_stride = n_h * HD;
params.k_head_stride = N * HD;
params.k_batch_stride = n_h * N * HD;
params.v_head_stride = HD * N;
params.v_batch_stride = n_h * HD * N;
params.o_head_stride = HD;
params.o_batch_stride = n_h * HD;
params.lse_head_stride = 1;
params.lse_batch_stride = n_h;
params.tma_o = nullptr;
int smem = 32768;
dim3 grid(1, n_h, batch);
fmha_6warp_multihead_kernel<HD, 128><<<grid, 192, smem>>>(params);
cudaDeviceSynchronize();
}
// Launch with TMA store (tma_o set)
{
// Create TMA descriptors for O
CUtensorMap* d_tma_o;
cudaMalloc(&d_tma_o, n_h * batch * sizeof(CUtensorMap));
for (int b = 0; b < batch; b++) {
for (int h = 0; h < n_h; h++) {
int idx = b * n_h + h;
bf16_t* o_head = d_o_tma + h * HD + b * n_h * HD;
CUtensorMap h_desc;
bool ok = create_tma_desc_2d_bf16(&h_desc, o_head, 1, HD, 1, HD);
if (!ok) { printf("TMA desc creation FAILED for head %d\n", h); return 1; }
cudaMemcpy(d_tma_o + idx, &h_desc, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
}
}
FmhaParams params;
params.q = d_q;
params.k = d_k;
params.v = d_v;
params.o = d_o_tma;
params.lse = d_lse_tma;
params.s_k = N;
params.scale = scale;
params.head_dim = HD;
params.q_head_stride = HD;
params.q_batch_stride = n_h * HD;
params.k_head_stride = N * HD;
params.k_batch_stride = n_h * N * HD;
params.v_head_stride = HD * N;
params.v_batch_stride = n_h * HD * N;
params.o_head_stride = HD;
params.o_batch_stride = n_h * HD;
params.lse_head_stride = 1;
params.lse_batch_stride = n_h;
params.tma_o = d_tma_o;
int smem = 32768;
dim3 grid(1, n_h, batch);
fmha_6warp_multihead_kernel<HD, 128><<<grid, 192, smem>>>(params);
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) {
printf("TMA kernel FAILED: %s\n", cudaGetErrorString(err));
cudaFree(d_tma_o);
return 1;
}
cudaFree(d_tma_o);
}
// Compare outputs
bf16_t* h_o_direct = new bf16_t[n_h * HD];
bf16_t* h_o_tma = new bf16_t[n_h * HD];
cudaMemcpy(h_o_direct, d_o_direct, n_h * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);
cudaMemcpy(h_o_tma, d_o_tma, n_h * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);
float* f_direct = new float[n_h * HD];
float* f_tma = new float[n_h * HD];
for (int i = 0; i < n_h * HD; i++) {
f_direct[i] = bf16_to_f32(h_o_direct[i]);
f_tma[i] = bf16_to_f32(h_o_tma[i]);
}
float cos = cosine_sim(f_direct, f_tma, n_h * HD);
printf("P6 TMA epilogue test (hd=%d, n_h=%d): cos=%.8f\n", HD, n_h, cos);
if (cos >= 0.999999f) {
printf("PASS: TMA and direct paths produce identical results\n");
} else {
printf("FAIL: TMA and direct paths differ\n");
// Print first few values for debugging
for (int h = 0; h < n_h; h++) {
printf(" Head %d: direct[0..3]=[%.4f,%.4f,%.4f,%.4f] tma[0..3]=[%.4f,%.4f,%.4f,%.4f]\n",
h, f_direct[h*HD], f_direct[h*HD+1], f_direct[h*HD+2], f_direct[h*HD+3],
f_tma[h*HD], f_tma[h*HD+1], f_tma[h*HD+2], f_tma[h*HD+3]);
}
}
delete[] h_o_direct;
delete[] h_o_tma;
delete[] f_direct;
delete[] f_tma;
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v);
cudaFree(d_o_direct); cudaFree(d_o_tma);
cudaFree(d_lse_direct); cudaFree(d_lse_tma);
return (cos >= 0.999999f) ? 0 : 1;
}