[JackRao123]

What

_compute_indexer_scores_chunk_with_global_rows computed indexer scores with a per-head fp32 PyTorch loop: bmm, then relu_, mul_ and add_ for each indexer head, all on SIMT rather than tensor cores. It serves two callers:

• _indexer_topk_from_score_chunks, whenever a chunk doesn't cover every row or carries explicit lengths: long causal calls that need chunking, and calls with varlen bounds;
• the full-score varlen branch of _indexer_topk_bshd.

This PR computes the scores with indexer_forward_wrapper instead, under the same name, and deletes the loop:

• Plain causal chunks: pass q_causal_offsets=row_start, which is the reference mask ((row_start + i + 1) // ratio).
• Chunks with explicit lengths: run with the tightest causal offset that covers every row, computed on device with no host sync. The exact per-row prefix is then applied, so masked positions match the reference exactly.

q_causal_offsets is available in the pinned nvidia-cudnn-frontend==1.29.0.

Numerics

• New test test_dsa_cudnn_indexer_scores.py: compares against a direct einsum reference over 9 cases: causal at ratio 1 and 4, unaligned row offsets, pooled bounds, empty rows, non-monotonic lengths, lengths beyond sk, packed resets.
• Result: the -inf mask is identical in every case; max abs diff ≤ 1.2e-4 on scores up to ~300.
• Environment: this branch on GB300 with nvidia-cudnn-frontend==1.29.0, 9/9 passing.

Performance

Measured on a downstream fork with a model whose indexer always passes explicit bounds (a 4:1 pooled-key indexer). 262k tokens on 8x GB300, CP4 EP8:

• Forward/backward step: 50.7 s → 32.7 s. The fp32 loop was about 15 s of it.
• Over 6 steps, loss and grad norm match the reference path within run-to-run noise.

Upstream models gain wherever they hit chunked scoring with explicit bounds, long causal calls, or the full-score varlen branch; I haven't benchmarked an upstream model.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[JackRao123]

Two follow-up commits:

• 6853e28e1: cuDNN frontend >= 1.27 rejects seqlen_q > seqlen_k * ratio, which pooled-key chunks with explicit seq_lens hit (e.g. 2312 query rows vs 578 keys at ratio 1). The ratio is now raised to ceil(sq / sk) before the causal offset is computed; seq_lens are still applied by the masked_fill_ after the kernel. Adds test cases with more rows than keys * ratio (they fail without the fix on cudnn-frontend 1.29.1).
• 5d2044ca8: updates the fake cuDNN indexer stubs in test_dsa_native_parity.py for the new q_causal_offsets argument.