diff --git a/tests/unit/test_p6_tma_store.cu b/tests/unit/test_p6_tma_store.cu new file mode 100644 index 00000000..b97a7b45 --- /dev/null +++ b/tests/unit/test_p6_tma_store.cu @@ -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 +#include +#include +#include +#include + +#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<<>>(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<<>>(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; +}