From 727c50945416ab6280009cd2fdb2327a027e269f Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 14:30:09 +0000
Subject: [PATCH] PV SS MMA with 8 K-tile accumulation

---
 tests/unit/test_pv_ss_128.cu | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/tests/unit/test_pv_ss_128.cu b/tests/unit/test_pv_ss_128.cu
index c1393dcb..8bc4fc19 100644
--- a/tests/unit/test_pv_ss_128.cu
+++ b/tests/unit/test_pv_ss_128.cu
@@ -37,7 +37,7 @@ test_pv_ss_128()
     for (int j = tid; j < 128; j += 128) {
         int core_k = j / 8, lc = j % 8;
         int dst_idx = core_k * 16 * 64 + 0 * 64 + 0 * 8 + lc;
-        sP[dst_idx] = f32_to_bf16(0.5f);
+        sP[dst_idx] = f32_to_bf16(0.5f);  // All 128 positions = 0.5
     }
     __syncthreads();
 
@@ -57,14 +57,20 @@ test_pv_ss_128()
     __syncthreads();
     uint32_t tb = *sTmemBase;
 
-    // PV SS MMA: A = P K-tile 0 at sP+0, B = V at sV
-    // K-tile 0 of (128,128): g_k=[0,1], offset=0, size=2048 BF16
-    uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sP), BLOCK_MN);
-    uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
-    uint32_t idesc = make_idesc(BLOCK_MN, HD);
-    if (tid == 0) umma_ss_f16(tb, dp, dv, idesc, false);
-    asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
-    __syncthreads();
+    // PV SS MMA: 8 K-tiles with accumulation
+    // K-tile kt of (128,128): g_k=[2*kt, 2*kt+1], offset = kt * 2048 BF16
+    // P row 0 = 0.5 for first 16 positions, 0.3 for positions 16-127
+    {
+        uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
+        uint32_t idesc = make_idesc(BLOCK_MN, HD);
+        for (int kt = 0; kt < 8; kt++) {
+            bf16_t* sp = sP + kt * 2048;
+            uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sp), BLOCK_MN);
+            if (tid == 0) umma_ss_f16(tb, dp, dv, idesc, kt > 0);
+            asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
+            __syncthreads();
+        }
+    }
 
     // Read O from TMEM
     if (wid == 0) {
@@ -82,7 +88,7 @@ test_pv_ss_128()
             printf("O[0,0..15] (raw): ");
             for (int d=0;d<HD;d++) printf("%.2f ", o_vals[d]);
             // Expected: sum(P[0,0..15] * V[j,0..15]) = 16 * 0.5 * 1.0 = 8.0 (raw, no MMA scale)
-            printf("\nExpected ~8.0 (raw, MMA scale unknown)\n");
+            printf("\nExpected ~64.0 (8 K-tiles × 16 values × 0.5 × 1.0, MMA scale 1.0)\n");
         }
     }