[dev] [DeepSeek-V4] Add compact BF16 and MXFP8 DSA indexer support

[hxbai]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds the SM100 compact cuDNN indexer forward + Top-K path for DeepSeek-V4 compressed sparse attention, including BF16 and MXFP8 execution, CUDA-graph-safe workspaces, and THD context-parallel support.

Rely on NVIDIA/cudnn-frontend#370

Key changes

• Dispatches the compact indexer on SM100 independently of whether indexer loss is disabled, sparse, or dense.
• Preserves the existing non-compact BF16 path on unsupported hardware or dependencies.
• Adds --dsa-indexer-precision={bf16,mxfp8} and validates the MXFP8 hardware, geometry, loss, and dependency requirements.
• Uses Transformer Engine for BF16-to-MXFP8 quantization and a Triton kernel to pack scale factors into the BSHD/THD layout required by cuDNN.
• Adds caller-owned BSHD and THD workspaces for CUDA graph capture, including candidate, output, softmax, offset, and MXFP8 quantization buffers.
• Adds compact THD support under context parallelism by padding K segments and mapping compact physical indices back to logical global indices.
• Propagates Megatron's deterministic mode to compact Top-K tie breaking.
• Updates the cuDNN Frontend development pin to the revision providing the required compact APIs.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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

[kunlunl]

/claude strict-review

[claude]

[IMPORTANT Implementation] THD MXFP8 scale prefixes are not refreshed across CUDA-graph replays, unlike cand_batch_offsets.

What: For the THD compact path with a workspace, cu_seqlens_q_scale_padded / cu_seqlens_k_scale_padded are read verbatim from the workspace (lines 1286–1289), which was populated once during eager warmup from the warmup-time cu_seqlens_q (see prepare_thd_compact_indexer_workspace, line 557–558). But quantize_indexer_mxfp8 is then called with the live cu_seqlens=cu_seqlens_q (line 1300) together with these stale padded prefixes. Inside _pack_indexer_mxfp8_scale_thd_kernel, seq_start/seq_len come from the live cu_seqlens_ptr while scale_row_start comes from the stale cu_seqlens_scale_padded_ptr — the two must be mutually consistent for the byte placement to be correct.

Why it matters: Just above, _refresh_thd_compact_cand_batch_offsets (line 1242) re-derives cand_batch_offsets from the live cu_seqlens_q on every dispatch, and the docstring at prepare_thd_compact_indexer_workspace line 548 explicitly states "Sequence boundaries may change between graph capture and replay while tensor shapes and maxima remain static." If that contract holds, then when per-segment boundaries differ on replay, the scale prefixes go stale: scales land at wrong physical rows, and because the per-sequence padding (make_indexer_mxfp8_scale_cu_seqlens rounds each segment up independently) can change the padded total, the preallocated q_scale/k_scale buffer can even be undersized → out-of-bounds writes. The BF16 path is unaffected (no scales), which is why existing THD graph tests wouldn't catch it.

Fix: Refresh the padded scale prefixes from the live cu_seqlens_q/cu_seqlens_kv on each dispatch (writing into caller-owned workspace storage, mirroring _refresh_thd_compact_cand_batch_offsets), and size q_scale/k_scale for the geometry-wide worst case rather than the warmup boundaries. Alternatively, if the THD MXFP8 path is only ever entered with statically-padded cu_seqlens_q (constant across replays), assert that invariant here and document why the scale prefixes need no refresh while cand_batch_offsets does.

[hxbai]

fixed

[claude]

Review Summary — Compact BF16 & MXFP8 DSA indexer support

Reviewed the full diff with focus on implementation correctness, TP/PP/SP/CP parallelism, CUDA-graph capture/replay, backward compatibility, Megatron Core process-group usage, and unused new identifiers.

Findings by severity

IMPORTANT (1)

• THD MXFP8 scale prefixes not refreshed across CUDA-graph replays (dsa_kernels.py:1284–1289) — inline comment posted. The compact THD path re-derives cand_batch_offsets from the live cu_seqlens_q on every dispatch but reuses the warmup-time cu_seqlens_q_scale_padded/cu_seqlens_k_scale_padded verbatim, while passing the live cu_seqlens_q into quantize_indexer_mxfp8. If per-segment boundaries change between capture and replay (a contract the code documents at line 548), scales are placed at incorrect physical rows and the preallocated scale buffers can be undersized. Either refresh the scale prefixes per-dispatch (as with cand_batch_offsets) or assert/document that the MXFP8 THD path only ever sees statically-padded cu_seqlens_q.

Areas checked and found clean

