ad tcgen.ld.red support to sm103a arch

Author: drisspg

flash_attn/cute/flash_fwd_sm100.py

@@ -209,6 +209,11 @@ def __init__(
         # despite the literal `is_sm103` name.
         is_sm103 = self.arch.is_family_of(Arch.sm_103f)
         self.is_sm103 = is_sm103
+        # SM103 tcgen05.ld.red fuses the S TMEM load with a per-x32-tile max computed
+        # in the TMEM controller, making row max nearly free. Only valid when the
+        # loaded values are used unmodified, so score_mod (incl. softcap) disables it;
+        # masked iterations additionally fall back to the software max (see softmax_step).
+        self.use_ldred_rowmax = is_sm103 and self.score_mod is None
         # enable_ex2_emu is derived: True if tuning config has freq > 0, else fallback to default logic
         _default_enable_ex2_emu = (self.head_dim_padded <= 128 or (self.head_dim_padded == 192 and self.use_2cta_instrs and not self.is_causal and not self.is_local)) and not is_sm103
         self.enable_ex2_emu = _default_enable_ex2_emu
@@ -1920,9 +1925,12 @@ def softmax_loop(
         )
         tStP = cute.make_tensor(tSAcc.iterator + self.tmem_s_to_p_offset, tStP_layout)
 
-        tmem_load_atom = cute.make_copy_atom(
-            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype
+        tmem_load_op = (
+            tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(32))
+            if const_expr(self.use_ldred_rowmax)
+            else tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32))
         )
+        tmem_load_atom = cute.make_copy_atom(tmem_load_op, self.qk_acc_dtype)
         thr_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tSAcc).get_slice(tidx)
         tStS_t2r = thr_tmem_load.partition_S(tSAcc)  # (((32,32),1),1,4)
 
@@ -2169,6 +2177,7 @@ def softmax_loop(
                         else:
                             mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase = softmax_step(
                                 mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase, n_block,
+                                mask_is_noop=True,
                             )
                     # Separate iterations with local masking on the left
                     if const_expr(self.is_local and block_info.window_size_left is not None):
@@ -2256,6 +2265,9 @@ def softmax_step(
         head_divmod=None,
         mask_fn: Optional[Callable] = None,
         is_first: bool = False,
+        # True when this call site applies no masking (no causal/local edge, no
+        # mask_mod, no seqlen boundary) — gates use of the tcgen05.ld.red hardware max.
+        mask_is_noop: cutlass.Constexpr[bool] = False,
     ) -> Tuple[cute.Int32, cute.Int32, cute.Int32]:
         """Perform a single step of the softmax computation on a block of attention scores.
 
@@ -2284,7 +2296,16 @@ def softmax_step(
         # Wait for Si
         pipeline_s_p_o.consumer_wait_w_index_phase(stage, mma_si_consumer_phase)
         tSrS_t2r = cute.make_rmem_tensor(thr_tmem_load.partition_D(tScS).shape, self.qk_acc_dtype)
-        cute.copy(thr_tmem_load, tStS_t2r, tSrS_t2r)
+        hw_row_max = Float32(-Float32.inf)
+        if const_expr(self.use_ldred_rowmax):
+            # ld.red returns each x32 tile's max in an extra register alongside the
+            # 32 values; the reduction tensor collapses the V-mode to a single element.
+            tSrS_red = cute.make_rmem_tensor(((1, 1), *tSrS_t2r.shape[1:]), self.qk_acc_dtype)
+            cute.copy(thr_tmem_load, tStS_t2r, (tSrS_t2r, tSrS_red))
+            for i in cutlass.range_constexpr(cute.size(tSrS_red.shape)):
+                hw_row_max = cute.arch.fmax(hw_row_max, tSrS_red[i])
+        else:
+            cute.copy(thr_tmem_load, tStS_t2r, tSrS_t2r)
         # tSrS_t2r = copy_utils.load_t2r(thr_tmem_load, tScS_shape, tStS_t2r)
         if cutlass.const_expr(self.score_mod is not None):
             self.apply_score_mod(
@@ -2304,7 +2325,12 @@ def softmax_step(
 
         if const_expr(mask_fn is not None):
             mask_fn(tSrS_t2r, n_block=n_block)
-        row_max, acc_scale = softmax.update_row_max(tSrS_t2r.load(), is_first)
+        # The hardware max is only valid when masking left the values untouched;
+        # masked iterations reduce over the post-mask values in software.
+        if const_expr(self.use_ldred_rowmax and mask_is_noop):
+            row_max, acc_scale = softmax.update_row_max_precomputed(hw_row_max, is_first)
+        else:
+            row_max, acc_scale = softmax.update_row_max(tSrS_t2r.load(), is_first)
 
         if const_expr(not is_first):
             # tSrScale_r2t = cute.make_rmem_tensor(thr_tmem_store_scale.partition_S(tScScale).shape, Float32)

flash_attn/cute/softmax.py

@@ -298,6 +298,18 @@ def update_row_max_from_local(
         self.row_max[0] = row_max_new
         return row_max_safe, acc_scale
 
+    @cute.jit
+    def update_row_max_precomputed(
+        self, hw_row_max: Float32, is_first: int
+    ) -> Tuple[Float32, Float32]:
+        """Row max already reduced in hardware (SM103 tcgen05.ld.red): skip the
+        software fmax tree — the TMEM controller computed the max during the S load."""
+        if cutlass.const_expr(is_first):
+            row_max_new = hw_row_max
+        else:
+            row_max_new = cute.arch.fmax(hw_row_max, self.row_max[0])
+        return self.update_row_max_from_local(row_max_new, is_first)
+
     @cute.jit
     def update_row_max(self, acc_S_row: cute.TensorSSA, is_first: int) -> Tuple[Float32, Float32]:
         if cutlass.const_expr(is_first):