diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py index 41fbb5106..e4ffad3c7 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py @@ -32,7 +32,7 @@ def _mma_partition_fragment_AB( - thr_mma: cute.core.ThrMma, + thr_mma: cute.ThrMma, sA: cute.Tensor, sB: cute.Tensor, swap_AB: bool, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py index 773b5ea33..ea1997711 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py @@ -78,8 +78,8 @@ def kernel(self, mIn: cute.Tensor, mOut: cute.Tensor, mLen: cute.Tensor, rows: I row = bidx * ROWS_PER_CTA + warp_id if row < rows: - local_vals = cute.make_fragment((chunk,), Int32) - local_is_valid = cute.make_fragment((chunk,), cutlass.Boolean) + local_vals = cute.make_rmem_tensor((chunk,), Int32) + local_is_valid = cute.make_rmem_tensor((chunk,), cutlass.Boolean) cnt_v = Int32(0) cnt_i = Int32(0) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py index 5187f2645..aef1bf8a5 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py @@ -962,8 +962,8 @@ def indexer_topk_kernel_per_row( assert self.top_k % vecsize_out == 0 nvec_per_thread = cutlass.const_expr(cute.ceil_div(self.top_k, vecsize_out * self.num_threads_per_cta)) - topk_vals = cute.make_fragment((vecsize_out, nvec_per_thread), self.dtype) - topk_indices = cute.make_fragment((vecsize_out, nvec_per_thread), cutlass.Int32) + topk_vals = cute.make_rmem_tensor((vecsize_out, nvec_per_thread), self.dtype) + topk_indices = cute.make_rmem_tensor((vecsize_out, nvec_per_thread), cutlass.Int32) stride = self.num_threads_per_cta * vecsize_out for i in cutlass.range(nvec_per_thread, unroll_full=True): @@ -1018,7 +1018,7 @@ def _get_tiled_copy(self): @cute.jit def predicate_tile(self, tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor: - tApA = cute.make_fragment( + tApA = cute.make_rmem_tensor( cute.make_layout( ( cute.size(tAcA, mode=[0, 1]), diff --git a/python/cudnn/deepseek_sparse_attention/utils/sm90/mma.py b/python/cudnn/deepseek_sparse_attention/utils/sm90/mma.py index b187bc5b3..022508afb 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/sm90/mma.py +++ b/python/cudnn/deepseek_sparse_attention/utils/sm90/mma.py @@ -91,7 +91,7 @@ def gemm_w_idx( def mma_partition_fragment_AB( - thr_mma: cute.core.ThrMma, + thr_mma: cute.ThrMma, sA: Optional[cute.Tensor], sB: Optional[cute.Tensor], swap_AB: bool, diff --git a/python/cudnn/native_sparse_attention/compression/fmha.py b/python/cudnn/native_sparse_attention/compression/fmha.py index cf874aa54..a28eede15 100644 --- a/python/cudnn/native_sparse_attention/compression/fmha.py +++ b/python/cudnn/native_sparse_attention/compression/fmha.py @@ -1756,7 +1756,7 @@ def softmax( cum_seqlen_q: Optional[cute.Tensor], cum_seqlen_k: Optional[cute.Tensor], scale_softmax_log2: Float32, - qk_thr_mma: cute.core.ThrMma, + qk_thr_mma: cute.ThrMma, tStS: cute.Tensor, tStSi: cute.Tensor, window_size_left: Optional[Int32], @@ -1791,7 +1791,7 @@ def softmax( :param scale_softmax_log2: Log2 scale factor for softmax operation :type scale_softmax_log2: Float32 :param qk_thr_mma: Thread MMA operation for QK matrix multiplication - :type qk_thr_mma: cute.core.ThrMma + :type qk_thr_mma: cute.ThrMma :param tStS: Shared tensor for softmax input/output :type tStS: cute.Tensor :param tStSi: Input tensor containing attention scores @@ -2056,7 +2056,7 @@ def softmax( @cute.jit def correction_rescale( self, - thr_mma: cute.core.ThrMma, + thr_mma: cute.ThrMma, tOtO: cute.Tensor, scale: Float32, ): @@ -2073,7 +2073,7 @@ def correction_rescale( 3. Store the rescaled results back to tensor memory :param thr_mma: Thread MMA operation for the computation - :type thr_mma: cute.core.ThrMma + :type thr_mma: cute.ThrMma :param tOtO: Tensor representing partial attention output to be rescaled :type tOtO: cute.Tensor :param scale: Scaling factor to apply to the partial results @@ -2134,7 +2134,7 @@ def correction_rescale( @cute.jit def correction_epilog( self, - thr_mma: cute.core.ThrMma, + thr_mma: cute.ThrMma, tOtO: cute.Tensor, mLSE: cute.Tensor | None, tTMEM_LOAD_VECrS: cute.Tensor, @@ -2160,7 +2160,7 @@ def correction_epilog( 5. Preparation for efficient TMA store operations :param thr_mma: Thread MMA operation for the computation - :type thr_mma: cute.core.ThrMma + :type thr_mma: cute.ThrMma :param tOtO: Tensor containing accumulated attention output :type tOtO: cute.Tensor :param mLSE: Tensor containing log-sum-exp values for LSE calculation diff --git a/python/cudnn/sdpa/utils.py b/python/cudnn/sdpa/utils.py index a1210a5ea..d4a1fedc5 100644 --- a/python/cudnn/sdpa/utils.py +++ b/python/cudnn/sdpa/utils.py @@ -165,7 +165,7 @@ def make_tiled_copy_B( # noqa: N802 def mma_make_fragment_A( # noqa: N802 smem: cute.Tensor, - thr_mma: cute.core.ThrMma, + thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False, # noqa: N803 ) -> cute.Tensor: """Wrapper for cute.mma_make_fragment.""" @@ -176,7 +176,7 @@ def mma_make_fragment_A( # noqa: N802 def mma_make_fragment_B( # noqa: N802 smem: cute.Tensor, - thr_mma: cute.core.ThrMma, + thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False, # noqa: N803 ) -> cute.Tensor: """Wrapper for cute.mma_make_fragment.""" @@ -207,7 +207,7 @@ def warp_reduce( ) -> cute.TensorSSA | cute.Numeric: """Reduction in warp.""" if cutlass.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) @@ -350,7 +350,7 @@ def exp2f(x: cute.TensorSSA | Float32) -> cute.TensorSSA | Float32: :rtype: cute.TensorSSA or Float32. """ if cutlass.const_expr(isinstance(x, cute.TensorSSA)): - res = cute.make_fragment(x.shape, Float32) + res = cute.make_rmem_tensor(x.shape, Float32) res.store(x) for i in cutlass.range_constexpr(cute.size(x.shape)): res[i] = cute.math.exp2(res[i], fastmath=True) @@ -399,7 +399,7 @@ def fmax(a: float | Float32, b: float | Float32, c: float | Float32 | None = Non def fmax_reduce(x: cute.TensorSSA, init_val: float | Float32 | None = None, arch: cutlass.Constexpr[int] = 80) -> Float32: """Fmax reduce wrapper.""" if cutlass.const_expr(arch < ARCH_SM100 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): @@ -413,7 +413,7 @@ def fmax_reduce(x: cute.TensorSSA, init_val: float | Float32 | None = None, arch return local_max[0] if cutlass.const_expr(init_val is None) else fmax(local_max[0], init_val) # [2025-06-15] x.reduce only seems to use 50% 3-input max and 50% 2-input max # We instead force the 3-input max. - res = cute.make_fragment(x.shape, Float32) + res = cute.make_rmem_tensor(x.shape, Float32) res.store(x) local_max = [ fmax(init_val, res[0], res[1]) if cutlass.const_expr(init_val is not None) else fmax(res[0], res[1]), @@ -437,7 +437,7 @@ def fadd_reduce(x: cute.TensorSSA, init_val: float | Float32 | None = None, arch if cutlass.const_expr(init_val is None): init_val = Float32.zero return x.reduce(cute.ReductionOp.ADD, init_val, 0) - 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])) if cutlass.const_expr(init_val is not None) else (res[0], res[1]) local_sum = [local_sum_0, (res[2], res[3]), (res[4], res[5]), (res[6], res[7])] @@ -503,7 +503,7 @@ def predicate_k( ) -> cute.Tensor: """Wrapper for predicate.""" # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if" - tApA = cute.make_fragment( # noqa: N806 + tApA = cute.make_rmem_tensor( # noqa: N806 cute.make_layout( (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])), stride=(cute.size(tAcA, mode=[2]), 0, 1), @@ -547,7 +547,7 @@ def shuffle_sync( mask = cute.arch.WARP_SIZE - width clamp = cute.arch.WARP_SIZE - 1 mask_and_clamp = mask << 8 | clamp - val = cute.make_fragment(1, type(value)) + val = cute.make_rmem_tensor(1, type(value)) val[0] = value val_i32 = cute.recast_tensor(val, cutlass.Int32) for i in cutlass.range_constexpr(cute.size(val_i32)):