fix(cute_dsl): migrate off APIs removed in nvidia-cutlass-dsl 4.6 (keeps 4.5.x compat)

flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py

@@ -1191,8 +1191,8 @@ def empty_warp(
     @cute.jit
     def load(
         self,
-        thr_mma_qk: cute.core.ThrMma,
-        thr_mma_pv: cute.core.ThrMma,
+        thr_mma_qk: cute.ThrMma,
+        thr_mma_pv: cute.ThrMma,
         mQ: cute.Tensor,
         mK: cute.Tensor,
         mV: cute.Tensor,
@@ -1388,8 +1388,8 @@ def load(
     @cute.jit
     def mma(
         self,
-        tiled_mma_qk: cute.core.ThrMma,
-        tiled_mma_pv: cute.core.ThrMma,
+        tiled_mma_qk: cute.ThrMma,
+        tiled_mma_pv: cute.ThrMma,
         sQ: cute.Tensor,
         sK: cute.Tensor,
         sV: cute.Tensor,
@@ -1666,7 +1666,7 @@ def softmax_loop(
         stage: int | Int32,
         softmax_scale_log2: Float32,
         softmax_scale: Float32,
-        thr_mma_qk: cute.core.ThrMma,
+        thr_mma_qk: cute.ThrMma,
         tStS: cute.Tensor,  # ((TILE_M, TILE_N), 1, 1, q_stage)
         sScale: cute.Tensor,
         mLSE: Optional[cute.Tensor],
@@ -1903,7 +1903,7 @@ def softmax_step(
         sm_stats_producer_phase: Int32,
         s0_s1_sequence_phase: Int32,
         softmax: SoftmaxSm100,
-        thr_mma_qk: cute.core.ThrMma,
+        thr_mma_qk: cute.ThrMma,
         pipeline_s_p_o: pipeline.PipelineAsync,
         pipeline_p_lastsplit: pipeline.PipelineAsync,
         pipeline_sm_stats: pipeline.PipelineAsync,
@@ -1948,7 +1948,7 @@ def softmax_step(
 
         # Wait for Si
         pipeline_s_p_o.consumer_wait_w_index_phase(stage, mma_si_consumer_phase)
-        tSrS_t2r = cute.make_fragment(
+        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)
@@ -1964,7 +1964,7 @@ def softmax_step(
         sm_stats_barrier.arrive_w_index(index=stage * 4 + warp_idx)
 
         softmax.scale_subtract_rowmax(tSrS_t2r, row_max)
-        tSrP_r2t_f32 = cute.make_fragment(
+        tSrP_r2t_f32 = cute.make_rmem_tensor(
             thr_tmem_store.partition_S(cute.make_identity_tensor(tScP_shape)).shape,
             Float32,
         )
@@ -2011,8 +2011,8 @@ def softmax_step(
     @cute.jit
     def correction_loop(
         self,
-        thr_mma_qk: cute.core.ThrMma,
-        thr_mma_pv: cute.core.ThrMma,
+        thr_mma_qk: cute.ThrMma,
+        thr_mma_pv: cute.ThrMma,
         tStS: cute.Tensor,
         tOtO: cute.Tensor,
         sScale: cute.Tensor,
@@ -2040,30 +2040,6 @@ def correction_loop(
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4
         mma_tile_coord_v = thr_mma_qk.thr_idx
 
-        tScS = thr_mma_qk.partition_C(cute.make_identity_tensor(self.mma_tiler_qk[:2]))
-        tStScale_layout = cute.composition(
-            tStS.layout, cute.make_layout((self.m_block_size, 1))
-        )
-        tStScales = tuple(
-            cute.make_tensor(
-                tStS.iterator + self.tmem_vec_offset[stage], tStScale_layout
-            )
-            for stage in range(self.s_stage)
-        )
-        tScScale = cute.composition(tScS, cute.make_layout((self.m_block_size, 1)))
-        tmem_load_v_atom = cute.make_copy_atom(
-            tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(1)), self.qk_acc_dtype
-        )
-        thr_tmem_load_vec = tcgen05.make_tmem_copy(
-            tmem_load_v_atom, tStScales[0]
-        ).get_slice(tidx)
-
-        [
-            thr_tmem_load_vec.partition_S(tStScales[stage])
-            for stage in range(self.s_stage)
-        ]
-        tSrScale_t2r_shape = thr_tmem_load_vec.partition_D(tScScale).shape
-
         # First iter: no correction is required
         # Notify mma warp that O has been rescaled
         for stage in cutlass.range(self.s_stage):
@@ -2117,7 +2093,6 @@ def correction_loop(
                 sm_stats_barrier.arrive_and_wait_w_index(index=1 * 4 + warp_idx)
                 sm_stats_consumer_phase ^= 1
 
-                cute.make_fragment(tSrScale_t2r_shape, Float32)
                 # q_stage=1 correction loop
                 if const_expr(mBlockNums is not None):
                     block_iter_count = (
@@ -2330,7 +2305,7 @@ def correction_loop(
     @cute.jit
     def correction_rescale(
         self,
-        thr_mma: cute.core.ThrMma,
+        thr_mma: cute.ThrMma,
         tOtO: cute.Tensor,
         tidx: Int32,
         scale: Float32,
@@ -2370,9 +2345,8 @@ def correction_rescale(
         tOtO_r2t = thr_tmem_store.partition_D(tOtO_i)
 
         frg_count = self.head_dim_v_padded // corr_tile_size
-        tOrO_frg = cute.make_fragment((tOrO_t2r_shape, frg_count), self.pv_acc_dtype)
         for i in cutlass.range_constexpr(frg_count):
-            tOrO_frg = cute.make_fragment(tOrO_t2r_shape, self.pv_acc_dtype)
+            tOrO_frg = cute.make_rmem_tensor(tOrO_t2r_shape, self.pv_acc_dtype)
             tOtO_t2r_i = cute.make_tensor(
                 tOtO_t2r.iterator + i * corr_tile_size, tOtO_t2r.layout
             )
@@ -2390,7 +2364,7 @@ def correction_rescale(
     @cute.jit
     def correction_epilogue(
         self,
-        thr_mma: cute.core.ThrMma,
+        thr_mma: cute.ThrMma,
         tOtO: cute.Tensor,
         tidx: Int32,
         stage: Int32,
@@ -2447,7 +2421,7 @@ def correction_epilogue(
         ):
             tOtO_t2r_i = tOtO_t2r[None, 0, 0, i]
             tOsO_r2s_i = tOsO_s2r[None, 0, 0, i]
-            tOrO_frg = cute.make_fragment(
+            tOrO_frg = cute.make_rmem_tensor(
                 tOcO_t2r[None, 0, 0, i].shape, self.pv_acc_dtype
             )
             cute.copy(tiled_tmem_load, tOtO_t2r_i, tOrO_frg)
@@ -2476,7 +2450,7 @@ def correction_epilogue(
     @cute.jit
     def correction_epilogue_combine(
         self,
-        thr_mma: cute.core.ThrMma,
+        thr_mma: cute.ThrMma,
         tOtO0: cute.Tensor,
         tOtO1: cute.Tensor,
         tidx: Int32,
@@ -2543,8 +2517,8 @@ def correction_epilogue_combine(
             tOtO1_t2r_i = tOtO1_t2r[None, 0, 0, i]
             tOsO_r2s_i = tOsO_s2r[None, 0, 0, i]
             frg_shape = tOcO_t2r[None, 0, 0, i].shape
-            tOrO0_frg = cute.make_fragment(frg_shape, self.pv_acc_dtype)
-            tOrO1_frg = cute.make_fragment(frg_shape, self.pv_acc_dtype)
+            tOrO0_frg = cute.make_rmem_tensor(frg_shape, self.pv_acc_dtype)
+            tOrO1_frg = cute.make_rmem_tensor(frg_shape, self.pv_acc_dtype)
             # When both scales are 0 (empty tile), skip tmem reads to avoid 0*NaN=NaN.
             is_zero_output = scale0 == Float32(0.0) and scale1 == Float32(0.0)
             if not is_zero_output:

flashinfer/cute_dsl/sparse/blk128/pack_gqa.py

@@ -93,7 +93,7 @@ def compute_ptr(
         num_threads: cutlass.Constexpr[int],
     ):
         num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row)
-        tPrPtr = cute.make_fragment(num_ptr_per_thread, cutlass.Int64)
+        tPrPtr = cute.make_rmem_tensor(num_ptr_per_thread, cutlass.Int64)
         for i in cutlass.range_constexpr(num_ptr_per_thread):
             row = i * num_threads + cRows[tidx % threads_per_row][0]
             idx = block * self.m_block_size + row

flashinfer/cute_dsl/sparse/blk128/utils.py

@@ -157,7 +157,7 @@ def make_tiled_copy_B(
 
 def mma_make_fragment_A(
     smem: cute.Tensor,
-    thr_mma: cute.core.ThrMma,
+    thr_mma: cute.ThrMma,
     swapAB: cutlass.Constexpr[bool] = False,
 ) -> cute.Tensor:
     if const_expr(swapAB):
@@ -168,7 +168,7 @@ def mma_make_fragment_A(
 
 def mma_make_fragment_B(
     smem: cute.Tensor,
-    thr_mma: cute.core.ThrMma,
+    thr_mma: cute.ThrMma,
     swapAB: cutlass.Constexpr[bool] = False,
 ) -> cute.Tensor:
     if const_expr(swapAB):
@@ -202,7 +202,7 @@ def warp_reduce(
     width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE,
 ) -> cute.TensorSSA | cute.Numeric:
     if const_expr(isinstance(val, cute.TensorSSA)):
-        res = cute.make_fragment(val.shape, val.dtype)
+        res = cute.make_rmem_tensor(val.shape, val.dtype)
         res.store(val)
         for i in cutlass.range_constexpr(cute.size(val.shape)):
             res[i] = warp_reduce(res[i], op, width)
@@ -257,7 +257,7 @@ def fmax_reduce(
     arch: cutlass.Constexpr[int] = 80,
 ) -> Float32:
     if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0):
-        res = cute.make_fragment(x.shape, Float32)
+        res = cute.make_rmem_tensor(x.shape, Float32)
         res.store(x)
         local_max = [res[0], res[1], res[2], res[3]]
         for i in cutlass.range_constexpr(4, cute.size(x.shape), 4):
@@ -274,7 +274,7 @@ def fmax_reduce(
             else fmax(local_max[0], init_val)
         )
     else:
-        res = cute.make_fragment(x.shape, Float32)
+        res = cute.make_rmem_tensor(x.shape, Float32)
         res.store(x)
         local_max_0 = (
             fmax(init_val, res[0], res[1])
@@ -307,7 +307,7 @@ def fadd_reduce(
             init_val = Float32.zero
         return x.reduce(cute.ReductionOp.ADD, init_val, 0)
     else:
-        res = cute.make_fragment(x.shape, Float32)
+        res = cute.make_rmem_tensor(x.shape, Float32)
         res.store(x)
         local_sum_0 = (
             cute.arch.add_packed_f32x2((init_val, 0.0), (res[0], res[1]))
@@ -356,7 +356,7 @@ def elem_pointer(
 @cute.jit
 def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
     # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if"
-    tApA = cute.make_fragment(
+    tApA = cute.make_rmem_tensor(
         cute.make_layout(
             (
                 cute.size(tAcA, mode=[0, 1]),
@@ -523,7 +523,7 @@ def cvt_f16(src: cute.Tensor, dst_or_dtype):
     if const_expr(isinstance(dst_or_dtype, type)):
         # dtype variant: create new tensor and call the tensor variant
         dtype = dst_or_dtype
-        dst = cute.make_fragment(src.shape, dtype)
+        dst = cute.make_rmem_tensor(src.shape, dtype)
         cvt_f16(src, dst)
         return dst
     else:
@@ -718,7 +718,7 @@ def domain_offset_aligned(
 @cute.jit
 def scalar_to_ssa(a: cute.Numeric, dtype) -> cute.TensorSSA:
     """Convert a scalar to a cute TensorSSA of shape (1,) and given dtype"""
-    vec = cute.make_fragment(1, dtype)
+    vec = cute.make_rmem_tensor(1, dtype)
     vec[0] = a
     return vec.load()
 

flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py

@@ -744,7 +744,7 @@ def _predicate_cpasync_rows(
         tCc: cute.Tensor,
         row_limit: Int32,
     ) -> cute.Tensor:
-        tPred = cute.make_fragment(
+        tPred = cute.make_rmem_tensor(
             cute.make_layout(
                 (
                     cute.size(tCc, mode=[0, 1]),
@@ -796,7 +796,7 @@ def _scale_copy_2d(
         tC: cute.Tensor,
         row_limit: Int32,
     ) -> None:
-        tP = cute.make_fragment(cute.make_layout(tS.shape), cutlass.Boolean)
+        tP = cute.make_rmem_tensor(cute.make_layout(tS.shape), cutlass.Boolean)
         for i in cutlass.range_constexpr(cute.size(tP)):
             tP[i] = cute.elem_less(tC[i][0][0][0], row_limit)
         for rest_m in cutlass.range_constexpr(cute.size(tS.shape[1])):