v4.4 tag release update. (#3032)

examples/python/CuTeDSL/hopper/dense_gemm.py

@@ -31,15 +31,12 @@ from typing import Tuple, Type
 import math
 import cuda.bindings.driver as cuda
 
-import torch
-
 import cutlass
 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
 
@@ -1006,7 +1003,10 @@ class HopperWgmmaGemmKernel:
                 tiled_copy_r2s, tRS_rD_out, tRS_sD[(None, None, None, epi_buffer)]
             )
 
-            cute.arch.fence_proxy("async.shared", space="cta")
+            cute.arch.fence_proxy(
+                "async.shared",
+                space="cta",
+            )
             # barrier for sync
             pipeline.sync(barrier_id=1)
 
@@ -1431,6 +1431,9 @@ def run(
     :rtype: float
     """
 
+    import torch
+    import cutlass.torch as cutlass_torch
+
     print("Running Hopper Dense GEMM with:")
     print(f"mnkl: {mnkl}")
     print(
@@ -1519,7 +1522,7 @@ def run(
 
     gemm = HopperWgmmaGemmKernel(acc_dtype, tile_shape_mn, cluster_shape_mn)
 
-    torch_stream = torch.cuda.Stream()
+    torch_stream = torch.cuda.current_stream()
     stream = cuda.CUstream(torch_stream.cuda_stream)
     # compile gemm kernel
     compiled_gemm = cute.compile(gemm, mA, mB, mC, stream)

examples/python/CuTeDSL/hopper/dense_gemm_persistent.py

@@ -31,14 +31,11 @@ from typing import Optional, Tuple, Type
 import math
 import cuda.bindings.driver as cuda
 
-import torch
-
 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
 
@@ -952,7 +949,10 @@ class HopperWgmmaGemmPersistentKernel:
                         tRS_sD[(None, None, None, epi_buffer)],
                     )
 
-                    cute.arch.fence_proxy("async.shared", space="cta")
+                    cute.arch.fence_proxy(
+                        "async.shared",
+                        space="cta",
+                    )
                     self.epilog_sync_barrier.arrive_and_wait()
 
                     gmem_coord = epi_tile_layout.get_hier_coord(epi_idx)
@@ -1465,6 +1465,8 @@ def run(
     :return: Execution time of the GEMM kernel in microseconds
     :rtype: float
     """
+    import torch
+    import cutlass.torch as cutlass_torch
 
     print("Running Hopper Persistent Dense GEMM with:")
     print(f"mnkl: {mnkl}")

examples/python/CuTeDSL/hopper/fmha.py

@@ -83,9 +83,6 @@ import sys
 import time
 from typing import Type, Tuple, Optional
 
-import torch
-
-
 import cuda.bindings.driver as cuda
 
 import cutlass
@@ -96,7 +93,6 @@ 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
@@ -598,6 +594,7 @@ class HopperFusedMultiHeadAttentionForward:
             k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner
         )
         # (MMA, MMA_K, MMA_D, PIPE)
+        # Adjust swizzle info to reuse smem
         sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner)
         sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer)
 
@@ -648,11 +645,19 @@ class HopperFusedMultiHeadAttentionForward:
 
         producer_warp_role = warp_idx % 4  # self.num_warps_per_warp_group
 
+        # Fence the mbarrier init to ensure all mbarrier initializations are visible
+        # to all threads. This is critical for FP8 performance - without this fence,
+        # the compiler may generate software polling loops instead of hardware waits.
+        cute.arch.mbarrier_init_fence()
+
         # 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
-        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
-        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
+        if cute.size(self.cluster_shape_mnk) > 1:
+            cute.arch.cluster_arrive_relaxed()
+            cute.arch.cluster_wait()
+        else:
+            cute.arch.sync_threads()
 
         if warp_idx == 0:
             cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q)
@@ -1164,7 +1169,10 @@ class HopperFusedMultiHeadAttentionForward:
                         tRS_sD[(None, None, None, epi_buffer, warp_group_idx - 1)],
                     )
 
-                    cute.arch.fence_proxy("async.shared", space="cta")
+                    cute.arch.fence_proxy(
+                        "async.shared",
+                        space="cta",
+                    )
                     pipeline.arrive_and_wait(
                         barrier_id=warp_group_idx,
                         num_threads=self.num_threads_per_warp_group,
@@ -1935,6 +1943,9 @@ def run(
     :return: Execution time of the FMHA kernel in microseconds
     :rtype: float
     """
+    import torch
+    import cutlass.torch as cutlass_torch
+
     print("Running Hopper SM90 FMHA test with:")
     print(f"  q_shape: {q_shape}")
     print(f"  k_shape: {k_shape}")