From d5b0941f2e70b606179a73f016badd6a85c52a07 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 14:29:13 +0000
Subject: [PATCH] PV SS MMA with (128,128) P layout

---
 tests/unit/test_pv_ss_128.cu | 102 +++++++++++++++++++++++++++++++++++
 1 file changed, 102 insertions(+)
 create mode 100644 tests/unit/test_pv_ss_128.cu

diff --git a/tests/unit/test_pv_ss_128.cu b/tests/unit/test_pv_ss_128.cu
new file mode 100644
index 00000000..c1393dcb
--- /dev/null
+++ b/tests/unit/test_pv_ss_128.cu
@@ -0,0 +1,102 @@
+/**
+ * PV SS MMA with P as (128, 128) and K-tile at offset.
+ * Extending test_pv_ss.cu to use the full P matrix layout.
+ */
+
+#include <cuda_runtime.h>
+#include <cstdio>
+#include <cmath>
+#include <cstring>
+
+#include "dsv4/kernels/attention/fmha_common.cuh"
+#include "dsv4/kernels/attention/fmha_umma_desc.cuh"
+
+using namespace dsv4::kernels::attention;
+
+static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
+
+constexpr int HD = 16, SK = 128, BLOCK_MN = 128;
+
+__global__ void __launch_bounds__(128)
+test_pv_ss_128()
+{
+    const int tid = threadIdx.x, wid = tid / 32, lane = tid % 32;
+
+    extern __shared__ char sbuf[];
+    uint32_t* sTmemBase = (uint32_t*)sbuf;
+    // P as (128, 128) canonical
+    bf16_t* sP = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15);
+    // V as (16, 16) canonical
+    bf16_t* sV = (bf16_t*)(((uintptr_t)(sP + 128 * 128) + 127) & ~(uintptr_t)127);
+
+    // Fill P: (128, 128) canonical, row 0 = all 0.5
+    // (128, 128): CORES_MN=16, CORES_K=16
+    // core(g_mn, g_k) at g_k * 16 * 64 + g_mn * 64 + lr * 8 + lc
+    for (int i = tid; i < 128 * 128; i += 128) sP[i] = 0;
+    __syncthreads();
+    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);
+    }
+    __syncthreads();
+
+    // Fill V: (16, 16) canonical, all 1.0
+    for (int i = tid; i < 256; i += 128) sV[i] = 0;
+    __syncthreads();
+    for (int i = tid; i < 256; i += 128) {
+        int r = i / 16, c = i % 16;
+        int ck = c / 8, lc = c % 8;
+        int tmn = r / 8, lr = r % 8;
+        sV[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = f32_to_bf16(1.0f);
+    }
+    __syncthreads();
+
+    // TMEM alloc
+    if (wid == 1) tmem_alloc(__cvta_generic_to_shared(sTmemBase), 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();
+
+    // Read O from TMEM
+    if (wid == 0) {
+        float o_vals[HD];
+        for (int n = 0; n < HD / 8; n++) {
+            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]),
+                  "=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
+                : "r"(tb + n*8));
+            asm volatile("tcgen05.wait::ld.sync.aligned;");
+            if (lane == 0) for (int c=0;c<8;c++) o_vals[n*8+c] = tmp[c];
+        }
+        if (lane == 0) {
+            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");
+        }
+    }
+
+    if (wid == 0) tmem_dealloc(tb, 128);
+}
+
+int main() {
+    printf("=== PV SS MMA with (128,128) P ===\n");
+    // SMEM: tmem(4+12) + sP(32768) + sV(512) + align
+    int smem = (4+16 + 128*128*2 + 256*2 + 256 + 127) & ~127;
+    printf("SMEM: %d bytes\n", smem);
+    test_pv_ss_128<<<1, 128, smem>>>();
+    cudaError_t err = cudaDeviceSynchronize();
+    if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }
+    printf("PASS\n");
+    return 0;
+}