Skip to content

DSA indexer_top_k: exclude OOB tile-fill lanes from radix selection #445

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

Closed
JackRao123 wants to merge 1 commit into from
Closed

DSA indexer_top_k: exclude OOB tile-fill lanes from radix selection #445

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
91 changes: 51 additions & 40 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 @@ -525,12 +525,18 @@ def indexer_topk_kernel_per_row(
-tXrX.element_type.inf,
)

# Skip out-of-bounds lanes: the last tile can extend past
# aligned_size, and its -inf-filled lanes must not be counted
# as real elements.
cur_tXcX = tXcX[None, None, None, tile_idx]
for i in cutlass.range(cute.size(tXrX), unroll_full=True):
bin_val = self.to_coarse_key(tXrX[i])
atomicAdd(
s_histogram.iterator + cutlass.Int32(bin_val),
val_one,
)
col = cur_tXcX[i // vec_size][1] + i % vec_size
if col < aligned_size:
bin_val = self.to_coarse_key(tXrX[i])
atomicAdd(
s_histogram.iterator + cutlass.Int32(bin_val),
val_one,
)

# for initial scalar load part.
for j in range(tidx, prologue_elems, self.num_threads_per_cta):
Expand Down Expand Up @@ -594,11 +600,13 @@ def indexer_topk_kernel_per_row(
)
for i in cutlass.range(cute.size(tXrX), unroll_full=True):
cur_tXcX = tXcX[None, None, None, tile_idx]
bin_val = self.to_coarse_key(tXrX[i])
if bin_val < threshold_bin:
pos = atomicAdd(s_counter.iterator, val_one)
idx = self.index_type(cur_tXcX[i // vec_size][1] + i % vec_size + vec_start)
s_indices[pos] = idx
col = cur_tXcX[i // vec_size][1] + i % vec_size
if col < aligned_size:
bin_val = self.to_coarse_key(tXrX[i])
if bin_val < threshold_bin:
pos = atomicAdd(s_counter.iterator, val_one)
idx = self.index_type(col + vec_start)
s_indices[pos] = idx

# for initial scalar load part.
for j in range(tidx, prologue_elems, self.num_threads_per_cta):
Expand Down Expand Up @@ -651,39 +659,42 @@ def indexer_topk_kernel_per_row(
raw_input = tXrX[i]
bin_val = self.to_coarse_key(raw_input)
cur_tXcX = tXcX[None, None, None, tile_idx]
idx = self.index_type(cur_tXcX[i // vec_size][1] + i % vec_size + vec_start)
if bin_val < threshold_bin:
pos = atomicAdd(s_counter.iterator, val_one)
s_indices[pos] = idx
elif bin_val == threshold_bin:
# pos = atomicAdd(s_num_input[0], 1)
pos = atomicAdd(s_num_input.iterator, val_one)
if cutlass.const_expr(self.enable_gmem_store):
if pos < self.indexer_topk_smem_input_size:
s_input_idx[0, pos] = idx
col = cur_tXcX[i // vec_size][1] + i % vec_size
# Skip out-of-bounds lanes filled with -inf by _fill_oob.
if col < aligned_size:
idx = self.index_type(col + vec_start)
if bin_val < threshold_bin:
pos = atomicAdd(s_counter.iterator, val_one)
s_indices[pos] = idx
elif bin_val == threshold_bin:
# pos = atomicAdd(s_num_input[0], 1)
pos = atomicAdd(s_num_input.iterator, val_one)
if cutlass.const_expr(self.enable_gmem_store):
if pos < self.indexer_topk_smem_input_size:
s_input_idx[0, pos] = idx
else:
buffer_pos = atomicAdd(
g_num_input.iterator,
val_one,
)
buffer[0, buffer_pos] = cutlass.Int32(cutlass.Uint32(idx))
ordered = self.to_ordered(raw_input)
sub_bin = (ordered >> self.first_refine_shift) & 0xFF
# atomicAdd(s_histogram[sub_bin], 1)
atomicAdd(
s_histogram.iterator + cutlass.Int32(sub_bin),
val_one,
)
else:
buffer_pos = atomicAdd(
g_num_input.iterator,
if pos < self.indexer_topk_smem_input_size:
s_input_idx[0, pos] = idx
ordered = self.to_ordered(raw_input)
sub_bin = (ordered >> self.first_refine_shift) & 0xFF
# atomicAdd(s_histogram[sub_bin], 1)
atomicAdd(
s_histogram.iterator + cutlass.Int32(sub_bin),
val_one,
)
buffer[0, buffer_pos] = cutlass.Int32(cutlass.Uint32(idx))
ordered = self.to_ordered(raw_input)
sub_bin = (ordered >> self.first_refine_shift) & 0xFF
# atomicAdd(s_histogram[sub_bin], 1)
atomicAdd(
s_histogram.iterator + cutlass.Int32(sub_bin),
val_one,
)
else:
if pos < self.indexer_topk_smem_input_size:
s_input_idx[0, pos] = idx
ordered = self.to_ordered(raw_input)
sub_bin = (ordered >> self.first_refine_shift) & 0xFF
# atomicAdd(s_histogram[sub_bin], 1)
atomicAdd(
s_histogram.iterator + cutlass.Int32(sub_bin),
val_one,
)

# for initial scalar load part.
for j in range(tidx, prologue_elems, self.num_threads_per_cta):
Expand Down
69 changes: 69 additions & 0 deletions test/python/fe_api/dsa/test_DSA_indexer_top_k.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,3 +143,72 @@ def test_DSA_indexer_top_k_wrapper(
values,
return_val,
)


@pytest.mark.L0
@pytest.mark.parametrize("trigger_dist", ["identical", "distinct"])
def test_DSA_indexer_top_k_oob_tile_lanes(trigger_dist):
"""Regression test: -inf-filled OOB lanes of a row's final vector tile must
not participate in the radix top-k.

The kernel reads each row in fixed-width vector tiles and fills the final
partial tile's out-of-bounds lanes with -inf. Those phantom lanes were
counted as real elements by the histogram and candidate-collection passes.
When a row's top-k threshold lands in the fp16 -inf coarse bin (fewer than
top_k values above ~-65504, since to_coarse_key converts fp32 scores to
fp16 first), the phantom lanes flood the threshold bin's candidate list and
overflow the per-row candidate buffers in the large-occupancy compile
(>148 rows), producing out-of-bounds shared/global writes
(cudaErrorIllegalAddress) or silently out-of-range selected indices.

This test's batch deterministically places row 0 in that regime: 149 rows
(>148 -> large_occupancy), a 4310-wide matrix, and a 4122-long row whose
values below the fp16 negative-overflow point plus the phantom lanes
outnumber the candidate-buffer capacity.
"""
try:
from cudnn import DSA
from cuda.bindings import driver as cuda
except ImportError:
pytest.skip("Environment not supported: cudnn[cutedsl] not installed")
if torch.cuda.get_device_capability()[0] < 9:
pytest.skip("indexer top-k kernel requires SM90+")

num_rows, num_cols, top_k, trigger_len, num_above = 149, 4310, 2048, 4122, 1122
scores = torch.empty(num_rows, num_cols, dtype=torch.float32)
scores[1:] = torch.arange(num_cols, dtype=torch.float32).expand(num_rows - 1, -1)
row0 = torch.empty(num_cols, dtype=torch.float32)
row0[:num_above] = torch.arange(num_above, dtype=torch.float32)
if trigger_dist == "identical":
row0[num_above:trigger_len] = -100000.0
else:
row0[num_above:trigger_len] = torch.linspace(-66000.0, -200000.0, trigger_len - num_above)
row0[trigger_len:] = float("-inf")
scores[0] = row0
seq_lens = torch.full((num_rows,), num_cols, dtype=torch.int32)
seq_lens[0] = trigger_len

input_values = scores.cuda()
seq_lens = seq_lens.cuda()
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
result = DSA.indexer_top_k_wrapper(
input_values,
seq_lens,
top_k,
next_n=1,
return_val=True,
stream=stream,
)
torch.cuda.synchronize()

indices = result["indices"].cpu()
values = result["values"].cpu()
for r in range(num_rows):
L = int(seq_lens[r].cpu())
k = min(top_k, L)
valid = indices[r] >= 0
assert int(valid.sum()) == k, f"row {r}: expected {k} valid indices, got {int(valid.sum())}"
assert int(indices[r].max()) < L, f"row {r}: out-of-range index {int(indices[r].max())} >= seq_len {L}"
got = values[r][valid].sort(descending=True).values
ref = torch.topk(scores[r, :L], k).values.sort(descending=True).values
torch.testing.assert_close(got, ref, atol=1e-6, rtol=1e-5)
Comment on lines +209 to +214

@coderabbitai coderabbitai Bot Jul 28, 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 | ⚡ Quick win

Assert selected indices are unique.

For trigger_dist="identical", duplicate in-range indices can still produce the same sorted values as torch.topk, masking a broken selection result.

Suggested assertion 
     valid = indices[r] >= 0
     assert int(valid.sum()) == k, f"row {r}: expected {k} valid indices, got {int(valid.sum())}"
-    assert int(indices[r].max()) < L, f"row {r}: out-of-range index {int(indices[r].max())} >= seq_len {L}"
+    selected_indices = indices[r][valid]
+    assert int(selected_indices.max()) < L, f"row {r}: out-of-range index {int(selected_indices.max())} >= seq_len {L}"
+    assert selected_indices.unique().numel() == k, f"row {r}: duplicate selected indices"
     got = values[r][valid].sort(descending=True).values
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
valid = indices[r] >= 0
assert int(valid.sum()) == k, f"row {r}: expected {k} valid indices, got {int(valid.sum())}"
assert int(indices[r].max()) < L, f"row {r}: out-of-range index {int(indices[r].max())} >= seq_len {L}"
got = values[r][valid].sort(descending=True).values
ref = torch.topk(scores[r, :L], k).values.sort(descending=True).values
torch.testing.assert_close(got, ref, atol=1e-6, rtol=1e-5)
valid = indices[r] >= 0
assert int(valid.sum()) == k, f"row {r}: expected {k} valid indices, got {int(valid.sum())}"
selected_indices = indices[r][valid]
assert int(selected_indices.max()) < L, f"row {r}: out-of-range index {int(selected_indices.max())} >= seq_len {L}"
assert selected_indices.unique().numel() == k, f"row {r}: duplicate selected indices"
got = values[r][valid].sort(descending=True).values
ref = torch.topk(scores[r, :L], k).values.sort(descending=True).values
torch.testing.assert_close(got, ref, atol=1e-6, rtol=1e-5)
🤖 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 `@test/python/fe_api/dsa/test_DSA_indexer_top_k.py` around lines 209 - 214, In
the row-validation assertions around the DSA top-k result, add a check that the
valid entries in indices[r] are unique before comparing values with torch.topk.
Use the existing valid mask and preserve the current count, range, and value
assertions.