diff --git a/python/fmha_sm100/cute/src/common/utils.py b/python/fmha_sm100/cute/src/common/utils.py
index 10b2b1f..d57bb05 100644
--- a/python/fmha_sm100/cute/src/common/utils.py
+++ b/python/fmha_sm100/cute/src/common/utils.py
@@ -153,7 +153,7 @@ def make_tiled_copy_B(
 
 
 def mma_make_fragment_A(
-    smem: cute.Tensor, thr_mma: cute.core.ThrMma, swapAB: cutlass.Constexpr[bool] = False
+    smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False
 ) -> cute.Tensor:
     if const_expr(swapAB):
         return mma_make_fragment_B(smem, thr_mma)
@@ -162,7 +162,7 @@ def mma_make_fragment_A(
 
 
 def mma_make_fragment_B(
-    smem: cute.Tensor, thr_mma: cute.core.ThrMma, swapAB: cutlass.Constexpr[bool] = False
+    smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False
 ) -> cute.Tensor:
     if const_expr(swapAB):
         return mma_make_fragment_A(smem, thr_mma)
diff --git a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py
index f5950f3..caad594 100644
--- a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py
+++ b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd.py
@@ -1939,8 +1939,8 @@ def _wg_mma_issue(
         self,
         tiled_mma_qk: cute.TiledMma,
         tiled_mma_pv: cute.TiledMma,
-        thr0_qk: cute.core.ThrMma,
-        thr0_pv: cute.core.ThrMma,
+        thr0_qk: cute.ThrMma,
+        thr0_pv: cute.ThrMma,
         tStS: cute.Tensor,
         tOrP: cute.Tensor,
         sK: cute.Tensor,
diff --git a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py
index ba930a6..846df38 100644
--- a/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py
+++ b/python/fmha_sm100/cute/src/sm100/fwd/atten_fwd_nvfp4_kv.py
@@ -2191,8 +2191,8 @@ def _wg_mma_issue(
         self,
         tiled_mma_qk: cute.TiledMma,
         tiled_mma_pv: cute.TiledMma,
-        thr0_qk: cute.core.ThrMma,
-        thr0_pv: cute.core.ThrMma,
+        thr0_qk: cute.ThrMma,
+        thr0_pv: cute.ThrMma,
         tStS: cute.Tensor,
         tOrP: cute.Tensor,
         sK: cute.Tensor,
diff --git a/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py b/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py
index 8ccc2cb..738920a 100644
--- a/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py
+++ b/python/fmha_sm100/cute/src/sm100/fwd_decode/atten_fwd.py
@@ -1899,7 +1899,7 @@ def softmax_loop(
         tmem_load_atom_pre: cute.CopyAtom,
         tmem_store_atom_pre: cute.CopyAtom,
         tmem_store_vec_atom_pre: cute.CopyAtom,
-        thr_mma_qk_pre: cute.core.ThrMma,
+        thr_mma_qk_pre: cute.ThrMma,
         pipeline_s_p_o: pipeline.PipelineAsync,
         pipeline_p_lastsplit: pipeline.PipelineAsync,
         pipeline_sm_stats: pipeline.PipelineAsync,