Skip to content

Fix cutlass DSL deprecation warnings in CuTe DSL kernels #376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
saltyminty merged 1 commit into from
Jul 10, 2026
Merged

Fix cutlass DSL deprecation warnings in CuTe DSL kernels #376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down
33 changes: 18 additions & 15 deletions python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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]
Expand All @@ -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()
Expand All @@ -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]
Expand All @@ -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)
Expand Down Expand Up @@ -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])
Expand Down Expand Up @@ -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()
Expand Down
1 change: 1 addition & 0 deletions python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 1 addition & 3 deletions python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
11 changes: 10 additions & 1 deletion python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
)
Comment thread
saltyminty marked this conversation as resolved.


def half_as_ushort(half_val):
"""Interpret FP16 value as uint16 bit pattern"""
return llvm.bitcast(cutlass.Uint16.mlir_type, half_val.ir_value())
Expand Down
1 change: 1 addition & 0 deletions python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
6 changes: 3 additions & 3 deletions python/cudnn/rmsnorm_rht_amax/kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
Expand Down