Pass tmem scalar fields as .ptr to TmemAllocator on SM100 (#2679)

The DSL now warns when a struct scalar is used directly as a pointer
("Use explicit struct.scalar.ptr for pointer instead"), so these fire
on every tmem_holding_buf / dealloc mbar access. Just pass .ptr like the
other SM100 kernels already do.

Author: pashu-cohere

flash_attn/cute/flash_bwd_mla_dq_dqv_sm100.py

@@ -354,7 +354,7 @@ class SharedStorage:
             mbar_ptr_KV_cpasync: cute.struct.MemRange[cutlass.Int64, self.num_stages_KV * 2]
             mbar_ptr_load_kv_epi: cute.struct.MemRange[cutlass.Int64, 2]
             # Tmem holding buffer
-            mbar_ptr_tmem_dealloc: cutlass.Int64
+            tmem_dealloc_mbar: cutlass.Int64
             tmem_holding_buf: cutlass.Int32
             # Clc pointers
             clc_ptr: cute.struct.Align[
@@ -569,11 +569,11 @@ def kernel(
         )
         # ---- Tensor memory dealloc barrier init ----
         tmem = utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.epilogue_warp_ids[0],
             is_two_cta=False,
-            two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
 
         # ---- Cluster arrive after barrier init ----

flash_attn/cute/flash_bwd_mla_sm100.py

@@ -293,7 +293,7 @@ def mbar_struct(num_stages):
                 self.num_stages_dPsum,
             ]
         )
-        mbar_ptr_tmem_dealloc_struct = Int64
+        tmem_dealloc_mbar_struct = Int64
         tmem_holding_buf_struct = Int32
 
         self.sched_stages = 1
@@ -314,7 +314,7 @@ class SharedStorage:
             mbar_ptr_dV_epi: mbar_ptr_dV_struct
             mbar_ptr_scaleP: mbar_ptr_scaleP_struct
             mbar_ptr_dPsum: mbar_ptr_dPsum_struct
-            mbar_ptr_tmem_dealloc: mbar_ptr_tmem_dealloc_struct
+            tmem_dealloc_mbar: tmem_dealloc_mbar_struct
             tmem_holding_buf: tmem_holding_buf_struct
             clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, clc_mbar_size]
             clc_response: cute.struct.MemRange[Int32, clc_response_size]
@@ -711,11 +711,11 @@ def kernel(
             num_threads=self.num_mma_threads + self.num_softmax_threads + self.num_epilogue_threads,
         )
         tmem = cutlass.utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.mma_warp_id,
             is_two_cta=self.use_2cta_instrs,
-            two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
 
         # ==== Prefetch TMA descriptors ====

flash_attn/cute/flash_bwd_sm100.py

@@ -782,7 +782,7 @@ class SharedStorage:
                     cutlass.Int64, self.dQaccum_reduce_stage // 2
                 ]
                 tmem_holding_buf: Int32
-                tmem_dealloc_mbar_ptr: cutlass.Int64
+                tmem_dealloc_mbar: cutlass.Int64
 
                 # 2-CTA
                 Qt_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage]
@@ -861,7 +861,7 @@ class SharedStorage:
                     cutlass.Int64, self.dQaccum_reduce_stage // 2
                 ]
                 tmem_holding_buf: Int32
-                tmem_dealloc_mbar_ptr: Int64
+                tmem_dealloc_mbar: Int64
 
                 sQ: cute.struct.Align[
                     cute.struct.MemRange[cute.Uint8, sQ_alloc_bytes],
@@ -1152,11 +1152,11 @@ def kernel(
             * len((self.mma_warp_id, *self.compute_warp_ids, *self.reduce_warp_ids)),
         )
         tmem = cutlass.utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.mma_warp_id,
             is_two_cta=self.use_2cta_instrs,
-            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
 
         # UMMA producers and AsyncThread consumers

flash_attn/cute/flash_fwd_mla_sm100.py

@@ -304,7 +304,7 @@ def mbar_struct(num_stages):
                 self.num_stages_bitmask,
             ]
         )
-        mbar_ptr_tmem_dealloc_struct = Int64
+        tmem_dealloc_mbar_struct = Int64
         tmem_holding_buf_struct = Int32
 
         self.sched_stages = 1
