v4.3 tag release update. (#2789)

examples/python/CuTeDSL/hopper/dense_gemm.py

@@ -38,6 +38,7 @@ import cutlass.cute as cute
 import cutlass.cute.testing as testing
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 from cutlass.cute.runtime import from_dlpack
 import cutlass.utils.hopper_helpers as sm90_utils
@@ -624,11 +625,11 @@ class HopperWgmmaGemmKernel:
             consumer_group=mainloop_pipeline_consumer_group,
             tx_count=tma_copy_bytes,
             cta_layout_vmnk=cta_layout_vmnk,
+            defer_sync=True,
         )
 
         #  Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         # ///////////////////////////////////////////////////////////////////////////////
         #  Generate smem tensor A/B
@@ -717,10 +718,7 @@ class HopperWgmmaGemmKernel:
         #  Cluster wait
         # ///////////////////////////////////////////////////////////////////////////////
         # cluster wait for barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            cute.arch.sync_threads()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
         # /////////////////////////////////////////////////////////////////////////////
         #  Prefetch
         # /////////////////////////////////////////////////////////////////////////////

examples/python/CuTeDSL/hopper/dense_gemm_persistent.py

@@ -37,6 +37,7 @@ import cutlass
 import cutlass.cute as cute
 import cutlass.cute.testing as testing
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.utils.hopper_helpers as sm90_utils
@@ -634,11 +635,11 @@ class HopperWgmmaGemmPersistentKernel:
             consumer_group=mainloop_pipeline_consumer_group,
             tx_count=tma_copy_bytes,
             cta_layout_vmnk=cute.make_layout((1, *cta_layout_mnk.shape)),
+            defer_sync=True,
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         # Generate smem tensor A/B
         sA = storage.sA.get_tensor(
@@ -718,10 +719,7 @@ class HopperWgmmaGemmPersistentKernel:
         k_tile_cnt = cute.size(gA_mkl, mode=[3])
 
         # Cluster wait for barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            cute.arch.sync_threads()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         is_dma_warp_group = warp_group_idx < self.num_dma_warp_groups
         if is_dma_warp_group:

examples/python/CuTeDSL/hopper/fmha.py

@@ -95,15 +95,18 @@ import cutlass.cute.testing as testing
 import cutlass.cute.nvgpu.warpgroup as warpgroup
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 from cutlass._mlir.dialects import math as _math
 
 import cutlass.utils.hopper_helpers as sm90_utils
 from cutlass.cute.runtime import from_dlpack
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(current_dir, ".."))
-from utils import fmha_helpers as fmha_utils
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+
+from helpers import fmha_helpers as fmha_utils
 
 
 class HopperFusedMultiHeadAttentionForward:
@@ -648,11 +651,8 @@ class HopperFusedMultiHeadAttentionForward:
         # We need this to guarantee that the Pipeline init is visible
         # To all producers and consumer blocks in the Cluster
         # and to finish smem init
-        if cute.size(self.cluster_shape_mnk) > 1:
-            cute.arch.cluster_arrive_relaxed()
-            cute.arch.cluster_wait()
-        else:
-            cute.arch.sync_threads()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
 
         if warp_idx == 0:
             cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q)
@@ -1646,6 +1646,7 @@ class HopperFusedMultiHeadAttentionForward:
             producer_group=load_q_producer_group,
             consumer_group=load_q_consumer_group,
             tx_count=self.tma_copy_q_bytes,
+            defer_sync=True,
         ).make_participants()
 
     def make_and_init_load_kv_pipeline(self, load_kv_mbar_ptr):
@@ -1663,6 +1664,7 @@ class HopperFusedMultiHeadAttentionForward:
             producer_group=load_kv_producer_group,
             consumer_group=load_kv_consumer_group,
             tx_count=self.tma_copy_kv_bytes,
+            defer_sync=True,
         ).make_participants()
 
     def make_and_init_tma_store_pipeline(self):
@@ -1686,6 +1688,7 @@ class HopperFusedMultiHeadAttentionForward:
                 pipeline.Agent.Thread,
                 self.num_threads_per_warp_group,
             ),
+            defer_sync=True,
         )
 
     @staticmethod