biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: uint64_t for SMEM pointer

Browse Source
This commit is contained in:
biondizzle
2026-05-28 06:39:19 +00:00
parent b4779e3f48
commit a391615f60
1 changed files with 3 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
dsv4/kernels/attention/fmha_epilogue_sm100.cuh
View File

@@ -32,9 +32,7 @@ fmha_decode_tmem(
// TMEM allocation: pass SMEM pointer for bookkeeping
// The actual TMEM columns are addressed by index (0, 1, 2, ...)
// We use sbuf's SMEM address (converted to u32) for the alloc call
uint32_t tmem_smem_ptr;
// Use NVIDIA's __cvta_generic_to_shared intrinsic or inline asm
// cvta.to.shared requires: cvta.to.shared.u64 dst, src (64-bit pointers)
uint64_t tmem_smem_ptr;
asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(tmem_smem_ptr) : "l"(sbuf));
// TMEM column count: each tcgen05.ld reads 4 FP32 per column (16 rows × 256 bits)
@@ -44,7 +42,7 @@ fmha_decode_tmem(
// 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);
if (wid == 0 && lane == 0) tmem_alloc((uint32_t)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]);
@@ -124,7 +122,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(tmem_smem_ptr, tmem_n);
if(wid==0 && lane==0) tmem_dealloc((uint32_t)tmem_smem_ptr, tmem_n);
}
} // namespace