biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test: HD=64 random data, 4 K-tiles, accumulate

Browse Source
This commit is contained in:
biondizzle
2026-05-28 11:47:56 +00:00
parent 73bd21ce01
commit 4483539f01
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/unit/test_umma_qk_hd64.cu
View File

@@ -29,7 +29,7 @@ test_umma_qk_hd64(const bf16_t* q, const bf16_t* k,
{
const int tid = threadIdx.x;
const int wid = tid / 32, lane = tid % 32;
const int n_ktiles = 1; // hd / MMA_K; // Force 1 K-tile for debugging
const int n_ktiles = hd / MMA_K; // 4 for hd=64
// SMEM: sQ (128, HD) canonical + sK (128, HD) canonical
// Each K-tile of (128, 16) = 4096 bytes
@@ -194,8 +194,8 @@ int main() {
float* h_s_out = (float*)calloc(128 * 16, sizeof(float));
float* h_s_scalar = (float*)calloc(SK, sizeof(float));
for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16_host(1.0f);
for (int i = 0; i < SK*HD; i++) h_k[i] = f32_to_bf16_host(1.0f);
for (int d = 0; d < HD; d++) h_q[d] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
for (int i = 0; i < SK*HD; i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
bf16_t *d_q, *d_k; float *d_s_out, *d_s_scalar;
cudaMalloc(&d_q, HD*sizeof(bf16_t)); cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));