dsa(indexer_backward): restore range_constexpr in 8 kernel_gemm epilogue loops

[zkyue]

Fixes #548 (SM100 DSA indexer backward: kernel_gemm
pure-kernel latency regressed 17.5-25.4% after the range_constexpr ->
range(..., unroll_full=True) migration).

What

Reverts exactly 8 loops in the kernel_gemm compute/reduce epilogues of
python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py
from cutlass.range(..., unroll_full=True) back to
cutlass.range_constexpr(...). No other changes; the one-shot barrier fix
from #426 is untouched. An analogous ninth migration site in
kernel_score_grad was measured separately and shows no runtime difference
(49.70 us vs 49.73 us at the anchor shape), so it is intentionally left
unchanged - see the issue.

Both forms fully unroll the loop body: range_constexpr unrolls at trace
time, while range(..., unroll_full=True) lowers through the DSL
structured-control-flow path. In this configuration the migrated form
produced the higher latencies in the issue; this PR restores the
pre-migration form for exactly these 8 loops. See the issue for the full
isolation trail (grad-signal-recipe insensitivity, 1.26.0 wheel anchor, and a
single-migration revert probe on the public commit 3fcee36).

Verification (B200 sm_100a, driver 590.48.01, CUDA 13.3, cutlass-dsl 4.6.1)

Performance - canonical seeded harness (B=1, H=64, D=128, S_k=4096), nsys
pure-kernel median, N=60 per variant, both variants interleaved per-iteration
in one profiling session; anchor shape S_q=8192/topk=1024:

build	kernel_gemm med	develop / build
develop @ 11c16ff	2056.9 us	1.000
this PR	1688.8 us	1.218
1.26.0 wheel (reference anchor, isolated session)	1691.8 us	-

Stated in both directions: develop has 21.8% higher latency than this PR;
this PR has 17.9% lower latency than develop. Across S_q {4096, 8192} x
topk {512, 1024, 1536, 2048} the same harness measures develop 17.5-25.4%
higher at every shape (full table with standard deviations in the issue).

Compile-time tradeoff (disclosed): range_constexpr raises time-to-first-call
with a cold DSL JIT cache from ~3.5 s to ~5.2 s at the anchor shape (one-time
per compiled shape); at 368 us saved per call it amortizes after ~4600 calls.

Numerics - no numerical regression observed:

• d_index_q, d_weights (kernel_gemm): bitwise identical to develop on
  the canonical harness at both S_q=8192/topk=1024 and S_q=4096/topk=512.
• d_index_k uses fp32 atomicAdd and is nondeterministic run-to-run: the
  develop-vs-PR difference (max_abs 6.7e-08, rms_rel 3.9e-07 at the anchor
  shape) is comparable in scale to same-build rerun differences
  (develop-develop max_abs 6.7e-08 / rms_rel 3.1e-07; PR-PR 6.0e-08 /
  3.3e-07).

Tests / lint:

• test/python/fe_api/dsa/test_DSA_indexer_backward.py and
  test_DSA_dense_indexer_backward.py: 3 passed.
• Full fe_api/dsa/: 81 passed, 7 skipped, 2 pre-existing failures in
  test_DSA_indexer_top_k.py that reproduce identically on unmodified develop
  (unrelated module).
• pre-commit run (pinned config: black 26.3.1, line-length 160): passed.

Summary by CodeRabbit

• Performance
  • Improved compile-time loop handling for sparse attention indexing operations.
  • Preserved existing computation, memory access, synchronization, and reduction behavior.

[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: 783f5bee-3a01-4a21-978e-50bc6bd49e24

📥 Commits

Reviewing files that changed from the base of the PR and between 721122b and a2478b4.

📒 Files selected for processing (1)

• python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py

---

📝 Walkthrough

Walkthrough

The SM100 indexer backward kernel replaces eight fully unrolled runtime loops with cutlass.range_constexpr. Computation, memory operations, synchronization, and reduction behavior remain unchanged.

Changes

SM100 indexer backward

Layer / File(s)	Summary
Compute warpgroup loop updates python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py	The dw_accum, dS, dQ, and dW loops now use compile-time iteration.
Reduction warpgroup loop update python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py	The dK pair-reduction loop now uses compile-time iteration.

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 identifies the restoration of range_constexpr in the eight kernel_gemm epilogue loops.
Description check	✅ Passed	The description provides the change, rationale, issue reference, compatibility impact, benchmarks, tests, and known failures; omitted template checkboxes are non-critical.
Linked Issues check	✅ Passed	The PR restores exactly the eight required loops, preserves the barrier fix, leaves the ninth loop unchanged, and documents performance, numerical, and compile-time effects for issue #548.
Out of Scope Changes check	✅ Passed	The changes are limited to the eight kernel_gemm loop bounds required by issue #548, with no unrelated code changes reported.
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.}

[Anerudhan]

@cudnn-ci-bot run oss

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-549-a2478b4
Pipeline: 62096947
Targets: oss