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 b7a243a2a..2e31ac413 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 @@ -1074,7 +1074,7 @@ def _compute_warpgroup( s_full_1_phase = Int32(0) dw_accum = cute.make_rmem_tensor(tSrS_shape, Float32) - for ei in cutlass.range(cute.size(dw_accum), unroll_full=True): + for ei in cutlass.range_constexpr(cute.size(dw_accum)): dw_accum[ei] = Float32(0.0) tSrS = cute.make_rmem_tensor(tSrS_shape, Float32) @@ -1129,16 +1129,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(cute.size(tSrS), unroll_full=True): + for ei in cutlass.range_constexpr(cute.size(tSrS)): tSrS_f16[ei] = self.q_dtype(tSrS[ei]) if bi % 2 == 0: - for ei in cutlass.range(cute.size(tSrS_f16), unroll_full=True): + for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): 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(cute.size(tSrS_f16), unroll_full=True): + for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): 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] @@ -1159,7 +1159,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(cute.size(tDQrDQ), unroll_full=True): + for ei in cutlass.range_constexpr(cute.size(tDQrDQ)): tDQrDQ_bf16[ei] = self.q_dtype(tDQrDQ[ei] * Float32(sm_scale)) cute.arch.fence_view_async_tmem_load() @@ -1169,7 +1169,7 @@ def _compute_warpgroup( sdQ_epi_slice, cute.make_layout((self.heads_padded, self.head_dim_padded)), ) - for ei in cutlass.range(cute.size(tDQrDQ_bf16), unroll_full=True): + for ei in cutlass.range_constexpr(cute.size(tDQrDQ_bf16)): h = cute.get(tCcDQ[ei], mode=[0, 0]) d = cute.get(tCcDQ[ei], mode=[0, 1]) sdQ_gemm_view[h, d] = tDQrDQ_bf16[ei] @@ -1188,7 +1188,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(cute.size(dw_accum), unroll_full=True): + for ei in cutlass.range_constexpr(cute.size(dw_accum)): if cute.get(tCcS[ei], mode=[0, 0]) == h: my_partial = my_partial + dw_accum[ei] total = cute.arch.warp_reduction_sum(my_partial) @@ -1257,7 +1257,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(cute.size(tDKrDK) // 2, unroll_full=True): + for pair in cutlass.range_constexpr(cute.size(tDKrDK) // 2): ei = pair * 2 n = cute.get(tCcDK[ei], mode=[0, 0]) d = cute.get(tCcDK[ei], mode=[0, 1])