407 lines
15 KiB
Plaintext
407 lines
15 KiB
Plaintext
/**
|
||
* Test multi-head FMHA kernel (6-warp, grid launch).
|
||
* Compile with -DHD_VAL=64 etc.
|
||
*
|
||
* Tests:
|
||
* 1. Multi-head MHA: n_h independent heads, each with its own K/V
|
||
* 2. Multi-head MQA: n_h Q heads sharing 1 K/V head
|
||
* 3. Batched: batch_size > 1
|
||
* 4. LSE output correctness
|
||
*/
|
||
|
||
#include <cuda_runtime.h>
|
||
#include <cstdio>
|
||
#include <cmath>
|
||
#include <cstdlib>
|
||
#include <cstring>
|
||
|
||
#ifndef HD_VAL
|
||
#define HD_VAL 64
|
||
#endif
|
||
|
||
#include "dsv4/kernels/attention/fmha_common.cuh"
|
||
#include "dsv4/kernels/attention/fmha_umma_desc.cuh"
|
||
|
||
using namespace dsv4::kernels::attention;
|
||
|
||
static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
|
||
static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; }
|
||
|
||
constexpr int HD = HD_VAL;
|
||
constexpr int SK = 128;
|
||
constexpr int TILE_SZ = 128 * MMA_K_BF16;
|
||
constexpr int V_SUB_SZ = 256;
|
||
constexpr int TMEM_N = (HD <= 128) ? 128 : 256;
|
||
|
||
#include "dsv4/kernels/attention/fmha_6warp_multihead.cuh"
|
||
|
||
// Reference: compute attention for one head
|
||
static 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
|
||
) {
|
||
float s[512]; // max s_k
|
||
for (int j = 0; j < s_k; j++) {
|
||
float dot = 0.0f;
|
||
for (int d = 0; d < hd; d++) dot += bf16_to_f32_host(q[d]) * bf16_to_f32_host(k[j * hd + d]);
|
||
s[j] = dot * scale;
|
||
}
|
||
float mx = -INFINITY;
|
||
for (int j = 0; j < s_k; j++) mx = fmaxf(mx, s[j]);
|
||
float sm = 0.0f;
|
||
for (int j = 0; j < s_k; j++) { s[j] = expf(s[j] - mx); sm += s[j]; }
|
||
for (int j = 0; j < s_k; j++) s[j] /= sm;
|
||
for (int d = 0; d < hd; d++) {
|
||
float ov = 0.0f;
|
||
for (int int_j = 0; int_j < s_k; int_j++) ov += s[int_j] * bf16_to_f32_host(v[d * s_k + int_j]);
|
||
o_ref[d] = ov;
|
||
}
|
||
if (lse_ref) *lse_ref = logf(sm) + mx;
|
||
}
|
||
|
||
static int test_mha(int n_h) {
|
||
printf("\n=== Test MHA: n_h=%d, HD=%d, SK=%d ===\n", n_h, HD, SK);
|
||
const float SCALE = 1.0f / sqrtf((float)HD);
|
||
int pass = 1;
|
||
|
||
// Allocate host tensors: Q(n_h, hd), K(n_h, SK*hd), V(n_h, hd*SK), O(n_h, hd)
|
||
bf16_t* h_q = (bf16_t*)malloc(n_h * HD * sizeof(bf16_t));
|
||
bf16_t* h_k = (bf16_t*)malloc(n_h * SK * HD * sizeof(bf16_t));
|
||
bf16_t* h_v = (bf16_t*)malloc(n_h * HD * SK * sizeof(bf16_t));
|
||
bf16_t* h_o = (bf16_t*)calloc(n_h * HD, sizeof(bf16_t));
|
||
float* h_lse = (float*)calloc(n_h, sizeof(float));
|
||
|
||
srand(42);
|
||
for (int i = 0; i < n_h * HD; i++) h_q[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
for (int i = 0; i < n_h * SK * HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
for (int i = 0; i < n_h * HD * SK; i++) h_v[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
|
||
bf16_t *d_q, *d_k, *d_v, *d_o;
|
||
float *d_lse;
|
||
cudaMalloc(&d_q, n_h * HD * sizeof(bf16_t));
|
||
cudaMalloc(&d_k, n_h * SK * HD * sizeof(bf16_t));
|
||
cudaMalloc(&d_v, n_h * HD * SK * sizeof(bf16_t));
|
||
cudaMalloc(&d_o, n_h * HD * sizeof(bf16_t));
|
||
cudaMalloc(&d_lse, n_h * sizeof(float));
|
||
cudaMemcpy(d_q, h_q, n_h * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemcpy(d_k, h_k, n_h * SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemcpy(d_v, h_v, n_h * HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemset(d_o, 0, n_h * HD * sizeof(bf16_t));
|
||
cudaMemset(d_lse, 0, n_h * sizeof(float));
|
||
|
||
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.q_head_stride = HD; // T=1, stride = 1 * hd
|
||
params.q_batch_stride = n_h * HD;
|
||
params.k_head_stride = SK * HD; // each head has its own K
|
||
params.k_batch_stride = n_h * SK * HD;
|
||
params.v_head_stride = HD * SK; // each head has its own V
|
||
params.v_batch_stride = n_h * HD * SK;
|
||
params.o_head_stride = HD;
|
||
params.o_batch_stride = n_h * HD;
|
||
params.lse_head_stride = 1;
|
||
params.lse_batch_stride = n_h;
|
||
|
||
int smem = (4 + 8 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + SK*4 + 256 + 127) & ~127;
|
||
|
||
if (smem > 48 * 1024) {
|
||
cudaFuncSetAttribute(fmha_6warp_multihead_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||
}
|
||
|
||
fmha_6warp_multihead_kernel<HD><<<dim3(1, n_h, 1), 192, smem>>>(params);
|
||
|
||
cudaError_t launch_err = cudaGetLastError();
|
||
cudaError_t sync_err = cudaSuccess;
|
||
if (launch_err != cudaSuccess) {
|
||
printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err));
|
||
pass = 0; goto cleanup;
|
||
}
|
||
|
||
sync_err = cudaDeviceSynchronize();
|
||
if (sync_err != cudaSuccess) {
|
||
printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err));
|
||
pass = 0; goto cleanup;
|
||
}
|
||
|
||
cudaMemcpy(h_o, d_o, n_h * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);
|
||
cudaMemcpy(h_lse, d_lse, n_h * sizeof(float), cudaMemcpyDeviceToHost);
|
||
|
||
// Verify each head
|
||
for (int h = 0; h < n_h; h++) {
|
||
float o_ref[512]; // max HD
|
||
float lse_ref;
|
||
reference_attention(
|
||
h_q + h * HD, h_k + h * SK * HD, h_v + h * HD * SK,
|
||
o_ref, &lse_ref, HD, SK, SCALE
|
||
);
|
||
|
||
float cs = 0, na = 0, nb = 0;
|
||
for (int d = 0; d < HD; d++) {
|
||
float a = bf16_to_f32_host(h_o[h * HD + d]), b = o_ref[d];
|
||
if (fabsf(b) > 1e-4f) { cs += a*b; na += a*a; nb += b*b; }
|
||
}
|
||
cs /= (sqrtf(na) * sqrtf(nb) + 1e-10f);
|
||
|
||
float lse_err = fabsf(h_lse[h] - lse_ref) / (fabsf(lse_ref) + 1e-10f);
|
||
printf(" Head %2d: cos=%.8f lse_err=%.6f (kernel=%.6f ref=%.6f)\n",
|
||
h, cs, lse_err, h_lse[h], lse_ref);
|
||
if (cs < 0.999f) {
|
||
printf(" HEAD %d FAILED (cos=%.6f < 0.999)\n", h, cs);
|
||
pass = 0;
|
||
}
|
||
}
|
||
|
||
printf("MHA test %s\n", pass ? "PASSED" : "FAILED");
|
||
|
||
cleanup:
|
||
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse);
|
||
free(h_q); free(h_k); free(h_v); free(h_o); free(h_lse);
|
||
return pass;
|
||
}
|
||
|
||
static int test_mqa(int n_q, int n_kv) {
|
||
printf("\n=== Test MQA: n_q=%d, n_kv=%d, HD=%d, SK=%d ===\n", n_q, n_kv, HD, SK);
|
||
const float SCALE = 1.0f / sqrtf((float)HD);
|
||
int pass = 1;
|
||
int q_per_kv = n_q / n_kv;
|
||
|
||
// Q: (n_q, hd), K: (n_kv, SK*hd), V: (n_kv, hd*SK), O: (n_q, hd)
|
||
bf16_t* h_q = (bf16_t*)malloc(n_q * HD * sizeof(bf16_t));
|
||
bf16_t* h_k = (bf16_t*)malloc(n_kv * SK * HD * sizeof(bf16_t));
|
||
bf16_t* h_v = (bf16_t*)malloc(n_kv * HD * SK * sizeof(bf16_t));
|
||
bf16_t* h_o = (bf16_t*)calloc(n_q * HD, sizeof(bf16_t));
|
||
|
||
srand(123);
|
||
for (int i = 0; i < n_q * HD; i++) h_q[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
for (int i = 0; i < n_kv * SK * HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
for (int i = 0; i < n_kv * HD * SK; i++) h_v[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
|
||
bf16_t *d_q, *d_k, *d_v, *d_o;
|
||
cudaMalloc(&d_q, n_q * HD * sizeof(bf16_t));
|
||
cudaMalloc(&d_k, n_kv * SK * HD * sizeof(bf16_t));
|
||
cudaMalloc(&d_v, n_kv * HD * SK * sizeof(bf16_t));
|
||
cudaMalloc(&d_o, n_q * HD * sizeof(bf16_t));
|
||
cudaMemcpy(d_q, h_q, n_q * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemcpy(d_k, h_k, n_kv * SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemcpy(d_v, h_v, n_kv * HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemset(d_o, 0, n_q * HD * sizeof(bf16_t));
|
||
|
||
// For MQA: n_q heads each with different Q, but shared K/V
|
||
// We launch n_q CTAs. Each CTA reads its own Q head but the SAME K/V
|
||
// (k_head_stride=0, v_head_stride=0 for pure MQA)
|
||
// But we need to map Q head h to KV head (h / q_per_kv)
|
||
// The kernel doesn't know about the mapping — it just uses blockIdx.y as head_idx
|
||
// and k_head_stride to index K. For MQA, we set k_head_stride=0 so all CTAs read the same K.
|
||
// For GQA, we'd need a different approach (grouped launches).
|
||
//
|
||
// Pure MQA test: 1 KV head, all Q heads share it
|
||
if (n_kv != 1) {
|
||
printf("MQA test requires n_kv=1 for stride=0 trick, skipping\n");
|
||
pass = 1; goto cleanup;
|
||
}
|
||
|
||
{
|
||
FmhaParams params;
|
||
params.q = d_q;
|
||
params.k = d_k;
|
||
params.v = d_v;
|
||
params.o = d_o;
|
||
params.lse = nullptr; // skip LSE for MQA test
|
||
params.s_k = SK;
|
||
params.scale = SCALE;
|
||
params.head_dim = HD;
|
||
params.q_head_stride = HD;
|
||
params.q_batch_stride = n_q * HD;
|
||
params.k_head_stride = 0; // MQA: all heads share same K
|
||
params.k_batch_stride = SK * HD;
|
||
params.v_head_stride = 0; // MQA: all heads share same V
|
||
params.v_batch_stride = HD * SK;
|
||
params.o_head_stride = HD;
|
||
params.o_batch_stride = n_q * HD;
|
||
params.lse_head_stride = 0;
|
||
params.lse_batch_stride = 0;
|
||
|
||
int smem = (4 + 8 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + SK*4 + 256 + 127) & ~127;
|
||
if (smem > 48 * 1024) {
|
||
cudaFuncSetAttribute(fmha_6warp_multihead_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||
}
|
||
|
||
fmha_6warp_multihead_kernel<HD><<<dim3(1, n_q, 1), 192, smem>>>(params);
|
||
|
||
cudaError_t launch_err = cudaGetLastError();
|
||
cudaError_t sync_err2 = cudaSuccess;
|
||
if (launch_err != cudaSuccess) {
|
||
printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err));
|
||
pass = 0; goto cleanup;
|
||
}
|
||
sync_err2 = cudaDeviceSynchronize();
|
||
if (sync_err2 != cudaSuccess) {
|
||
printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err2));
|
||
pass = 0; goto cleanup;
|
||
}
|
||
|
||
cudaMemcpy(h_o, d_o, n_q * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);
|
||
|
||
// All Q heads share the same K/V (k_head[0], v_head[0])
|
||
for (int h = 0; h < n_q; h++) {
|
||
float o_ref[512];
|
||
reference_attention(
|
||
h_q + h * HD, h_k, h_v,
|
||
o_ref, nullptr, HD, SK, SCALE
|
||
);
|
||
|
||
float cs = 0, na = 0, nb = 0;
|
||
for (int d = 0; d < HD; d++) {
|
||
float a = bf16_to_f32_host(h_o[h * HD + d]), b = o_ref[d];
|
||
if (fabsf(b) > 1e-4f) { cs += a*b; na += a*a; nb += b*b; }
|
||
}
|
||
cs /= (sqrtf(na) * sqrtf(nb) + 1e-10f);
|
||
printf(" Q-head %2d (shared KV): cos=%.8f\n", h, cs);
|
||
if (cs < 0.999f) { printf(" HEAD %d FAILED\n", h); pass = 0; }
|
||
}
|
||
}
|
||
|
||
printf("MQA test %s\n", pass ? "PASSED" : "FAILED");
|
||
|
||
cleanup:
|
||
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o);
|
||
free(h_q); free(h_k); free(h_v); free(h_o);
|
||
return pass;
|
||
}
|
||
|
||
static int test_batched(int n_h, int batch_size) {
|
||
printf("\n=== Test Batched: n_h=%d, batch=%d, HD=%d, SK=%d ===\n", n_h, batch_size, HD, SK);
|
||
const float SCALE = 1.0f / sqrtf((float)HD);
|
||
int pass = 1;
|
||
|
||
int total_q = batch_size * n_h;
|
||
bf16_t* h_q = (bf16_t*)malloc(total_q * HD * sizeof(bf16_t));
|
||
bf16_t* h_k = (bf16_t*)malloc(total_q * SK * HD * sizeof(bf16_t));
|
||
bf16_t* h_v = (bf16_t*)malloc(total_q * HD * SK * sizeof(bf16_t));
|
||
bf16_t* h_o = (bf16_t*)calloc(total_q * HD, sizeof(bf16_t));
|
||
|
||
srand(999);
|
||
for (int i = 0; i < total_q * HD; i++) h_q[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
for (int i = 0; i < total_q * SK * HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
for (int i = 0; i < total_q * HD * SK; i++) h_v[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
||
|
||
bf16_t *d_q, *d_k, *d_v, *d_o;
|
||
cudaMalloc(&d_q, total_q * HD * sizeof(bf16_t));
|
||
cudaMalloc(&d_k, total_q * SK * HD * sizeof(bf16_t));
|
||
cudaMalloc(&d_v, total_q * HD * SK * sizeof(bf16_t));
|
||
cudaMalloc(&d_o, total_q * HD * sizeof(bf16_t));
|
||
cudaMemcpy(d_q, h_q, total_q * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemcpy(d_k, h_k, total_q * SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemcpy(d_v, h_v, total_q * HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
||
cudaMemset(d_o, 0, total_q * HD * sizeof(bf16_t));
|
||
|
||
FmhaParams params;
|
||
params.q = d_q;
|
||
params.k = d_k;
|
||
params.v = d_v;
|
||
params.o = d_o;
|
||
params.lse = nullptr;
|
||
params.s_k = SK;
|
||
params.scale = SCALE;
|
||
params.head_dim = HD;
|
||
params.q_head_stride = HD;
|
||
params.q_batch_stride = n_h * HD;
|
||
params.k_head_stride = SK * HD;
|
||
params.k_batch_stride = n_h * SK * HD;
|
||
params.v_head_stride = HD * SK;
|
||
params.v_batch_stride = n_h * HD * SK;
|
||
params.o_head_stride = HD;
|
||
params.o_batch_stride = n_h * HD;
|
||
params.lse_head_stride = 0;
|
||
params.lse_batch_stride = 0;
|
||
|
||
int smem = (4 + 8 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + SK*4 + 256 + 127) & ~127;
|
||
if (smem > 48 * 1024) {
|
||
cudaFuncSetAttribute(fmha_6warp_multihead_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||
}
|
||
|
||
fmha_6warp_multihead_kernel<HD><<<dim3(1, n_h, batch_size), 192, smem>>>(params);
|
||
|
||
cudaError_t launch_err = cudaGetLastError();
|
||
cudaError_t sync_err3 = cudaSuccess;
|
||
int checked = 0, failed = 0;
|
||
if (launch_err != cudaSuccess) {
|
||
printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err));
|
||
pass = 0; goto cleanup;
|
||
}
|
||
sync_err3 = cudaDeviceSynchronize();
|
||
if (sync_err3 != cudaSuccess) {
|
||
printf("CUDA ERROR: %s\n", cudaGetErrorString(sync_err3));
|
||
pass = 0; goto cleanup;
|
||
}
|
||
|
||
cudaMemcpy(h_o, d_o, total_q * HD * sizeof(bf16_t), cudaMemcpyDeviceToHost);
|
||
|
||
// Verify a sample of heads across batches
|
||
for (int b = 0; b < batch_size; b++) {
|
||
for (int h = 0; h < n_h; h++) {
|
||
int idx = b * n_h + h;
|
||
float o_ref[512];
|
||
reference_attention(
|
||
h_q + idx * HD,
|
||
h_k + idx * SK * HD,
|
||
h_v + idx * HD * SK,
|
||
o_ref, nullptr, HD, SK, SCALE
|
||
);
|
||
|
||
float cs = 0, na = 0, nb = 0;
|
||
for (int d = 0; d < HD; d++) {
|
||
float a = bf16_to_f32_host(h_o[idx * HD + d]), b2 = o_ref[d];
|
||
if (fabsf(b2) > 1e-4f) { cs += a*b2; na += a*a; nb += b2*b2; }
|
||
}
|
||
cs /= (sqrtf(na) * sqrtf(nb) + 1e-10f);
|
||
checked++;
|
||
if (cs < 0.999f) {
|
||
printf(" FAIL batch=%d head=%d: cos=%.6f\n", b, h, cs);
|
||
failed++;
|
||
}
|
||
}
|
||
}
|
||
printf(" Checked %d heads, %d failed\n", checked, failed);
|
||
pass = (failed == 0);
|
||
printf("Batched test %s\n", pass ? "PASSED" : "FAILED");
|
||
|
||
cleanup:
|
||
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o);
|
||
free(h_q); free(h_k); free(h_v); free(h_o);
|
||
return pass;
|
||
}
|
||
|
||
int main() {
|
||
printf("========================================\n");
|
||
printf("Multi-head FMHA test suite (HD=%d)\n", HD);
|
||
printf("========================================\n");
|
||
|
||
int all_pass = 1;
|
||
|
||
// Test 1: MHA with 4 heads
|
||
all_pass &= test_mha(4);
|
||
|
||
// Test 2: MHA with 8 heads (covers Pro's hd=128 with 128 heads in smaller test)
|
||
all_pass &= test_mha(8);
|
||
|
||
// Test 3: MQA: 4 Q heads sharing 1 KV head
|
||
all_pass &= test_mqa(4, 1);
|
||
|
||
// Test 4: Batched: 4 heads × 2 batch
|
||
all_pass &= test_batched(4, 2);
|
||
|
||
printf("\n========================================\n");
|
||
printf("Overall: %s\n", all_pass ? "ALL PASSED" : "SOME FAILED");
|
||
printf("========================================\n");
|
||
return all_pass ? 0 : 1;
|
||
}
|