nvfp4-megamoe-kernel/tests/unit/test_p5_multitile.cu

/**
 * 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 <cstdio>
#include <cmath>
#include <cstring>

using namespace dsv4::kernels::attention;

// Host-side BF16 helpers
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);
}

// 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 += hbf16_to_f32(q[d]) * hbf16_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 += hbf16_to_f32(q[d]) * hbf16_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 * hbf16_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
    const 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] = 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);

    // 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<HD, 128>,
            cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
    }

    fmha_6warp_multihead_kernel<HD, 128><<<grid, block, smem>>>(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 = hbf16_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", hbf16_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; }
}