From a391615f602ddb2714fde7c5a17c54939b204640 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 06:39:19 +0000
Subject: [PATCH] fix: uint64_t for SMEM pointer

---
 dsv4/kernels/attention/fmha_epilogue_sm100.cuh | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh
index 46b9e831..8dfc157f 100644
--- a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh
+++ b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh
@@ -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