Fix IMA caused by OOB lanes in varlen indexer top-k

[hxbai]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.

Affected area

• FE OSS kernels or CuTeDSL

Summary

Fix the varlen indexer top-k kernel so that out-of-bounds lanes from predicated vector loads do not participate in radix histogram construction or candidate collection.

The change applies the existing vector-lane bounds predicate during:

• coarse histogram construction;
• direct candidate collection;
• initial refinement histogram and candidate spill.

A regression test covers rows where the valid input prefix contains fewer finite values than top_k, making the top-k threshold -inf.

Why

The kernel uses predicated vector loads and fills out-of-bounds register lanes with -inf. However, those lanes were still counted in the radix histogram and could subsequently be collected as top-k candidates.

When the real top-k threshold was also -inf, the kernel could not distinguish valid in-range -inf values from the synthetic OOB values. This produced thousands of invalid candidates, overflowed the per-row extra buffer, and resulted in an illegal memory access.

Filtering candidates by their logical input index fixes the issue while preserving valid top-k semantics for real in-range -inf values.

Related issues

None.

API and compatibility impact

None. There are no public API or input-contract changes.

The kernel now excludes internal OOB vector lanes from top-k selection. Valid in-range values, including -inf, retain their existing top-k semantics.

Testing

Tested on an NVIDIA B200 (SM100) with CUDA launch blocking enabled.

• python -m pytest -q --tb=short fe_api/dsa/test_DSA_indexer_top_k.py
  • 5 passed
• Replayed the original failing input:
  • shape: (633, 768)
  • dtype: float32
  • seq_lens=633
  • top_k=512
  • completed without IMA
• Replayed the original failing input for 100 consecutive iterations:
  • all iterations completed without IMA
• git diff --check
  • passed

Summary by CodeRabbit

• Bug Fixes

  • Corrected top-k selection to ignore padded vector lanes during histogram processing and index collection.
  • Prevented out-of-range padded values from influencing top-k results, including inputs containing negative infinity.

• Tests

  • Added coverage verifying returned indices remain within valid sequence bounds and match expected top-k values.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3debd53b-392c-467f-b9ef-28be0ca39f1f

📥 Commits

Reviewing files that changed from the base of the PR and between 3772b8b and 782c445.

📒 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

---

📝 Walkthrough

Walkthrough

The varlen top-k kernel now excludes vector-tail padding from histogram construction, index collection, and refinement. A CUDA regression test verifies that negative-infinity padding does not affect returned indices or selected values.

Changes

Varlen top-k padding 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	Stage 1 histogram updates, coarse index emission, and refinement operations now run only when each lane’s relative index is below aligned_size.
Validate negative-infinity padding behavior test/python/fe_api/dsa/test_DSA_indexer_top_k.py	Adds a CUDA-gated test that checks returned indices remain in range and selected values match torch.topk over real input columns.

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

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
Title check	✅ Passed	The title clearly states the main fix: preventing IMA from out-of-bounds lanes in the varlen indexer top-k path.
Description check	✅ Passed	The description matches the template and includes the required area, summary, why, related issues, compatibility, and testing sections.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[Anerudhan]

@cudnn-ci-bot run