From 771799e112a854f2a0d53f3a4a2f19c2d0521011 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 06:35:50 +0000 Subject: [PATCH] =?UTF-8?q?FMHA=20SM100:=20Fix=20TMEM=20operations=20?= =?UTF-8?q?=E2=80=94=20uint32=5Ft=20registers,=20correct=20PTX=20syntax?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TMEM load/store uses b32 (uint32_t) registers, NOT float. Bitcast float↔uint32_t for FP32 TMEM values. TMEM alloc takes SMEM pointer (not a return value). TMEM column addressing: col + row_group * tmem_n. --- dsv4/kernels/attention/fmha_common.cuh | 46 +++++++---- .../kernels/attention/fmha_epilogue_sm100.cuh | 82 +++++++++---------- 2 files changed, 65 insertions(+), 63 deletions(-) diff --git a/dsv4/kernels/attention/fmha_common.cuh b/dsv4/kernels/attention/fmha_common.cuh index c9089a6d..e19e87c2 100644 --- a/dsv4/kernels/attention/fmha_common.cuh +++ b/dsv4/kernels/attention/fmha_common.cuh @@ -1,6 +1,9 @@ /** * DSV4 FMHA shared definitions — base header. * BF16 type, TMEM ops, warp reductions, constants. + * + * TMEM operations use uint32_t registers (b32), NOT float. + * Bitcast between float and uint32_t for FP32 TMEM values. */ #pragma once @@ -20,6 +23,10 @@ __device__ __forceinline__ float bf16_to_f32(bf16_t h) { float f; asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); return f; } +// Float ↔ uint32_t bitcast (for TMEM operations) +__device__ __forceinline__ uint32_t f32_to_u32(float f) { uint32_t u; memcpy(&u,&f,4); return u; } +__device__ __forceinline__ float u32_to_f32(uint32_t u) { float f; memcpy(&f,&u,4); return f; } + constexpr int WARP = 32; constexpr int NTHREADS = 192; constexpr int NWARPS = 6; @@ -31,31 +38,34 @@ __device__ __forceinline__ float wsum(float v) { for(int o=16;o>0;o>>=1) v+=__shfl_xor_sync(0xFFFFFFFF,v,o); return v; } -// TMEM operations -__device__ uint32_t tmem_alloc(int n) { - uint32_t b = 0; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 %0, [%1], %2;" - : "=r"(b) : "r"(0), "r"(n)); - return b; +// TMEM operations — using uint32_t registers per CUTLASS reference + +__device__ void tmem_alloc(uint32_t smem_ptr, int num_cols) { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_ptr), "r"(num_cols)); } -__device__ void tmem_dealloc(uint32_t b, int n) { - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" - :: "r"(b), "r"(n)); + +__device__ void tmem_dealloc(uint32_t tmem_ptr, int num_cols) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" + :: "r"(tmem_ptr), "r"(num_cols)); } -__device__ void tmem_load_col(uint32_t col, int row_group, - float& r0, float& r1, float& r2, float& r3) { - uint32_t addr = col + row_group; + +/** Load 16 rows × 256 bits from TMEM column. 4 uint32_t registers per thread. */ +__device__ void tmem_load(uint32_t col_addr, + uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) { asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" - : "=f"(r0), "=f"(r1), "=f"(r2), "=f"(r3) : "r"(addr)); + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(col_addr)); } -__device__ void tmem_store_col(uint32_t col, int row_group, - float r0, float r1, float r2, float r3) { - uint32_t addr = col + row_group; + +/** Store 16 rows × 256 bits to TMEM column. 4 uint32_t registers per thread. */ +__device__ void tmem_store(uint32_t col_addr, + uint32_t r0, uint32_t r1, uint32_t r2, uint32_t r3) { asm volatile("tcgen05.st.sync.aligned.16x256b.x1.b32 [%0], {%1, %2, %3, %4};" - :: "r"(addr), "f"(r0), "f"(r1), "f"(r2), "f"(r3)); + :: "r"(col_addr), "r"(r0), "r"(r1), "r"(r2), "r"(r3)); } + __device__ void tmem_fence() { asm volatile("tcgen05.fence.cta_group::1.sync.aligned;" ::: "memory"); } -} // namespace dsv4::kernels::attention +} // namespace diff --git a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh index 851d8d49..1c767cf3 100644 --- a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh +++ b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh @@ -1,11 +1,6 @@ /** * DSV4 FMHA Phase 2 — TMEM accumulator + one-way correction epilogue. - * - * Priority 2 from ROADMAP: - * TMEM → regs (tcgen05.ld) → normalize → BF16 → GMEM - * - * D1.5 fix: O rescale in REGISTERS between KV tiles. - * Unblocks: D2 multi-CTA, NVFP4-1.2 (register slot for FP4 pack). + * Uses uint32_t TMEM registers (matching CUTLASS PTX syntax). */ #pragma once #include "fmha_common.cuh" @@ -29,33 +24,37 @@ fmha_decode_tmem( const bf16_t* vb = v + batch*bstride_kv; bf16_t* oh = o + batch*bstride_o + head*HD; - // TMEM allocation for O accumulator - // Each tcgen05.ld reads 4 FP32 per column per row-group - // For HD=64: need ceil(64/4)=16 columns - const int tmem_o_cols = (HD + 3) / 4; - int tmem_n = 1; while(tmem_n < tmem_o_cols + 4) tmem_n *= 2; - uint32_t tb = 0; - if (wid==0 && lane==0) tb = tmem_alloc(tmem_n); - tb = __shfl_sync(0xFFFFFFFF, tb, 0); - const uint32_t to = tb; - - // SMEM for Q and row_sums + // SMEM for Q + row_sums + TMEM allocation extern __shared__ char sbuf[]; float* sQ = (float*)sbuf; float* sRowSums = (float*)(sbuf + HD*sizeof(float)); + // Use remaining SMEM for TMEM allocation (tcgen05.alloc maps it) + uint32_t tmem_smem_ptr = 0; + asm("cvta.to.shared.u32 %0, %1;" : "=r"(tmem_smem_ptr) : "l"(sbuf)); + + // TMEM column count: each tcgen05.ld reads 4 FP32 per column (16 rows × 256 bits) + // For T=1 decode, we only use row-group 0 (16 rows). Each column holds 4 FP32 values. + // So HD values need ceil(HD/4) columns. + const int tmem_o_cols = (HD + 3) / 4; + // Round up to power of 2 for TMEM allocation + int tmem_n = 1; while(tmem_n < tmem_o_cols) tmem_n *= 2; + + if (wid == 0 && lane == 0) tmem_alloc(tmem_smem_ptr, tmem_n); + __syncthreads(); // Wait for TMEM alloc for (int d=tid; d row_max) { float rescale = expf(row_max - new_max); - // D1.5 FIX: Rescale O in TMEM (TMEM → regs → multiply → TMEM) - // This is the one-way path using SAME tcgen05.ld + tcgen05.st pair + // D1.5: Rescale O in TMEM (TMEM → regs → multiply → TMEM) for (int col=0; col