From 6e5401df3bf3de727987413e96d71448a1fc10af Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 15:51:55 +0000
Subject: [PATCH] auto: pre-test commit

---
 tests/unit/test_tmem_all_lanes.cu | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/tests/unit/test_tmem_all_lanes.cu b/tests/unit/test_tmem_all_lanes.cu
index cb0dbafd..8aef8ba1 100644
--- a/tests/unit/test_tmem_all_lanes.cu
+++ b/tests/unit/test_tmem_all_lanes.cu
@@ -173,7 +173,7 @@ test_tmem_all_lanes(const bf16_t* q, const bf16_t* k, const bf16_t* v,
     // and figure out the mapping on the host.
 
     if (wid == 0) {
-        for (int n = 0; n < 8; n++) {  // 8 reads of 8 columns = 64 columns
+        for (int n = 0; n < 16; n++) {  // 16 reads of 8 columns = 128 columns
             float tmp[8];
             asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
                 : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
@@ -183,10 +183,9 @@ test_tmem_all_lanes(const bf16_t* q, const bf16_t* k, const bf16_t* v,
             // Each lane writes its 8 values to GMEM
             // For column (n*8 + c), lane i's value goes to position lane in that column
             // Layout: tmem_dump[col * 32 + lane] = lane's value for column col
-            // (32 lanes, 1 value per column per lane)
             for (int c = 0; c < 8; c++) {
                 int col = n * 8 + c;
-                tmem_dump[col * 32 + lane] = tmp[c];
+                if (col < 128) tmem_dump[col * 32 + lane] = tmp[c];
             }
         }
     }
@@ -213,9 +212,9 @@ int main() {
     cudaMalloc(&d_q, HD*sizeof(bf16_t));
     cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));
     cudaMalloc(&d_v, HD*SK*sizeof(bf16_t));
-    // 64 columns × 32 lanes = 2048 FP32
-    cudaMalloc(&d_tmem_dump, 64 * 32 * sizeof(float));
-    cudaMemset(d_tmem_dump, 0, 64 * 32 * sizeof(float));
+    // 128 columns × 32 lanes = 4096 FP32
+    cudaMalloc(&d_tmem_dump, 128 * 32 * sizeof(float));
+    cudaMemset(d_tmem_dump, 0, 128 * 32 * sizeof(float));
     cudaMemcpy(d_q, h_q, HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
     cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
     cudaMemcpy(d_v, h_v, HD*SK*sizeof(bf16_t), cudaMemcpyHostToDevice);
@@ -227,8 +226,8 @@ int main() {
     cudaError_t err = cudaDeviceSynchronize();
     if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }
 
-    float* h_dump = (float*)malloc(64 * 32 * sizeof(float));
-    cudaMemcpy(h_dump, d_tmem_dump, 64 * 32 * sizeof(float), cudaMemcpyDeviceToHost);
+    float* h_dump = (float*)malloc(128 * 32 * sizeof(float));
+    cudaMemcpy(h_dump, d_tmem_dump, 128 * 32 * sizeof(float), cudaMemcpyDeviceToHost);
 
     // Reference
     float s[SK];
@@ -255,7 +254,7 @@ int main() {
         float target = o_ref[d];
         int best_col = -1, best_lane = -1;
         float best_diff = 1e10f;
-        for (int col = 0; col < 64; col++) {
+        for (int col = 0; col < 128; col++) {
             for (int ln = 0; ln < 32; ln++) {
                 float val = h_dump[col * 32 + ln];
                 float diff = fabsf(val - target);
@@ -272,11 +271,11 @@ int main() {
 
     // Print column summary: for each column, which lanes have non-zero values?
     printf("\n=== Non-zero lanes per column ===\n");
-    for (int col = 0; col < 64; col++) {
-        int nz = 0;
-        for (int ln = 0; ln < 32; ln++) if (fabsf(h_dump[col*32+ln]) > 1e-6f) nz++;
-        if (nz > 0) printf("  col %2d: %d non-zero lanes\n", col, nz);
-    }
+        for (int col = 0; col < 128; col++) {
+            int nz = 0;
+            for (int ln = 0; ln < 32; ln++) if (fabsf(h_dump[col*32+ln]) > 1e-6f) nz++;
+            if (nz > 0) printf("  col %3d: %d non-zero lanes, lane0=%10.6f\n", col, nz, h_dump[col*32+0]);
+        }
 
     cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_tmem_dump);
     free(h_q); free(h_k); free(h_v); free(h_dump);