From 3c7d9d930350552a14d591b87dbf385ec61e715c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 11:34:45 +0000 Subject: [PATCH] =?UTF-8?q?test:=20apply=201/sqrt(HD)=20scale=20to=20MMA?= =?UTF-8?q?=20output=20=E2=80=94=204x=20was=20the=20scale=20factor,=20not?= =?UTF-8?q?=20a=20bug!?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit/test_umma_qk.cu | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index 053b0a94..fa4b37bd 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -86,7 +86,7 @@ test_umma_qk_hd16(const bf16_t* q, const bf16_t* k, if (lane == 0 && n < 1) { // Only first 8 cols for debug for (int c = 0; c < 8; c++) { if (out_row < 128) { - s_out[out_row * 8 + c] = tmp[c]; + s_out[out_row * 8 + c] = tmp[c] * scale; // Apply 1/sqrt(HD) scale } } } @@ -118,11 +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 + // Test with random data — MMA output is UNSCALED (no 1/sqrt(HD)) + // We apply the scale factor in the read srand(42); - 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));