From 9e6ba25a98a2fc9d445066fc0d837c78f9cf9685 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 09:01:52 +0000 Subject: [PATCH] P5: standalone multi-tile CUDA test (2 KV tiles, hd=64) --- tests/unit/test_p5_multitile.cu | 147 ++++++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 tests/unit/test_p5_multitile.cu diff --git a/tests/unit/test_p5_multitile.cu b/tests/unit/test_p5_multitile.cu new file mode 100644 index 00000000..3f567775 --- /dev/null +++ b/tests/unit/test_p5_multitile.cu @@ -0,0 +1,147 @@ +/** + * P5: Minimal multi-tile FMHA test. + * Tests the multi-tile path (n_kv_tiles=2, N=256) against a CPU reference. + */ +#include "dsv4/kernels/attention/fmha_common.cuh" +#include "dsv4/kernels/attention/fmha_umma_desc.cuh" +#include "dsv4/kernels/attention/fmha_6warp_multihead.cuh" + +#include +#include +#include + +using namespace dsv4::kernels::attention; + +// CPU reference attention for single head +void reference_attention( + const bf16_t* q, const bf16_t* k, const bf16_t* v, + float* o_ref, float* lse_ref, + int hd, int s_k, float scale +) { + for (int d = 0; d < hd; d++) o_ref[d] = 0.0f; + float row_max = -INFINITY; + // Compute max + for (int j = 0; j < s_k; 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]); + dot *= scale; + if (dot > row_max) row_max = dot; + } + // Compute exp and weighted sum + float row_sum = 0.0f; + for (int j = 0; j < s_k; 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]); + dot *= scale; + float p = expf(dot - row_max); + row_sum += p; + for (int d = 0; d < hd; d++) o_ref[d] += p * bf16_to_f32(v[d * s_k + j]); + } + for (int d = 0; d < hd; d++) o_ref[d] /= row_sum; + *lse_ref = logf(row_sum) + row_max; +} + +int main() { + constexpr int HD = 64; + constexpr int SK = 256; // 2 KV tiles + constexpr float SCALE = 1.0f / sqrtf((float)HD); + + // Allocate host data + bf16_t h_q[HD], h_k[SK * HD], h_v[HD * SK]; + float h_o_ref[HD], h_lse_ref; + + srand(42); + for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16((float)(rand() % 100) / 100.0f); + for (int j = 0; j < SK * HD; j++) h_k[j] = f32_to_bf16((float)(rand() % 100) / 100.0f); + for (int j = 0; j < HD * SK; j++) h_v[j] = f32_to_bf16((float)(rand() % 100) / 100.0f); + + // CPU reference + reference_attention(h_q, h_k, h_v, h_o_ref, &h_lse_ref, HD, SK, SCALE); + + // GPU test + bf16_t *d_q, *d_k, *d_v, *d_o; + float *d_lse; + cudaMalloc(&d_q, HD * 2); + cudaMalloc(&d_k, SK * HD * 2); + cudaMalloc(&d_v, HD * SK * 2); + cudaMalloc(&d_o, HD * 2); + cudaMalloc(&d_lse, 4); + + cudaMemcpy(d_q, h_q, HD * 2, cudaMemcpyHostToDevice); + cudaMemcpy(d_k, h_k, SK * HD * 2, cudaMemcpyHostToDevice); + cudaMemcpy(d_v, h_v, HD * SK * 2, cudaMemcpyHostToDevice); + cudaMemset(d_o, 0, HD * 2); + cudaMemset(d_lse, 0, 4); + + FmhaParams params; + params.q = d_q; + params.k = d_k; + params.v = d_v; + params.o = d_o; + params.lse = d_lse; + params.s_k = SK; + params.scale = SCALE; + params.head_dim = HD; + params.n_kv_tiles = 2; // Multi-tile! + params.q_head_stride = 0; + params.q_batch_stride = 0; + params.k_head_stride = 0; + params.k_batch_stride = 0; + params.v_head_stride = 0; + params.v_batch_stride = 0; + params.o_head_stride = 0; + params.o_batch_stride = 0; + params.lse_head_stride = 0; + params.lse_batch_stride = 0; + + // SMEM size + constexpr int TILE_SZ = 128 * MMA_K_BF16; + constexpr int V_SUB_SZ = 256; + int smem = 4 + 4 + 4 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + 128*4 + HD*4 + 256 + 127; + smem &= ~127; + + dim3 grid(1, 1, 1); + dim3 block(NTHREADS); + + if (smem > 48 * 1024) { + cudaFuncSetAttribute(fmha_6warp_multihead_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + + fmha_6warp_multihead_kernel<<>>(params); + cudaError_t err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { + printf("Kernel launch FAILED: %s\n", cudaGetErrorString(err)); + return 1; + } + + // Compare + bf16_t h_o[HD]; + float h_lse; + cudaMemcpy(h_o, d_o, HD * 2, cudaMemcpyDeviceToHost); + cudaMemcpy(&h_lse, d_lse, 4, cudaMemcpyDeviceToHost); + + float cos = 0, norm_a = 0, norm_b = 0; + for (int d = 0; d < HD; d++) { + float a = h_o_ref[d]; + float b = bf16_to_f32(h_o[d]); + cos += a * b; + norm_a += a * a; + norm_b += b * b; + } + cos /= sqrtf(norm_a * norm_b + 1e-30f); + + printf("Multi-tile FMHA (HD=%d, SK=%d, n_kv_tiles=2):\n", HD, SK); + printf(" LSE: kernel=%.4f ref=%.4f\n", h_lse, h_lse_ref); + printf(" Cosine similarity: %.6f\n", cos); + printf(" Kernel O[0:5]:"); + for (int d = 0; d < 5; d++) printf(" %.4f", bf16_to_f32(h_o[d])); + printf("\n Ref O[0:5]:"); + for (int d = 0; d < 5; d++) printf(" %.4f", h_o_ref[d]); + printf("\n"); + + cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse); + + if (cos >= 0.999990) { printf("PASS\n"); return 0; } + else { printf("FAIL (cos < 0.999990)\n"); return 1; } +}