[CLI] Recover ssd and blockwise group gemm perf (#3344)

* cggemm

* cggemm

* mggemm

* quick fix for 13.3

* fix ssd

* typo

* update

Author: myu-guo

examples/python/CuTeDSL/cute/blackwell/kernel/attention/mamba2_ssd/mamba2_ssd.py

@@ -142,10 +142,10 @@ def __init__(
         )
 
         # Number of registers used by each warp
-        self.num_regs_uniform_warps = 24
+        self.num_regs_uniform_warps = 32
         self.num_regs_pre_inter_warps = 168
         self.num_regs_pre_intra_warps = 208
-        self.num_regs_epilogue_warps = 112
+        self.num_regs_epilogue_warps = 104
 
         # Shared storage
         self.shared_storage = None

examples/python/CuTeDSL/cute/blackwell/kernel/blockwise_gemm/blockwise_gemm.py

@@ -2706,18 +2706,6 @@ def run(
     torch_stream = torch.cuda.current_stream()
     # Get the raw stream pointer as a CUstream
     current_stream = cuda.CUstream(torch_stream.cuda_stream)
-    # try to check CUDA version to decide the opt level
-    try:
-        from cutlass import CUDA_VERSION
-
-        opt_level = (
-            3
-            if CUDA_VERSION.major < 13
-            or (CUDA_VERSION.major == 13 and CUDA_VERSION.minor < 1)
-            else 2
-        )
-    except ImportError:
-        opt_level = 3
     # Compile gemm kernel
     compiled_gemm = cute.compile(
         gemm,
@@ -2728,7 +2716,7 @@ def run(
         sfb_tensor,
         max_active_clusters,
         current_stream,
-        options=f"--opt-level {opt_level}",
+        options=f"--opt-level=2" if cutlass.__version__[0:3]=="4.6" else "",
     )
 
     # Execution

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

@@ -1664,13 +1664,11 @@ def kernel(
                         tTR_rSFA_subtile = tTR_rSFA[(None, None, None, subtile_idx)]
                         tTR_rSFB_subtile = tTR_rSFB[(None, None, None, subtile_idx)]
 
-                        acc_vec = tTR_rAcc.load()
-                        final_vec = tTR_rAcc_subtile.load()
-                        scale_a = tTR_rSFA_subtile.load()
-                        scale_b = tTR_rSFB_subtile.load()
-                        scale = scale_a * scale_b
-                        final_vec = acc_vec * scale + final_vec
-                        tTR_rAcc_subtile.store(final_vec.to(self.acc_dtype))
+                        scale = cute.make_rmem_tensor(tTR_rSFA_subtile.shape, self.acc_dtype)
+                        for fma_idx in cutlass.range_constexpr(cute.size(scale)):
+                            scale[fma_idx] = tTR_rSFA_subtile[fma_idx] * tTR_rSFB_subtile[fma_idx]
+                        for fma_idx in cutlass.range_constexpr(cute.size(tTR_rAcc_subtile)):
+                            tTR_rAcc_subtile[fma_idx] = tTR_rAcc[fma_idx] * scale[fma_idx] + tTR_rAcc_subtile[fma_idx]
 
                     #
                     # Async arrive accumulator buffer empty
@@ -3041,4 +3039,3 @@ def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
         args.use_cold_l2,
         args.fixed_m,
     )
-    print("PASS")

examples/python/CuTeDSL/cute/blackwell/kernel/blockwise_gemm/masked_grouped_gemm.py

@@ -875,7 +875,6 @@ def kernel(
             cute.slice_(self.cta_tile_shape_mnk, (0, None, None)),
             (None, None, None),
         )
-        k_tile_cnt = cute.size(gA_mkl, mode=[3])
 
         #
         # Partition global tensor for TiledMMA_A/B/C
@@ -1081,6 +1080,7 @@ def kernel(
             is_valid_tile = cutlass.Boolean(1)
             is_valid_tile = tile_info[3] == 1
 
+            k_tile_cnt = cute.arch.make_warp_uniform(cute.size(gA_mkl, mode=[3]))
             while is_valid_tile:
                 mma_tile_coord_mnl = (
                     tile_info[0] // cute.size(tiled_mma.thr_id.shape),
@@ -1200,6 +1200,7 @@ def kernel(
             is_valid_tile = cutlass.Boolean(1)
             is_valid_tile = tile_info[3] == 1
 
+            k_tile_cnt = cute.arch.make_warp_uniform(cute.size(gA_mkl, mode=[3]))
             while is_valid_tile:
                 #
                 # Prepare the mask for scaleA/scaleB
@@ -1389,6 +1390,7 @@ def kernel(
             is_valid_tile = cutlass.Boolean(1)
             is_valid_tile = tile_info[3] == 1
 
+            k_tile_cnt = cute.arch.make_warp_uniform(cute.size(gA_mkl, mode=[3]))
             while is_valid_tile:
                 # Peek (try_wait) AB buffer full for k_tile = 0
                 ab_consumer_state.reset_count()
@@ -1579,6 +1581,7 @@ def kernel(
             is_valid_tile = cutlass.Boolean(1)
             is_valid_tile = tile_info[3] == 1
 
+            k_tile_cnt = cute.arch.make_warp_uniform(cute.size(gA_mkl, mode=[3]))
             while is_valid_tile:
                 # initialize the final accumulator
                 tTR_rAcc_final.fill(0.0)
@@ -1666,13 +1669,11 @@ def kernel(
                         tTR_rSFA_subtile = tTR_rSFA[(None, None, None, subtile_idx)]
                         tTR_rSFB_subtile = tTR_rSFB[(None, None, None, subtile_idx)]
 
-                        acc_vec = tTR_rAcc.load()
-                        final_vec = tTR_rAcc_subtile.load()
-                        scale_a = tTR_rSFA_subtile.load()
-                        scale_b = tTR_rSFB_subtile.load()
-                        scale = scale_a * scale_b
-                        final_vec = acc_vec * scale + final_vec
-                        tTR_rAcc_subtile.store(final_vec.to(self.acc_dtype))
+                        scale = cute.make_rmem_tensor(tTR_rSFA_subtile.shape, self.acc_dtype)
+                        for fma_idx in cutlass.range_constexpr(cute.size(scale)):
+                            scale[fma_idx] = tTR_rSFA_subtile[fma_idx] * tTR_rSFB_subtile[fma_idx]
+                        for fma_idx in cutlass.range_constexpr(cute.size(tTR_rAcc_subtile)):
+                            tTR_rAcc_subtile[fma_idx] = tTR_rAcc[fma_idx] * scale[fma_idx] + tTR_rAcc_subtile[fma_idx]
 
                     #
                     # Async arrive accumulator buffer empty
@@ -1828,6 +1829,7 @@ def kernel(
 
             num_prev_subtiles = cutlass.Int32(0)
 
+            k_tile_cnt = cute.arch.make_warp_uniform(cute.size(gA_mkl, mode=[3]))
             while is_valid_tile:
                 mma_tile_coord_mnl = (
                     tile_info[0] // cute.size(tiled_mma.thr_id.shape),