From 871a177bc4e2b8887b92e8b288c3bcc8f1692f3a Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 10 Jul 2026 12:53:12 -0700 Subject: [PATCH 1/2] fix(cute_dsl): migrate off APIs removed in nvidia-cutlass-dsl 4.6 cutlass-dsl 4.6.0 removed the long-deprecated cute.make_fragment (use cute.make_rmem_tensor) and cute.core.ThrMma (use cute.ThrMma). Since requirements.txt floats on >=4.5.0, any fresh install now resolves to 4.6.0 and the blk128 sparse-attention and SM120 b12x GEMM paths raise AttributeError at trace time. Both replacement names already exist in 4.5.0/4.5.1/4.5.2 (verified against the v4.5.x tags and a live 4.5.2 install), so this is a pure rename with no version guard needed and no requirements change: the same code now works on 4.5.x and 4.6.x. - cute.make_fragment( -> cute.make_rmem_tensor( (18 sites, 4 files) - cute.core.ThrMma -> cute.ThrMma (13 sites, 2 files) Verified: a make_rmem_tensor smoke kernel traces, compiles and runs correctly on SM100 under both cutlass-dsl 4.5.2 and 4.6.0; repo-wide grep for the removed spellings is clean (flashinfer/, tests/, benchmarks/). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 --- .../cute_dsl/sparse/blk128/flash_fwd_sm100.py | 38 +++++++++---------- flashinfer/cute_dsl/sparse/blk128/pack_gqa.py | 2 +- flashinfer/cute_dsl/sparse/blk128/utils.py | 18 ++++----- .../dense_blockscaled_gemm_sm120_b12x.py | 4 +- 4 files changed, 31 insertions(+), 31 deletions(-) diff --git a/flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py b/flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py index efd2f9c8c51..1bd8d122fe4 100644 --- a/flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py +++ b/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, @@ -2117,7 +2117,7 @@ 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) + cute.make_rmem_tensor(tSrScale_t2r_shape, Float32) # q_stage=1 correction loop if const_expr(mBlockNums is not None): block_iter_count = ( @@ -2330,7 +2330,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 +2370,9 @@ 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) + tOrO_frg = cute.make_rmem_tensor((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 +2390,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 +2447,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 +2476,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 +2543,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: diff --git a/flashinfer/cute_dsl/sparse/blk128/pack_gqa.py b/flashinfer/cute_dsl/sparse/blk128/pack_gqa.py index db2818bc865..5be5a8d1aaa 100644 --- a/flashinfer/cute_dsl/sparse/blk128/pack_gqa.py +++ b/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 diff --git a/flashinfer/cute_dsl/sparse/blk128/utils.py b/flashinfer/cute_dsl/sparse/blk128/utils.py index be5123f87ae..1a0dcdec5ca 100644 --- a/flashinfer/cute_dsl/sparse/blk128/utils.py +++ b/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() diff --git a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py b/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py index 24c88fef195..3763776d69c 100644 --- a/flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py +++ b/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])): From cb7ec1b52edf93277388d9bc381572d1bfdda888 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Tue, 14 Jul 2026 14:34:26 -0700 Subject: [PATCH 2/2] chore(cute_dsl): drop dead scale-fragment scaffolding in blk128 correction warp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review flagged an unused cute.make_rmem_tensor call and a shadowed outer tOrO_frg allocation. Both predate the 4.6 rename (they were equally dead as cute.make_fragment). The unused allocation was the last reference into a vestigial tmem scale-load scaffold (tStScale_layout / tStScales / tScScale / thr_tmem_load_vec chain), so the whole reference-less chain is removed rather than just the flagged line; every deleted symbol has zero remaining uses in the enclosing scope (ruff F841-clean). Trace-time-pure deletions only — no copies, barriers, or arch ops removed. Re-verified on SM100: tests/attention/test_vsa_block_sparse.py is 30 passed / 9 skipped under both cutlass-dsl 4.5.2 and 4.6.1, matching the pre-cleanup runs. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 --- .../cute_dsl/sparse/blk128/flash_fwd_sm100.py | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py b/flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py index 1bd8d122fe4..5e4497536f6 100644 --- a/flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py +++ b/flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py @@ -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_rmem_tensor(tSrScale_t2r_shape, Float32) # q_stage=1 correction loop if const_expr(mBlockNums is not None): block_iter_count = ( @@ -2370,7 +2345,6 @@ 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_rmem_tensor((tOrO_t2r_shape, frg_count), self.pv_acc_dtype) for i in cutlass.range_constexpr(frg_count): tOrO_frg = cute.make_rmem_tensor(tOrO_t2r_shape, self.pv_acc_dtype) tOtO_t2r_i = cute.make_tensor(