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

test: all-ones data, expected S[0,j]=16.0 for every j

Browse Source
This commit is contained in:
biondizzle
2026-05-28 11:32:56 +00:00
parent f5a0966afc
commit 013f370046
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/unit/test_umma_qk.cu
View File

@@ -51,7 +51,7 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k,
for (int i = tid; i < 4096; i += 128) sQ_pad[i] = 0;
__syncthreads();
// Descriptors
// Construct descriptors
uint32_t sQ_smem = __cvta_generic_to_shared(sQ);
uint32_t sK_smem = __cvta_generic_to_shared(sK);
uint64_t desc_q = make_umma_desc_kmajor_none(sQ_smem, 128);
@@ -118,9 +118,11 @@ int main() {
float* h_s_out = (float*)calloc(128*8, sizeof(float));
float* h_s_scalar = (float*)calloc(SK, sizeof(float));
// Test with ALL-ONES data: Q[0,d]=1.0, K[i,d]=1.0
// Expected: S[0,j] = sum(1*1, d=0..15) = 16.0 for all j
srand(42);
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);
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);
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));