[buptzyb]

Summary

Main counterpart of #8534.

cuDNN frontend 1.26 (NVIDIA/cudnn-frontend#316) changed two defaults that the DSA cuDNN call sites rely on. main pins FE 1.29.1, so both issues are live on main. Both are silent: no error, just wrong gradients or a wrong top-k selection.

1. Sparse indexer backward: top-k ids are treated as per-batch local

Since FE 1.26, indexer_backward_wrapper treats top-k ids as local to each batch unless topk_indices_global=True. The sparse-indexer-loss path passes global ids (b * S_k + local). With BSHD and more than one sequence per micro-batch, every batch after the first loses its sparse indexer gradient on FE 1.29.

The wrapper is now called through _call_sparse_indexer_backward. It sets topk_indices_global=True whenever the installed wrapper accepts that parameter (or takes **kwargs), and keeps the old call otherwise.

2. Causal offsets: the default alignment moved from bottom-right to top-left

When q_causal_offsets is not given, FE <= 1.25 aligns the causal mask bottom-right (seqlen_k * ratio - seqlen_q), and FE >= 1.26 aligns it top-left (offset 0). On main, three callers relied on the old default:

• The single-sequence packed CP segment scorer (_indexer_topk_single_packed_cp_segments). It crops each query chunk's key prefix so that bottom-right masking gives each row its absolute causal length.
• The multi-sequence packed CP THD scorer (_indexer_topk_multi_packed_cp_thd). Its varlen call gives each query segment a longer key prefix and passed no offsets. On dev, [dev](perf): DSA pre-attention: plan-driven scorer dispatch, host-built CP layout metadata #6206 already passes these offsets; main does not have [dev](perf): DSA pre-attention: plan-driven scorer dispatch, host-built CP layout metadata #6206.
• The BSHD dense score recompute and dense indexer backward.

With FE 1.26 or newer, the bottom-right offsets are now passed explicitly: a uniform offset for BSHD calls, and per-segment offsets (segment_k_lengths - segment_q_lengths) for the THD call. Older frontends and zero-offset calls keep their exact call and kernel compile key. A top-left frontend whose wrapper lacks q_causal_offsets fails closed.

Differences from the dev PR

• main has no utils.filter_kwargs_for_callable, so the wrapper signature is probed locally (_wrapper_accepts_kwarg).
• The THD call site above (item 2, second bullet) is fixed here; on dev it is already correct.
• The tests call main's _indexer_topk_single_packed_cp_segments signature.

Test plan

• CPU tests in test_dsa_cudnn_fe_compat.py, 23 passed. They use fakes of both frontend generations (dsa_cudnn_fe_fakes.py) and cover:
  • version parsing and the 1.25/1.26 gate;
  • fail-closed behavior;
  • the cropped segment scorer and the multi-sequence packed CP THD scorer selecting identical top-k under either default;
  • a negative control for each scorer: with the offsets removed on FE 1.31, the selection changes.
• In test_dsa_native_parity.py, the fake indexer_backward_wrapper used by the sparse-loss tests requires topk_indices_global=True.
• CPU: test_dsa_native_parity.py and test_attention_variant_dsa.py give the same results on this branch and on main. The 3 failures and 34 errors are identical on both and come from the CPU environment.
• GPU check on GB300 with cuDNN FE 1.29.1, main's pin. One standalone script ran against main and against this branch, calling the affected functions directly:

  • Sparse indexer backward, BSHD micro-batch 2. The reference is the same rows run as two b=1 calls, scaled for the row-mean loss.

    • main: batch 1 gets all-zero dQ/dW/dK.
    • This branch: both batches match (dK relative error ≤ 8e-6).

  • Top-k recall against an fp32 torch top-k over each row's absolute causal keys (mean / min):

    Scorer	main	This branch
    Single-sequence segment scorer, CP=1	0.706 / 0.266	1.000 / 1.000
    Single-sequence segment scorer, CP=2, rank 0	0.587 / 0.086	1.000 / 1.000
    Single-sequence segment scorer, CP=2, rank 1	0.332 / 0.109	1.000 / 1.000
    Multi-sequence packed CP THD, CP=2, rank 0 (documents of 1024 and 2048 tokens)	0.579 / 0.039	1.000 / 1.000
    Multi-sequence packed CP THD, CP=2, rank 1	0.289 / 0.055	1.000 / 1.000

    On main, the selected keys stay inside each row's bounds, but tens of thousands of them score below the reference k-th score: top-k ran over a top-left-masked key range.

• CHECK_ONLY=true BASE_REF=main tools/autoformat.sh (black 26.3.0, isort, pylint, ruff)

Overlap with open PRs

• feat(dsa): add FP32 RoPE and FP8 scoring options for the cuDNN indexer #8403 edits the same indexer-forward call sites in dsa_cudnn_kernels.py and conflicts textually; whichever lands second rebases.
• perf(dsa): compute indexer scores with the cuDNN kernel #8154 merges cleanly with this change.
• Move DSA, CSA, absorbed MLA and DeepSeek-v4 attention into megatron/core/ops with deprecated aliases (3/4) #8348 moves dsa_cudnn_kernels.py; the move would carry these hunks.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.