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]

📝 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
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 main fix: preventing illegal memory access from out-of-bounds lanes in varlen top-k.
Description check	✅ Passed	The description follows the template well and includes the required sections, rationale, compatibility impact, and testing details.

✨ 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

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-410-782c445
Pipeline: 59853073

[JackRao123]

We independently hit this exact bug in production and can confirm the diagnosis — thanks for fixing it.

Our data point: a GLM-5.2 context-parallel (CP32) training run on B200s crashed with Xid 43 (illegal memory access) mid-forward_backward, roughly once per several hours of long-context training. CUDA_LAUNCH_BLOCKING=1 attributes the fault to cute_dsl_topk_wrapper; the captured input (a (862, 4310) fp32 score chunk, top_k=2048) reproduces the crash deterministically in a fresh process — same call, same iteration, on both B200 (SM100) and B300 (SM103); compute-sanitizer shows the invalid 2-byte shared writes from the candidate-buffer overflow.

One extra trigger route worth noting alongside yours: our rows had plenty of finite values — the threshold landed in the -inf bin because to_coarse_key converts fp32 scores to fp16 first, so every score below −65504 collapses into the fp16 -inf bin with the phantom lanes. (2242 real + ~4072 phantom candidates vs 512 + 4310 buffer capacity in the large-occupancy compile.) Your fix covers both routes since it filters by logical index.

We verified an equivalent patch on our side: previously-crashing inputs pass with exact torch.topk parity, sanitizer-clean, across B200 and B300. Happy to contribute our regression test (which includes the hard-crash variant) if useful, but whichever lands first works for us.