From 6e6be455c227e62de9887c143033a767ba47b550 Mon Sep 17 00:00:00 2001 From: Hyaloid Date: Thu, 25 Jun 2026 20:12:33 +0800 Subject: [PATCH] fix: IMA on indexer_topk_wrapper --- .../indexer_top_k_decode_varlen.py | 73 +++++++++++++++---- .../indexer_top_k_varlen_util.py | 30 ++++++-- 2 files changed, 84 insertions(+), 19 deletions(-) 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 d8b21341c..4ed33410a 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 @@ -14,6 +14,8 @@ """SM90+ CuTe DSL indexer top-K decode kernel.""" +import math + import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -687,18 +689,63 @@ def cute_dsl_topk_wrapper( buffer_numbers = 2 else: buffer_numbers = 1 - # Note: zeros will trigger an elementwise_add kernel. - buffer_torch = torch.empty(num_rows, buffer_numbers, num_cols, dtype=torch.int32, device="cuda") - g_global_counter_torch = None - # TVM FFI uses env stream automatically - compiled_kernel( - input_values, - None, # indices, used for merge blocks kernel of the multi-cta. - buffer_torch, - g_global_counter_torch, - seq_lens, - output_indices_torch, - output_values_torch, - ) + # Decode-varlen IMA workaround. + elems_per_row = buffer_numbers * num_cols + int32_max = (1 << 31) - 1 + total_elems = num_rows * elems_per_row + + if total_elems <= int32_max: + # int32 fast path: single launch, unchanged from before chunking. + buffer_torch = torch.empty( + num_rows, + buffer_numbers, + num_cols, + dtype=torch.int32, + device="cuda", + ) + # TVM FFI uses env stream automatically + compiled_kernel( + input_values, + None, # indices, used for merge blocks kernel of the multi-cta. + buffer_torch, + None, # g_global_counter_torch + seq_lens, + output_indices_torch, + output_values_torch, + ) + return output_indices_torch, output_values_torch + + # Fallback + if elems_per_row > 0: + max_chunk_rows = int32_max // elems_per_row + 1 + else: + max_chunk_rows = num_rows + input_elem_bytes = max(1, dtype.width // 8) + align_rows = max(1, 32 // input_elem_bytes) + row_step = (next_n * align_rows) // math.gcd(next_n, align_rows) + if max_chunk_rows < row_step: + chunk_rows = num_rows + else: + chunk_rows = (max_chunk_rows // row_step) * row_step + for row_lo in range(0, num_rows, chunk_rows): + row_hi = min(row_lo + chunk_rows, num_rows) + batch_lo = row_lo // next_n + batch_hi = row_hi // next_n + chunk_extra = torch.empty( + row_hi - row_lo, + buffer_numbers, + num_cols, + dtype=torch.int32, + device="cuda", + ) + compiled_kernel( + input_values[row_lo:row_hi], + None, + chunk_extra, + None, + seq_lens[batch_lo:batch_hi], + output_indices_torch[row_lo:row_hi], + output_values_torch[row_lo:row_hi] if return_val else None, + ) return output_indices_torch, output_values_torch 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..0ac48b950 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 @@ -407,12 +407,12 @@ def indexer_topk_kernel_per_row( # Note, for multi-cta version, each ctas must have its own extra_buffer. if cutlass.const_expr(self.enable_gmem_store): if cutlass.const_expr(self.enable_multi_cta): - grid_dim_x, grid_dim_y, _ = cute.arch.grid_dim() + _, grid_dim_y, _ = cute.arch.grid_dim() bidx_val, bidy_val, _ = cute.arch.block_idx() - buffer_row_id = bidx_val * grid_dim_y + bidy_val - buffer = extra_buffer[buffer_row_id, None, None] + buffer_row_id = cutlass.Int64(bidx_val) * cutlass.Int64(grid_dim_y) + cutlass.Int64(bidy_val) else: - buffer = extra_buffer[bidx, None, None] + buffer_row_id = cutlass.Int64(bidx) + buffer = self._slice_row_64bit(extra_buffer, buffer_row_id) # for initial scalar load part. row_ptr = score.iterator + row_start @@ -442,7 +442,7 @@ def indexer_topk_kernel_per_row( idX = cute.make_identity_tensor((shape[0], aligned_size)) input_ptr = input.iterator + vec_start input_addr_u64 = input_ptr.toint() - input_ptr_aligned = cute.make_ptr(self.dtype, input_addr_u64, assumed_align=align_bytes) + input_ptr_aligned = cute.make_ptr(self.dtype, input_addr_u64, input.memspace, assumed_align=align_bytes) input_tensor = cute.make_tensor( input_ptr_aligned, @@ -452,7 +452,7 @@ def indexer_topk_kernel_per_row( # slice for CTAs gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, None)) for mT in (input_tensor, idX)] # Note, we use gX_aligned here to avoid the alignment issue when the input is not aligned. - gX_aligned_ptr = cute.make_ptr(self.dtype, gX.iterator.toint(), assumed_align=align_bytes) + gX_aligned_ptr = cute.make_ptr(self.dtype, gX.iterator.toint(), gX.memspace, assumed_align=align_bytes) gX_aligned = cute.make_tensor(gX_aligned_ptr, cute.make_layout(gX.shape, stride=gX.stride)) self.num_sub_tiles = gX.shape[2] @@ -990,6 +990,24 @@ def indexer_topk_kernel_per_row( if cutlass.const_expr(self.return_val): cute.autovec_copy(topk_vals[None, i], mValues_store[None, col]) + def _slice_row_64bit(self, buffer, row_id): + row_offset_elems = row_id * cute.size(buffer.stride[0]) + elem_bytes = buffer.element_type.width // 8 + row_offset_bytes = row_offset_elems * elem_bytes + row_ptr = cute.make_ptr( + buffer.element_type, + buffer.iterator.toint() + row_offset_bytes, + buffer.memspace, + assumed_align=elem_bytes, + ) + return cute.make_tensor( + row_ptr, + cute.make_layout( + (buffer.shape[1], buffer.shape[2]), + stride=(buffer.stride[1], buffer.stride[2]), + ), + ) + def _get_tiled_copy(self): threads_per_row = self.num_threads_per_cta tiler_mn = (