Add fused SBHD compressed sparse attention

[FDecaYed]

Add fused SBHD compressed sparse attention

Summary

Add the high-performance fused SBHD backend for DeepSeek-V4 CSA on top of the current-main DSA backend contract.

Scope

• Add CSA-namespaced sparse-attention and fused-indexer helpers in the post-[dev] Refactor CSA structure: Move CSA implementation helpers into csa_utils dir #6372 csa_utils/fused_sparse_attention.py layout.
• Dispatch CompressedSparseAttention through the fused SBHD path when the canonical DSA backend router selects cuDNN.
• Preserve the full compressed + window + sink teacher semantics required by [DSv4] Use the full CSA denominator for unfused indexer loss #5960 in the fused indexer-loss integration.
• Resolve an omitted CLI backend to cudnn for dsv4_hybrid and to none for ordinary DSA; an explicit none remains authoritative.
• Reject TileLang for DSv4, require SM90+ for the cuDNN path, and guard the unsupported SM90 dense-indexer-loss combination.
• Add focused layout/indexer/loss/namespace tests and an SM90+-gated real fused-versus-native forward parity test.

Production diff: 6 files, +1,206/-3. Test diff: 2 files, +182/-4.

Non-goals

• Replacing or duplicating ordinary dsa.py, dsa_kernels.py, dsa_cudnn_kernels.py, dsa_indexer_loss.py, dsa_layout.py, or dsa_masking.py.
• Adding ambiguous CSA implementations under DSA symbol names.
• Adding a separate fused compressor implementation.
• TileLang support for DSv4, packed THD, context parallelism, HybridModel construction, or recipes.

Review boundary and dependencies

Requires #6402 (D2). The local branch carries the complete D2 line through 0ecb394a67d2; this remains an independent PR targeting main.

• Own commit: 5f0d068ad3d289ac2bc44ce1b3056c0d04fcdcf3
• Exact own diff: 8 files, +1,388/-7
• Temporary carried commit boundary: 0ecb394a67d2ae2b89c20d8d489c48fa5b5b1ef9

Please review the own commit only. After D2 merges, it will be replayed directly onto then-current main.

Provenance

Uses current-main ordinary DSA as source of truth, #6020 for the DSA/CSA ownership contract, #6372 for the current CSA helper layout, and #5960 for teacher correctness. #5795 is a historical feature/test oracle only.

Validation

Passed locally:

• clean worktree and no dependency/lock-file changes;
• signed-commit verification and required Signed-off-by trailer;
• Ruff, Python bytecode compilation, and git diff --check on the complete own diff;
• isort and Black 26.3.1 on the new core/kernel/test files;
• focused API/namespace and obsolete-import scans.

The two existing training argument files retain their parent branch's whole-file Black/isort debt; D3 adds no imports there and does not reformat unrelated code. Runtime pytest is blocked on this macOS ARM host because the cuDNN Frontend package is Linux-only and CUDA/SM90 hardware is unavailable. The real fused/native parity test is included and gated on CUDA, SM90+, cuDNN DSA, and FlashMLA; it and full functional CI must run after publication.

Required GPU evidence before merge:

• SM90 ratio-4 training/backward with sparse indexer loss;
• SM90 ratio-4 inference with compact top-K;
• SM100 dense indexer loss and per-token-loss scaling;
• ratio-0 and ratio-128 no-indexer paths;
• output, loss, and gradient parity against explicit dsa_kernel_backend='none' using the intended DSv4 production shape (d_v=512 and supported head/top-K values).

Sibling rebase constraint

D3 and #6405 (D4) are independent and can merge in either order, but both touch megatron/training/argument_utils.py and megatron/training/arguments.py. If D4 merges first, D3’s rebase must preserve D4’s C/H/W variant inference and compression-ratio normalization while adding the backend-default behavior here.

Local artifact

D3-fused-sbhd-csa.patch, generated from the own commit only. Verify it against SHA256SUMS.

Before publication

After D2 merges, replay only D3’s own commit onto current main, re-audit backend APIs and optional dependencies, preserve any already-merged D4 training-argument logic, then run the focused SM90+ kernel parity test and full functional CI.

[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.