FMHA SM100: Fix TMEM operations — uint32_t registers, correct PTX syntax
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.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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<HD; d+=NTHREADS) sQ[d] = bf16_to_f32(qh[d]);
|
||||
__syncthreads();
|
||||
|
||||
// Init TMEM O to zero
|
||||
// Initialize TMEM O to zero (all row groups)
|
||||
for (int col=tid; col<tmem_o_cols; col+=NTHREADS) {
|
||||
for (int rg=0; rg<8; rg++) tmem_store_col(to+col, rg, 0,0,0,0);
|
||||
for (int rg=0; rg<8; rg++) {
|
||||
tmem_store(col + rg * tmem_n, 0, 0, 0, 0);
|
||||
}
|
||||
}
|
||||
tmem_fence(); __syncthreads();
|
||||
__syncthreads();
|
||||
|
||||
float row_max = -INFINITY, row_sum = 0.0f;
|
||||
|
||||
// Single-thread softmax + P@V with TMEM O accumulation
|
||||
if (tid == 0) {
|
||||
for (int c=0; c<s_k; c++) {
|
||||
float s_val = 0.0f;
|
||||
@@ -67,15 +66,14 @@ fmha_decode_tmem(
|
||||
if (new_max > 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<tmem_o_cols; col++) {
|
||||
float r0,r1,r2,r3;
|
||||
tmem_load_col(to+col, 0, r0,r1,r2,r3);
|
||||
r0*=rescale; r1*=rescale; r2*=rescale; r3*=rescale;
|
||||
tmem_store_col(to+col, 0, r0,r1,r2,r3);
|
||||
uint32_t u0,u1,u2,u3;
|
||||
tmem_load(col, u0,u1,u2,u3);
|
||||
float r0=u32_to_f32(u0)*rescale, r1=u32_to_f32(u1)*rescale;
|
||||
float r2=u32_to_f32(u2)*rescale, r3=u32_to_f32(u3)*rescale;
|
||||
tmem_store(col, f32_to_u32(r0), f32_to_u32(r1), f32_to_u32(r2), f32_to_u32(r3));
|
||||
}
|
||||
tmem_fence();
|
||||
row_sum *= rescale; row_max = new_max;
|
||||
}
|
||||
|
||||
@@ -90,33 +88,27 @@ fmha_decode_tmem(
|
||||
float v2=(d2<HD)?bf16_to_f32(vb[d2*s_k+c]):0.0f;
|
||||
float v3=(d3<HD)?bf16_to_f32(vb[d3*s_k+c]):0.0f;
|
||||
|
||||
float r0,r1,r2,r3;
|
||||
tmem_load_col(to+col, 0, r0,r1,r2,r3);
|
||||
r0+=p_val*v0; r1+=p_val*v1; r2+=p_val*v2; r3+=p_val*v3;
|
||||
tmem_store_col(to+col, 0, r0,r1,r2,r3);
|
||||
uint32_t u0,u1,u2,u3;
|
||||
tmem_load(col, u0,u1,u2,u3);
|
||||
float r0=u32_to_f32(u0)+p_val*v0, r1=u32_to_f32(u1)+p_val*v1;
|
||||
float r2=u32_to_f32(u2)+p_val*v2, r3=u32_to_f32(u3)+p_val*v3;
|
||||
tmem_store(col, f32_to_u32(r0), f32_to_u32(r1), f32_to_u32(r2), f32_to_u32(r3));
|
||||
}
|
||||
tmem_fence();
|
||||
}
|
||||
sRowSums[0] = row_sum;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ================================================================
|
||||
// One-way Correction Epilogue
|
||||
// TMEM → regs → normalize → BF16 → GMEM
|
||||
//
|
||||
// This is the MoE epilogue pattern adapted for FMHA:
|
||||
// 1. tcgen05.ld: Load O from TMEM to registers (FP32)
|
||||
// 2. O[i] /= row_sum (normalize in registers)
|
||||
// 3. cvt.rn.bf16.f32: Cast to BF16
|
||||
// 4. st.global: Write to GMEM
|
||||
// One-way Correction Epilogue: TMEM → regs → normalize → BF16 → GMEM
|
||||
// ================================================================
|
||||
if (tid == 0) {
|
||||
float inv_sum = 1.0f / sRowSums[0];
|
||||
for (int col=0; col<tmem_o_cols; col++) {
|
||||
float r0,r1,r2,r3;
|
||||
tmem_load_col(to+col, 0, r0,r1,r2,r3);
|
||||
r0*=inv_sum; r1*=inv_sum; r2*=inv_sum; r3*=inv_sum;
|
||||
uint32_t u0,u1,u2,u3;
|
||||
tmem_load(col, u0,u1,u2,u3);
|
||||
float r0=u32_to_f32(u0)*inv_sum, r1=u32_to_f32(u1)*inv_sum;
|
||||
float r2=u32_to_f32(u2)*inv_sum, r3=u32_to_f32(u3)*inv_sum;
|
||||
|
||||
int d0=col*4+0, d1=col*4+1, d2=col*4+2, d3=col*4+3;
|
||||
if(d0<HD) oh[d0]=f32_to_bf16(r0);
|
||||
@@ -127,7 +119,7 @@ fmha_decode_tmem(
|
||||
}
|
||||
|
||||
if(lse_out && tid==0) lse_out[batch*gridDim.y+head] = logf(row_sum) + row_max;
|
||||
if(wid==0 && lane==0) tmem_dealloc(tb, tmem_n);
|
||||
if(wid==0 && lane==0) tmem_dealloc(tmem_smem_ptr, tmem_n);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
Reference in New Issue
Block a user