Skip to content

fix: IMA on indexer_topk_wrapper #312

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
Anerudhan merged 1 commit into from
Jun 26, 2026
Merged

fix: IMA on indexer_topk_wrapper #312

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
73 changes: 60 additions & 13 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 @@ -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
Expand Down Expand Up @@ -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
30 changes: 24 additions & 6 deletions 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 @@ -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)
Comment on lines +410 to +415

@coderabbitai coderabbitai Bot Jun 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Description: Look for FE API tests covering the varlen TopK extra-buffer / IMA overflow path.
# Expectation: At least one test under test/python/fe_api references the decode-varlen TopK path
# and covers large row/column sizes or multi-CTA execution.

fd -p 'test/python/fe_api' -t f | xargs -r rg -n -C3 \
  'indexer_topk|topk|deepseek|decode_varlen|enable_multi_cta|extra_buffer|IMA|int64|Int64'

Repository: NVIDIA/cudnn-frontend

Length of output: 50377

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Locate any FE API tests or utilities that mention the DeepSeek sparse-attention indexer path
# or the specific knobs from the comment.
rg -n -C 2 \
  -g 'test/python/fe_api/**' \
  'deepseek_sparse_attention|indexer_top_k|indexer_topk|enable_gmem_store|enable_multi_cta|multi_cta|extra_buffer|buffer_row_id|_slice_row_64bit|Int64|int64' .

Repository: NVIDIA/cudnn-frontend

Length of output: 11006

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Check whether the FE API test tree contains any DeepSeek sparse-attention coverage
# and map the nearby test files if present.
fd -t f test/python/fe_api | rg 'deepseek|dsa|nsa|topk|indexer' || true

Repository: NVIDIA/cudnn-frontend

Length of output: 537

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Locate any FE API tests or utilities that mention the DeepSeek sparse-attention indexer path
# or the specific knobs from the comment.
rg -n -C 2 \
  -g 'test/python/fe_api/**' \
  'deepseek_sparse_attention|indexer_top_k|indexer_topk|enable_gmem_store|enable_multi_cta|multi_cta|extra_buffer|buffer_row_id|_slice_row_64bit|Int64|int64' .

Repository: NVIDIA/cudnn-frontend

Length of output: 11006

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Check whether the FE API test tree contains any DeepSeek sparse-attention coverage
# and map the nearby test files if present.
fd -t f test/python/fe_api | rg 'deepseek|dsa|nsa|topk|indexer' || true

Repository: NVIDIA/cudnn-frontend

Length of output: 537

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Look for the specific knobs in FE API tests and the DSA indexer_top_k tests.
rg -n -C 2 'enable_gmem_store|enable_multi_cta|extra_buffer|buffer_row_id|_slice_row_64bit' test/python/fe_api
printf '\n---\n'
rg -n -C 2 'enable_gmem_store|enable_multi_cta|extra_buffer|buffer_row_id|_slice_row_64bit' test/python/fe_api/dsa

Repository: NVIDIA/cudnn-frontend

Length of output: 159

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the DSA indexer_top_k test setup to see what it actually covers.
sed -n '1,220p' test/python/fe_api/dsa/test_DSA_indexer_top_k.py

Repository: NVIDIA/cudnn-frontend

Length of output: 3972

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Look for the specific knobs in FE API tests and the DSA indexer_top_k tests.
rg -n -C 2 'enable_gmem_store|enable_multi_cta|extra_buffer|buffer_row_id|_slice_row_64bit' test/python/fe_api
printf '\n---\n'
rg -n -C 2 'enable_gmem_store|enable_multi_cta|extra_buffer|buffer_row_id|_slice_row_64bit' test/python/fe_api/dsa

Repository: NVIDIA/cudnn-frontend

Length of output: 159

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the DSA indexer_top_k test setup to see what it actually covers.
sed -n '1,220p' test/python/fe_api/dsa/test_DSA_indexer_top_k.py

Repository: NVIDIA/cudnn-frontend

Length of output: 3972

Add FE API coverage for the 64-bit extra-buffer row path
test/python/fe_api/dsa/test_DSA_indexer_top_k.py covers the generic DSA.IndexerTopK API, but not the enable_gmem_store/enable_multi_cta branch or the buffer_row_id = bidx * grid_dim_y + bidy path here. Add a regression case that drives large extra_buffer row offsets so _slice_row_64bit is exercised.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py`
around lines 410 - 415, Add FE API test coverage for the 64-bit extra-buffer row
selection path in IndexerTopK. Extend
test/python/fe_api/dsa/test_DSA_indexer_top_k.py with a regression case that
enables the enable_gmem_store and enable_multi_cta branch, drives large
extra_buffer row offsets, and verifies the buffer_row_id computation using
bidx_val and grid_dim_y so _slice_row_64bit is exercised. Make sure the new test
goes through DSA.IndexerTopK and specifically hits the multi-CTA row indexing
path rather than only the generic API path.

Source: Path instructions


# for initial scalar load part.
row_ptr = score.iterator + row_start
Expand Down Expand Up @@ -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,
Expand All @@ -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]
Expand Down Expand Up @@ -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 = (
Expand Down