DSA indexer_top_k: fall back to scalar stores for odd top_k

[JackRao123]

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

Cap the DSA indexer_top_k output-store vector width at 1 when top_k is odd, so odd top_k values JIT-compile instead of failing assert self.top_k % vecsize_out == 0. Adds top_k=705 (odd, greater than threads-per-CTA) to the indexer_top_k test parametrization as a regression case.

Why

Fixes #406. check_support() and the documented constraint permit any top_k in (0, 2048], but the output phase selected vecsize_out = 2 whenever top_k > num_threads_per_cta, and the width-2 vectorized index/value store has no scalar tail — so odd top_k above that threshold asserted at compile time. Real workloads hit this when clamping top_k = min(index_topk, visible_keys) on context-parallel shards (see the issue for the exact [235, 705] / top_k=705 production case).

Scalar stores (vecsize_out == 1) are the path every top_k <= num_threads_per_cta config already takes; this change extends that path to odd top_k rather than adding a tail case to the vectorized writer. Only the output-phase store width changes — the selection phases and all even-top_k codegen are untouched.

Related issues

Fixes #406

API and compatibility impact

None for previously-working configs: vecsize_out is unchanged for even top_k and for odd top_k <= num_threads_per_cta. Previously-asserting odd top_k values now compile, using scalar output stores.

Testing

• black --line-length 160 on both changed files.
• I don't have an SM90/SM100 machine available in this environment, so I could not run the GPU tests myself. The added top_k=705 parametrization exercises the previously-asserting configuration in both test_DSA_indexer_top_k_compile_execute and test_DSA_indexer_top_k_wrapper (default shapes give 256–512 threads per CTA, so 705 previously resolved vecsize_out = 2 and failed the assert). Happy to iterate on CI results or run additional cases on a B200 and report back.

Summary by CodeRabbit

• Bug Fixes

  • Fixed top-k attention processing for odd top-k values, improving correctness when results exceed standard vectorized widths.

• Tests

  • Expanded coverage to include an additional odd top-k scenario, helping prevent regressions in variable-length attention processing.

[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: 69523202-abc3-4fc1-9751-8c8b8b414412

📥 Commits

Reviewing files that changed from the base of the PR and between 3041f3e and f0c4b32.

📒 Files selected for processing (2)

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

---

📝 Walkthrough

Walkthrough

The indexer output phase now uses scalar stores for odd top_k values. DSA parameterized tests add top_k=705 to cover the regression case.

Changes

Odd top-k support

Layer / File(s)	Summary
Output store parity handling and regression coverage python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py, test/python/fe_api/dsa/dsa_utils.py	vecsize_out uses 2 for even top_k and 1 for odd top_k; DSA tests now include top_k=705 and document the regression condition.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: odd top_k now falls back to scalar stores.
Description check	✅ Passed	The description fills all required sections and clearly explains the change, rationale, impact, and testing.
Linked Issues check	✅ Passed	The change addresses #406 by handling odd top_k values with scalar stores and adding a 705 regression test.
Out of Scope Changes check	✅ Passed	The changes are tightly scoped to the reported DSA top_k bug and its regression test.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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