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

[JackRao123]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran black --line-length 160 on the changed files.

Fixes #446

Affected area

FE OSS kernels or CuTeDSL (DSA indexer_top_k).

Summary

The CuTe-DSL radix top-k kernel reads each row in fixed-width vector tiles. The final partial tile's out-of-bounds lanes are filled with -inf by _fill_oob so the predicated copy is memory-safe — but the histogram and candidate-collection loops then iterate every fragment lane without a bounds check, counting those phantom -inf lanes as real elements.

This is harmless while a row's top-k threshold sits above the phantom bin, because phantom -inf lanes sort last and are dropped. However, to_coarse_key computes the coarse radix bin from the fp16 conversion of the fp32 score, so every score below fp16's −65504 minimum collapses into the fp16 -inf bin — the same bin the phantom lanes occupy. When a row's top-k threshold lands inside that bin (fewer than top_k values above ~−65504), the threshold bin's candidate list becomes real_count + up-to-tile-width phantom lanes:

• in the large_occupancy compile (num_rows > 148), the per-row candidate buffers (512-entry smem + num_cols gmem) overflow → out-of-bounds shared/global writes → cudaErrorIllegalAddress;
• in any configuration, phantom lanes can be selected as winners, yielding silently out-of-range indices.

Observed in production

A GLM-5.2 context-parallel (CP32) training run on B200s crashed with Xid 43 (illegal memory access) mid-forward_backward — once in 5 hours, sequence-length-independent, no ECC errors. The crash is deterministically data-dependent: the captured score tensor reproduces it in a fresh process, on both B200 and B300, at the same call for the same data. CUDA_LAUNCH_BLOCKING=1 attributes the fault to this kernel; compute-sanitizer reports invalid 2-byte shared writes; bisection isolated the trigger to a single row (seq_len 4122 of width 4310, 2242 values below −65504) plus the >148-row compile.

Fix

Skip lanes whose tile column coordinate is out of bounds (col < aligned_size) in the three vectorized histogram/collection loops — the same bound the predicated copy already computes via predicate_tile. The scalar prologue/leftover loops are exact-bounded and unchanged. Selection semantics for in-bounds elements are untouched.

Testing

• New L0 regression test test_DSA_indexer_top_k_oob_tile_lanes deterministically drives a row into the flood regime (149 rows → large_occupancy; 4310-wide; a 4122-long row with < top_k values above the fp16 negative-overflow point), in two variants: identical collapsed values (fails pre-patch with out-of-range indices) and distinct collapsed values (pre-patch: cudaErrorIllegalAddress). Both pass post-patch.
• GPU validation was performed against the v1.26.0 release (the affected version; the patched region is byte-identical on develop) on B200 and B300: previously-crashing inputs complete with in-range indices, compute-sanitizer reports no invalid accesses, and 12/12 randomized parity trials against torch.topk pass (rows 37–862, widths 4097–51720, fp16-overflow and exact-tie distributions, both kernel compiles).
• Kernel microbench, (862, 51720) k=2048: 0.115 → 0.129 ms/call (+12% on a 0.1 ms kernel).

Related

• Complements DSA indexer_top_k: fall back to scalar stores for odd top_k #407 (odd top_k store width) — same kernel, disjoint code (that fix is the phase-3 output store; this one is phases 1–2 selection).
• fix: IMA on indexer_topk_wrapper #312 fixed a different IMA (int32 scratch overflow); this one is the data-dependent candidate flood.

[coderabbitai]

📝 Walkthrough

Walkthrough

The varlen indexer top-k kernel now excludes out-of-bounds lanes from histogram updates, threshold processing, and index writes. A parametrized SM90+ regression test verifies valid indices and agreement with torch.topk.

Changes

OOB tile lane handling

Layer / File(s)	Summary
Guard padded lanes in top-k processing python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py	Histogram updates, threshold comparisons, index computation, and buffering now run only for lanes where col < aligned_size.
Validate OOB tile lane behavior test/python/fe_api/dsa/test_DSA_indexer_top_k.py	Adds parametrized SM90+ coverage for near-threshold OOB-lane scenarios, checking index validity, bounds, and torch.topk value agreement.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested labels: cat-bug

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly matches the core change: excluding out-of-bounds tile-fill lanes from radix top-k selection.
Description check	✅ Passed	The description is mostly complete and aligned with the template, covering area, summary, related issue, and testing.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with 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.

Inline comments:
In `@test/python/fe_api/dsa/test_DSA_indexer_top_k.py`:
- Around line 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.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 96e8c3db-c70d-4bec-9fd4-c95d2024b54e

📥 Commits

Reviewing files that changed from the base of the PR and between dd5ce86 and 5cb18d3.

📒 Files selected for processing (2)

• python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py
• test/python/fe_api/dsa/test_DSA_indexer_top_k.py

[coderabbitai]

🎯 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.

[JackRao123]

Closing as a duplicate of #410, which predates this PR and carries an essentially identical fix (same OOB-lane guards in the same three loops). Production evidence and an additional trigger route (fp16-collapse of scores below −65504 into the -inf coarse bin) are noted in the #410 comments. Whichever lands first works for us.