• New config field dsa_indexer_precision: Literal["bf16", "mxfp8"] — auto-converts to --dsa-indexer-precision via ArgumentGroupFactory (Literal origin handled correctly), defaults to "bf16" (backward-compatible), and __post_init__ gates MXFP8 on apply_dsa_kernel_fusion, SM100+, 64 heads / head_dim 128, sparse-loss-only, and a compatible cuDNN Frontend wrapper signature (inspect.signature parameter check). No API/checkpoint-format break.
• Changed return signatures (_indexer_topk_core → 4-tuple, indexer_topk 2/3-tuple, compute_cp_indexer_topk 3-tuple, FusedIndexerSparseAttnFromTopkFunc.forward +3 args) — all production callers updated; the 2-tuple unpacks at csa.py:2373 and the CP caller are safe because return_softmax defaults False. Backward returns exactly match the extended forward arg count (18 args → 19 None/grad entries incl. ctx).
• max_seqlen_kv += 2 / + 2 invisible-padding-row accounting — consistent between build_thd_compact_k_layout (one appended zero row per real K segment + capacity-tail distribution), the workspace-prep max_seqlen_k args, and compute_cp_indexer_topk (line 386). The appended rows are causally unreachable so returned local ids are unchanged.
• MXFP8 Triton scale packing — BSHD and THD kernels map logical (row, scale_group) to the Blackwell F8 128×4 physical layout; bounds are masked (valid/in_bounds), THD batch lookup uses a correctly-sized binary search (SEARCH_STEPS = bit_length), and out-of-range groups/tokens load 0. Input validation in pack_indexer_mxfp8_scale/quantize_indexer_mxfp8 is thorough (dtype, contiguity, device, 128×4 divisibility, THD-vs-BSHD kwarg exclusivity).
• CUDA-graph capture safety — workspace matches()/validate() check device/dtype/shape/contiguity (and sequence values only when not capturing); sizing helpers and .item() syncs are confined to eager prepare_* paths and explicitly forbidden under capture; the compressor GEMM fp8-disable is applied in both SBHD (csa.py:960) and THD (csa.py:1062) paths per the HEAD fix.
• Loss restructure (step 5) — early loss_coeff <= 0 short-circuit, padding-row masking without GPU→CPU sync (via data_ptr comparison of unpadded vs padded cu_seqlens_q), and compact_predict reused directly when the compact softmax is available, with a correct dense-recompute fallback. calculate_per_token_loss divisor handling matches the reference compute_dsa_indexer_loss.
• Megatron Core process groups — no new direct parallel_state.get_*_group() reads introduced in the touched megatron/core files; CP paths accept an explicit cp_group / pg_collection.cp.
• Unused new identifiers — spot-checked new args/vars (compact_workspace, precision, deterministic, return_softmax, compact_predict, q_padding_mask, workspace_max_seqlen_k, topk_max_seqlen_k, the quantization helpers) all have live read/use paths.

Overall risk assessment

Low–moderate. The change is large but well-isolated behind the new mxfp8 precision and the compact-wrapper availability checks; BF16 behavior is preserved and the default is unchanged. The one IMPORTANT finding is confined to the MXFP8 + THD + CUDA-graph-replay combination — please confirm the boundary-static invariant or add the per-dispatch refresh before landing. Extensive new unit/functional test coverage (workspace warmup/reuse, capture/replay, MXFP8 dispatch and dequant reference, CP parity) accompanies the change.

[hxbai]

/claude strict-review

[claude]

Review Summary — Compact BF16 & MXFP8 DSA Indexer (#5992)

I performed a strict correctness/distributed/performance/back-compat review of the six production files (indexer_quantization.py, csa.py, csa_cp_utils.py, csa_kernels.py, transformer_config.py, arguments.py) plus the accompanying tests.

Findings by severity

• CRITICAL: 0
• IMPORTANT: 0
• SUGGESTION: 0

No blocking issues found.

What was verified

• Dead code / unused args (mandatory check): every new parameter, workspace field, and return-tuple element is consumed. compact_predict, deterministic, precision, return_softmax, use_compact, compact_workspace, indexer_precision, the _*_compact_indexer_workspaces lists, and _active_* slots all have live readers.
• Autograd arity: FusedCSAIndexerSparseAttnFromTopkFunc.forward takes 19 inputs; backward returns exactly 19 gradients (3 attn grads + None + 3 indexer grads + 12 None). Matches.
• Variable return arity: indexer_topk returns a 2- or 3-tuple gated on return_softmax; all three production callers (csa.py:1990, csa.py:2431, csa_cp_utils.py:439) unpack the correct shape (the two non-CP callers omit return_softmax, defaulting to False).
• CUDA-graph capture safety: workspaces are sized/prepared only during eager warmup (host syncs guarded by is_current_stream_capturing()), validated during capture via metadata-only matches/validate (no .item()), and the compact kernel's output buffers are asserted to alias the caller-owned storage by data_ptr().
• Buffer sizing: THD scale capacity (indexer_mxfp8_thd_scale_capacity) and candidate buffer (max(cand_floats, total_q·max_seqlen_k)) are provable upper bounds over any in-shape sequence layout.
• Triton scale packers: BSHD kernel masks source loads (other=0) and bounds stores; THD kernel's grid + 512-byte tiling makes every store in-bounds by construction; the binary search over cu_seqlens_scale_padded uses bit_length() steps.
• Precision fallback: MXFP8 hard-errors when compact support is absent; BF16 warns and falls back to the dense forward → radix Top-K path, which returns dense scores so the sparse-loss branch can still gather predict.
• Config validation: dsa_indexer_precision MXFP8 gate correctly requires cudnn backend, SM100+, 64 heads / head_dim 128, and sparse loss; cuDNN Frontend wrapper compatibility is probed via inspect.signature (import present).
• Megatron Core process groups: no new direct parallel_state.get_*_group() reads introduced.
• Rename hygiene: _build_cp_indexer_layout → build_cp_indexer_layout with no stale references anywhere in the tree.

The implementation is defensively coded and internally consistent. Nice work.