v4.4 tag release update. (#3032)

examples/python/CuTeDSL/blackwell/fmha.py

@@ -33,16 +33,14 @@ import sys
 import time
 from typing import Type, Tuple, Union, Optional
 
-import torch
-import torch.nn.functional as F
 import cuda.bindings.driver as cuda
+import torch
 
 import cutlass
 import cutlass.cute as cute
 import cutlass.cute.nvgpu.tcgen05 as tcgen05
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
-import cutlass.torch as cutlass_torch
 import cutlass.utils.blackwell_helpers as sm100_utils
 import cutlass.cute.testing as testing
 from cutlass.cute.runtime import from_dlpack
@@ -174,8 +172,7 @@ class BlackwellFusedMultiHeadAttentionForward:
         self.load_warp_id = 13
         self.epilogue_warp_id = 14
         self.empty_warp_id = 15
-        SM100_TMEM_CAPACITY_COLUMNS = 512
-        self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
+        self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100")
 
         self.threads_per_warp = 32
         self.threads_per_cta = self.threads_per_warp * len(
@@ -1637,6 +1634,8 @@ class BlackwellFusedMultiHeadAttentionForward:
         :type atom_args: tuple
         :param tensor_args: Tuple containing softmax related tensors
         :type tensor_args: tuple
+        :param fused_mask: Compute trip counts and apply masking for attention blocks
+        :type fused_mask: fmha_utils.FusedMask
         :return: Updated state values (row_max, row_sum, and pipeline related arguments)
         :rtype: tuple
         """
@@ -1729,20 +1728,16 @@ class BlackwellFusedMultiHeadAttentionForward:
             tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)
         )
         for j in range(frg_cnt):
-            for k in range(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2):
-                tTMEM_LOADrS_frg[k, j], tTMEM_LOADrS_frg[k + 1, j] = (
-                    cute.arch.fma_packed_f32x2(
-                        (tTMEM_LOADrS_frg[k, j], tTMEM_LOADrS_frg[k + 1, j]),
-                        (scale, scale),
-                        (minus_row_max_scale, minus_row_max_scale),
-                    )
+            for k in cutlass.range(
+                cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True
+            ):
+                tTMEM_LOADrS_frg[k, j] = (
+                    tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale
                 )
                 tTMEM_LOADrS_frg[k, j] = cute.math.exp2(
                     tTMEM_LOADrS_frg[k, j], fastmath=True
                 )
-                tTMEM_LOADrS_frg[k + 1, j] = cute.math.exp2(
-                    tTMEM_LOADrS_frg[k + 1, j], fastmath=True
-                )
+
             s_vec = tTMEM_LOADrS_frg[None, j].load()
             tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype))
         # Sequence barrier arrive
@@ -1859,6 +1854,8 @@ class BlackwellFusedMultiHeadAttentionForward:
         :type s0_s1_sequence_pipeline: pipeline.PipelineAsync
         :param tile_sched_params: Parameters for tile scheduling
         :type tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams
+        :param fused_mask: Compute trip counts and apply masking for attention blocks
+        :type fused_mask: fmha_utils.FusedMask
         """
         tidx, _, _ = cute.arch.thread_idx()
         thread_idx = tidx % (
@@ -2204,11 +2201,8 @@ class BlackwellFusedMultiHeadAttentionForward:
             )
 
             cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i)
-            for j in range(0, cute.size(tTMrO_i), 2):
-                tTMrO_i[j], tTMrO_i[j + 1] = cute.arch.mul_packed_f32x2(
-                    (tTMrO_i[j], tTMrO_i[j + 1]),
-                    (scale, scale),
-                )
+            for j in cutlass.range(cute.size(tTMrO_i), vectorize=True):
+                tTMrO_i[j] = tTMrO_i[j] * scale
             cute.copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i)
 
     @cute.jit
@@ -2310,11 +2304,8 @@ class BlackwellFusedMultiHeadAttentionForward:
                 tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype
             )
             cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO)
-            for j in range(0, cute.size(tTMrO), 2):
-                tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2(
-                    (tTMrO[j], tTMrO[j + 1]),
-                    (scale, scale),
-                )
+            for j in range(cute.size(tTMrO), vectorize=True):
+                tTMrO[j] = tTMrO[j] * scale
             tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype)
             o_vec = tTMrO.load()
             tSMrO.store(o_vec.to(self.o_dtype))
@@ -2327,7 +2318,10 @@ class BlackwellFusedMultiHeadAttentionForward:
                 mLSE[row_idx + cuseqlen_q, blk_coord[2]] = lse
 
         # fence view async shared
-        cute.arch.fence_proxy("async.shared", space="cta")
+        cute.arch.fence_proxy(
+            "async.shared",
+            space="cta",
+        )
 
 
 def run(
@@ -2442,6 +2436,7 @@ def run(
     print(f"  iterations: {iterations}")
     print(f"  skip_ref_check: {skip_ref_check}")
     print(f"  use_cold_l2: {use_cold_l2}")
+    import cutlass.torch as cutlass_torch
 
     # Unpack parameters
     b, s_q, h_q, d = q_shape
@@ -2950,7 +2945,7 @@ def run(
         else:
             lse_tensor = None
 
-        return testing.JitArguments(
+        args = testing.JitArguments(
             q_tensor_workspace.iterator,
             k_tensor_workspace.iterator,
             v_tensor_workspace.iterator,
@@ -2970,6 +2965,15 @@ def run(
             ),
             current_stream,
         )
+        args.add_to_scope(
+            [
+                q_tensor_workspace,
+                k_tensor_workspace,
+                v_tensor_workspace,
+                o_tensor_workspace,
+            ]
+        )
+        return args
 
     workspace_count = 1
     if use_cold_l2: