From a3542f9c520dae3c24a1b8cc82b97a97af517ada Mon Sep 17 00:00:00 2001
From: HuPengsheet <1554414201@qq.com>
Date: Wed, 17 Jun 2026 00:48:35 +0800
Subject: [PATCH] fix code

---
 .../cute/blackwell/kernel/attention/fmha/fmha.py  | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py b/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py
index d7887f009d..4846cd7f44 100644
--- a/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py
+++ b/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py
@@ -284,6 +284,11 @@ def __init__(
         self.num_regs_other = 32
         self.buffer_align_bytes = 1024
         self.arch = BaseDSL._get_dsl().get_arch_enum()
+        # SM101 (e.g. Thor) shares the SM100 TMEM load path; other arches use LdRed.
+        self.uses_sm100_softmax_tmem_load = (
+            (self.arch >= Arch.sm_100 and self.arch <= Arch.sm_100f)
+            or (self.arch >= Arch.sm_101 and self.arch <= Arch.sm_101f)
+        )
 
         if self.arch >= Arch.sm_103:
             assert self.enable_ex2_emulation == False, (
@@ -2260,9 +2265,7 @@ def softmax_step(
         old_row_max = row_max
         skip_softmax = cutlass.Boolean(False)
         if whether_apply_mask:
-            if cutlass.const_expr(
-                self.arch >= Arch.sm_100 and self.arch <= Arch.sm_100f
-            ):
+            if cutlass.const_expr(self.uses_sm100_softmax_tmem_load):
                 cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS)
             else:
                 tTMEM_LOADrMax = cute.make_rmem_tensor(
@@ -2308,9 +2311,7 @@ def softmax_step(
             inplace_producer.commit()
             inplace_producer.advance()
         else:
-            if cutlass.const_expr(
-                self.arch >= Arch.sm_100 and self.arch <= Arch.sm_100f
-            ):
+            if cutlass.const_expr(self.uses_sm100_softmax_tmem_load):
                 cute.copy(
                     tiled_tmem_load,
                     tTMEM_LOADtS[None, 0, None, None],
@@ -2731,7 +2732,7 @@ def softmax(
         )
         tmem_p_offset = self.tmem_p0_offset if stage == 0 else self.tmem_p1_offset
         tStS_P = cute.make_tensor(tStS.iterator + tmem_p_offset, tStS_P_layout)
-        if cutlass.const_expr(self.arch >= Arch.sm_100 and self.arch <= Arch.sm_100f):
+        if cutlass.const_expr(self.uses_sm100_softmax_tmem_load):
             tmem_load_atom = cute.make_copy_atom(
                 tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)),
                 self.qk_acc_dtype,