@@ -325,7 +325,7 @@ class SharedStorage:
             mbar_ptr_V_cpasync: mbar_ptr_V_struct
             mbar_ptr_sm_stats: mbar_sm_stats_struct
             mbar_ptr_bitmask: mbar_bitmask_struct
-            mbar_ptr_tmem_dealloc: mbar_ptr_tmem_dealloc_struct
+            tmem_dealloc_mbar: tmem_dealloc_mbar_struct
             tmem_holding_buf: tmem_holding_buf_struct
             clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, clc_mbar_size]
             clc_response: cute.struct.MemRange[Int32, clc_response_size]
@@ -820,11 +820,11 @@ def kernel(
             num_threads=self.num_mma_threads + self.num_softmax_threads + self.num_epilogue_threads,
         )
         tmem = cutlass.utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.mma_warp_id,
             is_two_cta=self.use_2cta_instrs,
-            two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
 
         # ==== Prefetch TMA descriptors ====

flash_attn/cute/flash_fwd_sm100.py

@@ -694,7 +694,7 @@ class SharedStorage:
             mbar_O_epi: cute.struct.MemRange[Int64, self.q_stage * 2]
             mbar_s0_s1_sequence: cute.struct.MemRange[Int64, 2 * 2]
             # Tmem dealloc cluster barrier
-            tmem_dealloc_mbar_ptr: Int64
+            tmem_dealloc_mbar: Int64
             # Tmem holding buffer
             tmem_holding_buf: Int32
             # Smem tensors
@@ -883,11 +883,11 @@ def kernel(
         )
         # Tensor memory dealloc barrier init
         tmem = cutlass.utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.mma_warp_id,
             is_two_cta=self.use_2cta_instrs,
-            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
 
         ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread)

flash_attn/cute/sm100_hd256_2cta_fmha_backward_dkdvkernel.py

@@ -660,7 +660,7 @@ class SharedStorage:
                 cutlass.Int64, self.mma_compute_dKdV_stage * 2
             ]
             tmem_holding_buf: cutlass.Int32
-            tmem_dealloc_mbar_ptr: cutlass.Int64
+            tmem_dealloc_mbar: cutlass.Int64
             clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
             clc_response: cute.struct.MemRange[Int32, 4]
             # Smem tensors
@@ -1029,11 +1029,11 @@ def dkdv_bwd(
         )
 
         tmem = utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.load_warp_id,
             is_two_cta=True,
-            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
 
         tmem.allocate(self.tmem_alloc_cols)

flash_attn/cute/sm100_hd256_2cta_fmha_backward_dqkernel.py

@@ -520,7 +520,7 @@ class SharedStorage:
                 cutlass.Int64, self.load_compute_sum_OdO_stage * 2
             ]
             # A CTA-wide "TMEM lifetime" barrier used to safely deallocate TMEM after all users finish.
-            tmem_dealloc_mbar_ptr: Int64
+            tmem_dealloc_mbar: Int64
             # Tmem holding buffer
             tmem_holding_buf: Int32
             # CLC pipeline barriers and response buffer
@@ -751,11 +751,11 @@ def kernel(
 
         # Tensor memory dealloc barrier init
         tmem = utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=self.tmem_alloc_barrier,
             allocator_warp_id=self.epilogue_warp_ids[0],
             is_two_cta=True,
-            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
         tmem.allocate(self.tmem_alloc_cols)
         tmem.wait_for_alloc()

flash_attn/cute/sm100_hd256_2cta_fmha_forward.py

@@ -510,7 +510,7 @@ class SharedStorage:
                 Int64, self.mma_corr_stage * 2
             ]  # mma_corr_{producer,consumer}
             # A CTA-wide "TMEM lifetime" barrier used to safely deallocate TMEM after all users finish.
-            tmem_dealloc_mbar_ptr: Int64
+            tmem_dealloc_mbar: Int64
             # Tmem holding buffer
             tmem_holding_buf: Int32
             # CLC pipeline barriers and response buffer
@@ -676,11 +676,11 @@ def kernel(
         ).make_participants()
         # Tensor memory dealloc barrier init
         tmem = utils.TmemAllocator(
-            storage.tmem_holding_buf,
+            storage.tmem_holding_buf.ptr,
             barrier_for_retrieve=self.tmem_alloc_barrier,
             allocator_warp_id=self.correction_warp_ids[0],
             is_two_cta=True,
-            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
         )
         tmem.allocate(self.tmem_alloc_cols)
         tmem.wait_for_alloc()