v4.4 tag release update. (#3032)

examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py

@@ -30,13 +30,11 @@ import argparse
 from typing import Type, Tuple, Union
 
 import cuda.bindings.driver as cuda
-import torch
 
 import cutlass
 import cutlass.cute as cute
 import cutlass.cute.testing as testing
 from cutlass.cute.nvgpu import cpasync, tcgen05
-import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
@@ -80,7 +78,7 @@ This GEMM kernel supports the following features:
 
 This GEMM works as follows:
 1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
-2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations.
+2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy operations.
 2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction.
 3. EPILOGUE warp:
     - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld.
@@ -1034,7 +1032,10 @@ class BlockwiseContiguousGroupedGemmKernel:
                     )
 
                 # fence view async shared
-                cute.arch.fence_proxy("async.shared", space="cta")
+                cute.arch.fence_proxy(
+                    "async.shared",
+                    space="cta",
+                )
                 self.sched_sync_barrier.arrive_and_wait()
                 # commit tile info pipeline
                 tile_info_pipeline.producer_commit(tile_info_producer_state)
@@ -1150,7 +1151,10 @@ class BlockwiseContiguousGroupedGemmKernel:
                 for idx in cutlass.range(4, unroll_full=True):
                     tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                 is_valid_tile = tile_info[3] == 1
-                cute.arch.fence_proxy("async.shared", space="cta")
+                cute.arch.fence_proxy(
+                    "async.shared",
+                    space="cta",
+                )
                 tile_info_pipeline.consumer_release(tile_info_consumer_state)
                 tile_info_consumer_state.advance()
 
@@ -1322,7 +1326,10 @@ class BlockwiseContiguousGroupedGemmKernel:
                 for idx in cutlass.range(4, unroll_full=True):
                     tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                 is_valid_tile = tile_info[3] == 1
-                cute.arch.fence_proxy("async.shared", space="cta")
+                cute.arch.fence_proxy(
+                    "async.shared",
+                    space="cta",
+                )
                 tile_info_pipeline.consumer_release(tile_info_consumer_state)
                 tile_info_consumer_state.advance()
 
@@ -1479,7 +1486,10 @@ class BlockwiseContiguousGroupedGemmKernel:
                 for idx in cutlass.range(4, unroll_full=True):
                     tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                 is_valid_tile = tile_info[3] == 1
-                cute.arch.fence_proxy("async.shared", space="cta")
+                cute.arch.fence_proxy(
+                    "async.shared",
+                    space="cta",
+                )
                 tile_info_pipeline.consumer_release(tile_info_consumer_state)
                 tile_info_consumer_state.advance()
 
@@ -1715,7 +1725,10 @@ class BlockwiseContiguousGroupedGemmKernel:
                 for idx in cutlass.range(4, unroll_full=True):
                     tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                 is_valid_tile = tile_info[3] == 1
-                cute.arch.fence_proxy("async.shared", space="cta")
+                cute.arch.fence_proxy(
+                    "async.shared",
+                    space="cta",
+                )
                 tile_info_pipeline.consumer_release(tile_info_consumer_state)
                 tile_info_consumer_state.advance()
 
@@ -1884,7 +1897,10 @@ class BlockwiseContiguousGroupedGemmKernel:
                         tRS_sC[(None, None, None, c_buffer)],
                     )
                     # Fence and barrier to make sure shared memory store is visible to TMA store
-                    cute.arch.fence_proxy("async.shared", space="cta")
+                    cute.arch.fence_proxy(
+                        "async.shared",
+                        space="cta",
+                    )
                     self.epilog_sync_barrier.arrive_and_wait()
 
                     #
@@ -1914,7 +1930,10 @@ class BlockwiseContiguousGroupedGemmKernel:
                 for idx in cutlass.range(4, unroll_full=True):
                     tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
                 is_valid_tile = tile_info[3] == 1
-                cute.arch.fence_proxy("async.shared", space="cta")
+                cute.arch.fence_proxy(
+                    "async.shared",
+                    space="cta",
+                )
                 tile_info_pipeline.consumer_release(tile_info_consumer_state)
                 tile_info_consumer_state.advance()
 
@@ -2595,6 +2614,8 @@ class BlockwiseContiguousGroupedGemmKernel:
 
 
 def create_mask(num_groups, expect_m, fixed_m=False, m_aligned=128):
+    import torch
+
     valid_m = 0
     group_m_list = []
     gidx_mapping = []
@@ -2632,6 +2653,9 @@ def create_tensors(
     scale_dtype,
     fixed_m=False,
 ):
+    import torch
+    import cutlass.torch as cutlass_torch
+
     torch.manual_seed(1111)
 
     valid_m, group_m_list, _gidx_mapping = create_mask(l, m, fixed_m)
@@ -2702,6 +2726,9 @@ def run(
     fixed_m: bool = False,
     **kwargs,
 ):
+    import torch
+    import cutlass.torch as cutlass_torch
+
     """
     Prepare A/B/C tensors, launch GPU kernel, and reference checking.
     """