P5: standalone TMA multi-tile test with 128B-aligned memory

tests/unit/test_p5_tma_multitile.cu

@@ -0,0 +1,147 @@
+/**
+ * P5: Test multi-tile TMA FMHA kernel with proper alignment.
+ */
+#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_tma_multirow_multitile.cuh"
+
+#include <cstdio>
+#include <cmath>
+#include <cstring>
+
+using namespace dsv4::kernels::attention;
+
+static float hbf16_to_f32(uint16_t h) {
+    uint32_t u = ((uint32_t)h) << 16;
+    float f; memcpy(&f, &u, 4); return f;
+}
+static uint16_t hf32_to_bf16(float f) {
+    uint32_t u; memcpy(&u, &f, 4); return (uint16_t)(u >> 16);
+}
+
+int main() {
+    constexpr int HD = 64;
+    constexpr int SK = 256;
+    const float SCALE = 1.0f / sqrtf((float)HD);
+
+    // Allocate 128B-aligned GPU memory
+    bf16_t *d_q, *d_k, *d_v, *d_o;
+    float *d_lse;
+    cudaMalloc(&d_q, HD * 2 + 128);
+    cudaMalloc(&d_k, SK * HD * 2 + 128);
+    cudaMalloc(&d_v, HD * SK * 2 + 128);
+    cudaMalloc(&d_o, HD * 2 + 128);
+    cudaMalloc(&d_lse, 4 + 128);
+
+    // Align pointers
+    d_q = (bf16_t*)(((uintptr_t)d_q + 127) & ~(uintptr_t)127);
+    d_k = (bf16_t*)(((uintptr_t)d_k + 127) & ~(uintptr_t)127);
+    d_v = (bf16_t*)(((uintptr_t)d_v + 127) & ~(uintptr_t)127);
+    d_o = (bf16_t*)(((uintptr_t)d_o + 127) & ~(uintptr_t)127);
+    d_lse = (float*)(((uintptr_t)d_lse + 127) & ~(uintptr_t)127);
+
+    // Fill with random data
+    srand(42);
+    bf16_t h_q[HD], h_k[SK * HD], h_v[HD * SK];
+    for (int d = 0; d < HD; d++) h_q[d] = hf32_to_bf16((float)(rand() % 100) / 100.0f);
+    for (int j = 0; j < SK * HD; j++) h_k[j] = hf32_to_bf16((float)(rand() % 100) / 100.0f);
+    for (int j = 0; j < HD * SK; j++) h_v[j] = hf32_to_bf16((float)(rand() % 100) / 100.0f);
+
+    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);
+
+    // Create TMA descriptors
+    CUtensorMap h_tma_k, h_tma_v;
+    if (!create_tma_desc_2d_bf16(&h_tma_k, d_k, SK, HD, 128, 16)) {
+        printf("K TMA desc creation FAILED\n"); return 1;
+    }
+    if (!create_tma_desc_2d_bf16(&h_tma_v, d_v, HD, SK, 16, 16)) {
+        printf("V TMA desc creation FAILED\n"); return 1;
+    }
+
+    CUtensorMap *d_tma_k, *d_tma_v;
+    cudaMalloc(&d_tma_k, sizeof(CUtensorMap));
+    cudaMalloc(&d_tma_v, sizeof(CUtensorMap));
+    cudaMemcpy(d_tma_k, &h_tma_k, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
+    cudaMemcpy(d_tma_v, &h_tma_v, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
+
+    // Launch multi-tile kernel
+    FmhaTmaMultiRowMultiTileParams params;
+    params.q = d_q;
+    params.tma_k = d_tma_k;
+    params.tma_v = d_tma_v;
+    params.o = d_o;
+    params.lse = d_lse;
+    params.s_k = SK;
+    params.T = 1;
+    params.n_h = 1;
+    params.scale = SCALE;
+    params.q_head_stride = 0;
+    params.q_batch_stride = 0;
+    params.o_head_stride = 0;
+    params.o_batch_stride = 0;
+    params.lse_head_stride = 0;
+    params.lse_batch_stride = 0;
+
+    int smem = 4 + 16 + 128*16*2 + 128*16*2 + 128*16*2 + 128*16*2 + 16*16*2
+               + 128*256*4 + 128*4 + 128*4 + 128*4 + 128*4 + 256 + 127;
+    smem &= ~127;
+    dim3 grid(1, 1, 1);
+    dim3 block(NTHREADS);
+
+    cudaFuncSetAttribute(fmha_6warp_tma_multirow_multitile_kernel<HD>,
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+
+    fmha_6warp_tma_multirow_multitile_kernel<HD><<<grid, block, smem>>>(params);
+    cudaError_t err = cudaDeviceSynchronize();
+    if (err != cudaSuccess) {
+        printf("Kernel FAILED: %s\n", cudaGetErrorString(err));
+        return 1;
+    }
+
+    // Compare with CPU reference
+    float h_o_ref[HD], h_lse_ref;
+    float row_max = -INFINITY;
+    for (int j = 0; j < SK; j++) {
+        float dot = 0.0f;
+        for (int d = 0; d < HD; d++) dot += hbf16_to_f32(h_q[d]) * hbf16_to_f32(h_k[j * HD + d]);
+        dot *= SCALE;
+        if (dot > row_max) row_max = dot;
+    }
+    float row_sum = 0.0f;
+    for (int d = 0; d < HD; d++) h_o_ref[d] = 0.0f;
+    for (int j = 0; j < SK; j++) {
+        float dot = 0.0f;
+        for (int d = 0; d < HD; d++) dot += hbf16_to_f32(h_q[d]) * hbf16_to_f32(h_k[j * HD + d]);
+        dot *= SCALE;
+        float p = expf(dot - row_max);
+        row_sum += p;
+        for (int d = 0; d < HD; d++) h_o_ref[d] += p * hbf16_to_f32(h_v[d * SK + j]);
+    }
+    for (int d = 0; d < HD; d++) h_o_ref[d] /= row_sum;
+    h_lse_ref = logf(row_sum) + row_max;
+
+    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, na = 0, nb = 0;
+    for (int d = 0; d < HD; d++) {
+        float a = h_o_ref[d], b = hbf16_to_f32(h_o[d]);
+        cos += a * b; na += a * a; nb += b * b;
+    }
+    cos /= sqrtf(na * nb + 1e-30f);
+
+    printf("Multi-tile TMA FMHA (HD=%d, SK=%d):\n", HD, SK);
+    printf("  LSE: kernel=%.4f ref=%.4f\n", h_lse, h_lse_ref);
+    printf("  Cosine: %.6f\n", cos);
+    printf("  %s\n", cos >= 0.999990 ? "PASS" : "FAIL");
+
+    cudaFree(d_tma_k); cudaFree(d_tma_v);
+    return (cos >= 0.999990) ? 0 : 1;
+}