From 013f370046a00d08a54a2d02c5135e41777294c3 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 11:32:56 +0000 Subject: [PATCH] test: all-ones data, expected S[0,j]=16.0 for every j --- tests/unit/test_umma_qk.cu | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index 8334d8e0..053b0a94 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -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));