diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py index 50daf52cc..99a332578 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py @@ -484,24 +484,27 @@ def __call__( # All GEMMs: SS path tmma1 = _make_trivial_tiled_mma( self.q_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.K, + self.q_dtype, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.K, self.acc_dtype, cta_group, self.gemm1_tiler[:2], ) tmma2 = _make_trivial_tiled_mma( self.q_dtype, - tcgen05.OperandMajorMode.MN, - tcgen05.OperandMajorMode.MN, + self.q_dtype, + cute.nvgpu.OperandMajorMode.MN, + cute.nvgpu.OperandMajorMode.MN, self.acc_dtype, cta_group, self.gemm2_tiler[:2], ) tmma3 = _make_trivial_tiled_mma( self.q_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.MN, + self.q_dtype, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.MN, self.acc_dtype, cta_group, self.gemm3_tiler[:2], diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py index 1d575d954..8b657cafb 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py @@ -231,24 +231,27 @@ def __call__( # All GEMMs: SS path (A & B from SMEM, accumulator in TMEM) tmma1 = _make_trivial_tiled_mma( self.q_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.K, + self.q_dtype, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.K, self.acc_dtype, cta_group, self.gemm1_tiler[:2], ) tmma2 = _make_trivial_tiled_mma( self.q_dtype, - tcgen05.OperandMajorMode.MN, - tcgen05.OperandMajorMode.MN, + self.q_dtype, + cute.nvgpu.OperandMajorMode.MN, + cute.nvgpu.OperandMajorMode.MN, self.acc_dtype, cta_group, self.gemm2_tiler[:2], ) tmma3 = _make_trivial_tiled_mma( self.q_dtype, - tcgen05.OperandMajorMode.K, - tcgen05.OperandMajorMode.MN, + self.q_dtype, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.MN, self.acc_dtype, cta_group, self.gemm3_tiler[:2], @@ -1063,7 +1066,7 @@ def _compute_warpgroup( s_full_1_phase = Int32(0) dw_accum = cute.make_rmem_tensor(tSrS_shape, Float32) - for ei in cutlass.range_constexpr(cute.size(dw_accum)): + for ei in cutlass.range(cute.size(dw_accum), unroll_full=True): dw_accum[ei] = Float32(0.0) tSrS = cute.make_rmem_tensor(tSrS_shape, Float32) @@ -1118,16 +1121,16 @@ def _compute_warpgroup( # Phase 2: Convert dS f32→bf16, write to sdS via coordinate mapping. tSrS_f16 = cute.make_rmem_tensor(tSrS.shape, self.q_dtype) - for ei in cutlass.range_constexpr(cute.size(tSrS)): + for ei in cutlass.range(cute.size(tSrS), unroll_full=True): tSrS_f16[ei] = self.q_dtype(tSrS[ei]) if bi % 2 == 0: - for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): + for ei in cutlass.range(cute.size(tSrS_f16), unroll_full=True): h = cute.get(tCcS[ei], mode=[0, 0]) n = cute.get(tCcS[ei], mode=[0, 1]) sdS_gemm_view_0[h, n] = tSrS_f16[ei] else: - for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): + for ei in cutlass.range(cute.size(tSrS_f16), unroll_full=True): h = cute.get(tCcS[ei], mode=[0, 0]) n = cute.get(tCcS[ei], mode=[0, 1]) sdS_gemm_view_1[h, n] = tSrS_f16[ei] @@ -1148,7 +1151,7 @@ def _compute_warpgroup( cute.copy(tiled_tmem_load_dq, tDqDq_t2r, tDQrDQ) tDQrDQ_bf16 = cute.make_rmem_tensor(tDQrDQ.shape, self.q_dtype) - for ei in cutlass.range_constexpr(cute.size(tDQrDQ)): + for ei in cutlass.range(cute.size(tDQrDQ), unroll_full=True): tDQrDQ_bf16[ei] = self.q_dtype(tDQrDQ[ei] * Float32(sm_scale)) cute.arch.fence_view_async_tmem_load() @@ -1158,7 +1161,7 @@ def _compute_warpgroup( sdQ_epi_slice, cute.make_layout((self.heads_padded, self.head_dim_padded)), ) - for ei in cutlass.range_constexpr(cute.size(tDQrDQ_bf16)): + for ei in cutlass.range(cute.size(tDQrDQ_bf16), unroll_full=True): h = cute.get(tCcDQ[ei], mode=[0, 0]) d = cute.get(tCcDQ[ei], mode=[0, 1]) sdQ_gemm_view[h, d] = tDQrDQ_bf16[ei] @@ -1177,7 +1180,7 @@ def _compute_warpgroup( for h_local in cutlass.range_constexpr(HEADS_PER_WARP): h = warp_base_h + h_local my_partial = Float32(0.0) - for ei in cutlass.range_constexpr(cute.size(dw_accum)): + for ei in cutlass.range(cute.size(dw_accum), unroll_full=True): if cute.get(tCcS[ei], mode=[0, 0]) == h: my_partial = my_partial + dw_accum[ei] total = cute.arch.warp_reduction_sum(my_partial) @@ -1246,7 +1249,7 @@ def _reduce_warpgroup( # local→global when topk_indices_global=False); gmem fallback # mirrors that conversion via const_expr branch. batch_offset_l2g = Int32(0) if const_expr(self.topk_indices_global) else batch_idx * (seqlen_k // batch_size) - for pair in cutlass.range_constexpr(cute.size(tDKrDK) // 2): + for pair in cutlass.range(cute.size(tDKrDK) // 2, unroll_full=True): ei = pair * 2 n = cute.get(tCcDK[ei], mode=[0, 0]) d = cute.get(tCcDK[ei], mode=[0, 1]) @@ -1516,7 +1519,7 @@ class SharedStorage: if tidx == 0: block_sum = Float32(0.0) - for i in cutlass.range_constexpr(self.THREADS_PER_CTA): + for i in cutlass.range(self.THREADS_PER_CTA, unroll_full=True): block_sum += thread_sums[i] thread_sums[0] = block_sum cute.arch.sync_threads() diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.py index 2b3689432..8c5adef2d 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.py @@ -216,6 +216,7 @@ def __call__( self.q_major_mode = cutlass.utils.LayoutEnum.from_tensor(mQ).mma_major_mode() tiled_mma_qk = _make_trivial_tiled_mma( + self.q_dtype, self.q_dtype, self.k_major_mode, # A operand major mode (K) self.q_major_mode, # B operand major mode (Q) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py index 41ff88173..52ec85758 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py @@ -21,12 +21,10 @@ import cutlass.cute as cute import cutlass.utils as utils import torch -from cutlass.utils.distributed import atomicAdd - from cudnn.deepseek_sparse_attention.utils.compiler import compile_options from .block_scan import block_prefix_sum_kernel -from .indexer_top_k_varlen_util import IndexerTopKKernelVarlen +from .indexer_top_k_varlen_util import IndexerTopKKernelVarlen, atomicAdd class ComputeDynamicCTAOffsets: 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 57c7f9f37..dbbe6cce2 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 @@ -17,7 +17,7 @@ import cutlass.cute as cute import torch from cutlass._mlir.dialects import llvm -from cutlass.utils.distributed import atomicAdd +from cutlass.cutlass_dsl import dsl_user_op from .block_scan import block_prefix_sum_kernel, fence_acq_rel_cta @@ -26,6 +26,15 @@ """ +@dsl_user_op +def atomicAdd(dst_ptr: cute.Pointer, val: cutlass.Int32, *, loc=None, ip=None) -> cutlass.Int32: + """System-scope relaxed atomic add (drop-in for the deprecated + ``cutlass.utils.distributed.atomicAdd``).""" + return cute.arch.atomic_add( + dst_ptr.llvm_ptr, val, sem="relaxed", scope="sys", loc=loc, ip=ip + ) + + def half_as_ushort(half_val): """Interpret FP16 value as uint16 bit pattern""" return llvm.bitcast(cutlass.Uint16.mlir_type, half_val.ir_value()) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.py b/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.py index e841ed014..797ba9975 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.py @@ -237,6 +237,7 @@ def __call__( self.k_major_mode = cutlass.utils.LayoutEnum.from_tensor(mK).mma_major_mode() tiled_mma_qk = _make_trivial_tiled_mma( + self.q_dtype, self.q_dtype, self.k_major_mode, self.q_major_mode, diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py b/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py index fec8ec7b6..79f6b208c 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py @@ -249,7 +249,8 @@ def __call__( tiled_mma_qk = _make_trivial_tiled_mma( self.q_dtype, - tcgen05.OperandMajorMode.K, + self.q_dtype, + cute.nvgpu.OperandMajorMode.K, self.q_major_mode, self.qk_acc_dtype, cta_group, diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py index b04685b78..73e4be7ee 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py @@ -271,6 +271,7 @@ def _setup_attributes(self): ) tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -280,6 +281,7 @@ def _setup_attributes(self): self.mma_inst_shape_mn, ) tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -706,6 +708,7 @@ def __call__( # ---- TMA atoms ---- tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -715,6 +718,7 @@ def __call__( self.mma_inst_shape_mn, ) tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py index 30cad9534..f15ace3a6 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py @@ -236,6 +236,7 @@ def _setup_attributes(self): ) tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -245,6 +246,7 @@ def _setup_attributes(self): self.mma_inst_shape_mn, ) tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -680,6 +682,7 @@ def __call__( # ---- TMA atoms ---- tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -689,6 +692,7 @@ def __call__( self.mma_inst_shape_mn, ) tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py index 7e415accb..eecd0c805 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py @@ -256,6 +256,7 @@ def _setup_attributes(self) -> None: def _create_tiled_mma(self): return sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, @@ -267,6 +268,7 @@ def _create_tiled_mma(self): def _create_tiled_mma_sfb(self): return sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, self.a_dtype, self.a_major_mode, self.b_major_mode, diff --git a/python/cudnn/rmsnorm_rht_amax/kernel.py b/python/cudnn/rmsnorm_rht_amax/kernel.py index 2445cae94..af0993a65 100644 --- a/python/cudnn/rmsnorm_rht_amax/kernel.py +++ b/python/cudnn/rmsnorm_rht_amax/kernel.py @@ -191,7 +191,7 @@ def kernel(self, m_x: cute.Tensor, m_w: cute.Tensor, m_o: cute.Tensor, m_amax: c w = t_xr_w.load().to(Float32) y = x * rstd * w - for elem_idx in cutlass.range_constexpr(cfg.ept): + for elem_idx in cutlass.range(cfg.ept, unroll_full=True): reg[elem_idx] = y[elem_idx] for block_idx in cutlass.range_constexpr(cfg.num_vec_blocks): @@ -208,14 +208,14 @@ def kernel(self, m_x: cute.Tensor, m_w: cute.Tensor, m_o: cute.Tensor, m_amax: c for cross_stage in cutlass.range_constexpr(cfg.num_cross_stages): xor_mask = cutlass.Int32(1 << cross_stage) is_lower = (tid & xor_mask) == cutlass.Int32(0) - for elem_idx in cutlass.range_constexpr(cfg.ept): + for elem_idx in cutlass.range(cfg.ept, unroll_full=True): partner = shuffle_sync_bfly(reg[elem_idx], offset=xor_mask) if is_lower: reg[elem_idx] = reg[elem_idx] + partner else: reg[elem_idx] = partner - reg[elem_idx] - for elem_idx in cutlass.range_constexpr(cfg.ept): + for elem_idx in cutlass.range(cfg.ept, unroll_full=True): scaled = reg[elem_idx] * inv_sqrt_had abs_val = fabs_f32(scaled) running_max = fmax_f32(running_max, abs_val)