230 lines
7.6 KiB
Plaintext
230 lines
7.6 KiB
Plaintext
/**
|
|
* 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");
|
|
|
|
// Also test: 2 separate single-tile launches + Python merge
|
|
bf16_t *d_k1, *d_v1, *d_k2, *d_v2;
|
|
bf16_t *d_o1, *d_o2;
|
|
float *d_lse1, *d_lse2;
|
|
cudaMalloc(&d_k1, 128 * HD * 2);
|
|
cudaMalloc(&d_k2, 128 * HD * 2);
|
|
cudaMalloc(&d_v1, HD * 128 * 2);
|
|
cudaMalloc(&d_v2, HD * 128 * 2);
|
|
cudaMalloc(&d_o1, HD * 2);
|
|
cudaMalloc(&d_o2, HD * 2);
|
|
cudaMalloc(&d_lse1, 4);
|
|
cudaMalloc(&d_lse2, 4);
|
|
|
|
cudaMemcpy(d_k1, d_k, 128 * HD * 2, cudaMemcpyDeviceToDevice);
|
|
cudaMemcpy(d_k2, d_k + 128 * HD, 128 * HD * 2, cudaMemcpyDeviceToDevice);
|
|
cudaMemcpy(d_v1, d_v, HD * 128 * 2, cudaMemcpyDeviceToDevice);
|
|
cudaMemcpy(d_v2, d_v + HD * 128, HD * 128 * 2, cudaMemcpyDeviceToDevice);
|
|
|
|
// Run single-tile on first half
|
|
FmhaParams p1 = params;
|
|
p1.k = d_k1; p1.v = d_v1; p1.o = d_o1; p1.lse = d_lse1;
|
|
p1.s_k = 128; p1.n_kv_tiles = 1;
|
|
fmha_6warp_multihead_kernel<HD, 128><<<grid, block, smem>>>(p1);
|
|
|
|
// Run single-tile on second half
|
|
FmhaParams p2 = params;
|
|
p2.k = d_k2; p2.v = d_v2; p2.o = d_o2; p2.lse = d_lse2;
|
|
p2.s_k = 128; p2.n_kv_tiles = 1;
|
|
fmha_6warp_multihead_kernel<HD, 128><<<grid, block, smem>>>(p2);
|
|
|
|
cudaDeviceSynchronize();
|
|
|
|
// Read single-tile results
|
|
bf16_t h_o1[HD], h_o2[HD];
|
|
float h_lse1, h_lse2;
|
|
cudaMemcpy(h_o1, d_o1, HD * 2, cudaMemcpyDeviceToHost);
|
|
cudaMemcpy(h_o2, d_o2, HD * 2, cudaMemcpyDeviceToHost);
|
|
cudaMemcpy(&h_lse1, d_lse1, 4, cudaMemcpyDeviceToHost);
|
|
cudaMemcpy(&h_lse2, d_lse2, 4, cudaMemcpyDeviceToHost);
|
|
|
|
// Python merge: O = (exp(lse1)*O1 + exp(lse2)*O2) / (exp(lse1) + exp(lse2))
|
|
float e1 = expf(h_lse1), e2 = expf(h_lse2);
|
|
float h_o_merge[HD];
|
|
for (int d = 0; d < HD; d++) {
|
|
float o1 = hbf16_to_f32(h_o1[d]) * e1;
|
|
float o2 = hbf16_to_f32(h_o2[d]) * e2;
|
|
h_o_merge[d] = (o1 + o2) / (e1 + e2);
|
|
}
|
|
|
|
// Compare merge vs reference
|
|
float cos_merge = 0, nm_a = 0, nm_b = 0;
|
|
for (int d = 0; d < HD; d++) {
|
|
cos_merge += h_o_ref[d] * h_o_merge[d];
|
|
nm_a += h_o_ref[d] * h_o_ref[d];
|
|
nm_b += h_o_merge[d] * h_o_merge[d];
|
|
}
|
|
cos_merge /= sqrtf(nm_a * nm_b + 1e-30f);
|
|
|
|
printf(" Single-tile merge: cos=%.6f lse1=%.4f lse2=%.4f\n", cos_merge, h_lse1, h_lse2);
|
|
|
|
// Compare in-kernel multi-tile vs Python merge
|
|
float cos_vs_merge = 0; nm_a = 0; nm_b = 0;
|
|
for (int d = 0; d < HD; d++) {
|
|
float a = h_o_merge[d];
|
|
float b = hbf16_to_f32(h_o[d]);
|
|
cos_vs_merge += a * b;
|
|
nm_a += a * a;
|
|
nm_b += b * b;
|
|
}
|
|
cos_vs_merge /= sqrtf(nm_a * nm_b + 1e-30f);
|
|
printf(" In-kernel vs Python merge: cos=%.6f\n", cos_vs_merge);
|
|
|
|
cudaFree(d_k1); cudaFree(d_k2); cudaFree(d_v1); cudaFree(d_v2);
|
|
cudaFree(d_o1); cudaFree(d_o2); cudaFree(d_lse1); cudaFree(d_lse2);
|
|
|
|
if (cos >= 0.999990) { printf("PASS\n"); return 0; }
|
|
else { printf("FAIL (cos < 0.999990)\n"); return 1; }
|
|
}
|