diff --git a/tests/unit/test_pv_accum.cu b/tests/unit/test_pv_accum.cu new file mode 100644 index 00000000..a2a9db17 --- /dev/null +++ b/tests/unit/test_pv_accum.cu @@ -0,0 +1,101 @@ +/** + * PV SS MMA: 2 K-tiles accumulated from (128,128) P. + * Detailed TMEM output read to diagnose the 56/64 split. + */ + +#include +#include +#include +#include + +#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_accum() +{ + const int tid = threadIdx.x, wid = tid / 32, lane = tid % 32; + + extern __shared__ char sbuf[]; + uint32_t* sTmemBase = (uint32_t*)sbuf; + bf16_t* sP = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15); + bf16_t* sV = (bf16_t*)(((uintptr_t)(sP + 128 * 128) + 127) & ~(uintptr_t)127); + + // Fill P: (128, 128) canonical, row 0 = all 0.5 + 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; + sP[core_k * 16 * 64 + 0 * 64 + 0 * 8 + lc] = 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, tmn = r / 8, lr = r % 8; + sV[ck * 2 * 64 + tmn * 64 + lr * 8 + lc] = f32_to_bf16(1.0f); + } + __syncthreads(); + + if (wid == 1) tmem_alloc(__cvta_generic_to_shared(sTmemBase), 128); + __syncthreads(); + uint32_t tb = *sTmemBase; + + // Zero TMEM (16 columns) + if (wid == 0) { + for (int n = 0; n < 2; n++) { + float z0=0,z1=0,z2=0,z3=0,z4=0,z5=0,z6=0,z7=0; + asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0],{%1,%2,%3,%4,%5,%6,%7,%8};" :: "r"(tb+n*8),"f"(z0),"f"(z1),"f"(z2),"f"(z3),"f"(z4),"f"(z5),"f"(z6),"f"(z7)); + } + tmem_fence_store(); + } + __syncthreads(); + + // 2 K-tiles, all accumulate=true + { + 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 < 2; 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, true); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + __syncthreads(); + } + } + + // Read ALL 16 TMEM columns, show row 0 and row 32 + if (wid == 0) { + // 32x32b.x8: lane 0 = row 0, lane 8 = row 32 + for (int group = 0; group < 2; group++) { + 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 + group*8)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + if (lane == 0) { printf("Row0 cols%d-%d: ", group*8, group*8+7); for(int c=0;c<8;c++) printf("%.1f ", tmp[c]); printf("\n"); } + if (lane == 8) { printf("Row32 cols%d-%d: ", group*8, group*8+7); for(int c=0;c<8;c++) printf("%.1f ", tmp[c]); printf("\n"); } + } + } + + if (wid == 0) tmem_dealloc(tb, 128); +} + +int main() { + printf("=== PV SS MMA Accumulation Debug ===\n"); + int smem = (4+16 + 128*128*2 + 256*2 + 256 + 127) & ~127; + test_pv_accum<<<1, 128, smem>>>(); + cudaError_t err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; } + return 0; +}