[RayenTian]

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

What does this PR do?

Adds two orthogonal, opt-in precision options for the cuDNN DSA (DeepSeek Sparse Attention) lightning indexer. Both default to the existing behavior, so models that don't set them are unaffected.

• dsa_indexer_rope_fp32 (bool, default False) — keeps the indexer query/key preprocessing (key LayerNorm → RoPE) in FP32 end-to-end, with no intermediate BF16 round-trip before the scoring matmul.
• dsa_indexer_scoring_precision ("bf16" | "fp8", default "bf16") — runs the cuDNN indexer scoring matmul in BF16 (unchanged) or FP8 (E4M3 operands with UE8M0 per-vector scales). "fp8" requires dsa_kernel_backend="cudnn" and dsa_indexer_loss_coeff=0; SM90+ is enforced by the existing cuDNN backend dependency check.

Motivation

In RL post-training (e.g. NeMo-RL) that uses vLLM as the rollout/generation backend, the indexer's top-k selection must match between training and generation. The top-k boundary is sensitive to the precision of both the q/k preprocessing and the scoring matmul:

• Megatron's default path materializes BF16 between the key LayerNorm and RoPE, and runs the scoring matmul in BF16.
• The FP8 inference indexer keeps preprocessing in FP32 and scores in E4M3.

That precision gap flips some top-k picks near the boundary, so training recomputes log-probs over a different key set than generation actually selected — a train/inference mismatch that hurts RL convergence.

These two options make that precision configurable so the training side can reproduce the FP8 inference indexer's top-k selection bit-for-bit. The switches are orthogonal: dsa_indexer_rope_fp32 removes the BF16 round-trip in preprocessing, dsa_indexer_scoring_precision controls the scoring matmul precision. Megatron Core only exposes the mechanism (both default off); enabling them for a given model/runtime is a downstream concern — e.g. NeMo-RL turns them on for its vLLM rollout path.

Implementation

• All cuDNN indexer scoring calls are routed through a single _indexer_forward choke point in dsa_cudnn_kernels.py, which dispatches on scoring_precision (BF16 cast vs. E4M3 quantization + precision="fp8" / q_scale / k_scale). scoring_precision is threaded through every top-k path (including the pure-torch global-row path) rather than monkey-patching the cuDNN namespace at runtime.
• forward_before_topk keeps q/k in FP32 under dsa_indexer_rope_fp32.
• Config validation gates "fp8" to the cuDNN backend with zero indexer loss.
• The default path (BF16, no rope_fp32) is byte-identical to before — covered by a regression test.

Note: this adds an FP8 dispatch through the cuDNN DSA external library. Bit-exact determinism coverage under tests/unit_tests/determinism/kernels/ is being evaluated (see the determinism pre-check below).

Test

End-to-end (training)

Split Path

Validated both options on a real MoE + MLA + DSA training path, using a toy GLM-5 (glm_moe_dsa) model built through Megatron-Bridge:

• Model/setup: 2 decoder blocks, hidden 1024, 8 routed experts, DSA indexer with index_topk=256, seq_length=2048 (> index_topk, so the sparse top-k truncation is actually exercised), random weights, mock data, 2×H100 (TP=1 / PP=1 / EP=2).
• Runs: three matched configurations trained for a few hundred steps each, differing only by the new switches:
  1. baseline — bf16 scoring, no rope_fp32;
  2. dsa_indexer_rope_fp32=true;
  3. dsa_indexer_scoring_precision=fp8 + dsa_indexer_loss_coeff=0 (cuDNN backend).
  4. dsa_indexer_rope_fp32=true + dsa_indexer_scoring_precision=fp8 + dsa_indexer_loss_coeff=0 (cuDNN backend).
• Checked:
  • the options propagate to the DSA indexer config (dsa_kernel_backend, scoring_precision, rope_fp32, loss_coeff);
  • the fp8 scoring path is actually hit at runtime — the E4M3 quantization/round-trip of q/k runs inside the cuDNN indexer top-k path (verified with temporary instrumentation on the scoring choke points, both the single-chunk _indexer_forward and the chunked global-row path);
  • under rope_fp32, the indexer q/k carry FP32 through preprocessing;
  • the indexer top-k selection changes under fp8 vs. bf16 on identical inputs, while the default path is unchanged;
  • all runs train stably (loss decreases, no NaN / divergence).

Fused Path

Validated the same two options on the fully fused cuDNN DSA path (run_fused_dsa_attention → FusedIndexerSparseAttnFunc, which fuses indexer scoring + FlashMLA sparse attention + the optional indexer loss into one autograd Function), again on a toy GLM-5 (glm_moe_dsa) through Megatron-Bridge:

• Reaching the fused path: the fused hook only engages for absorbed MLA with a FlashMLA-compatible compressed-KV width, so the toy was sized with kv_lora_rank=512 (key width 576) and dsa_indexer_n_heads=32; 2×H100 (TP=1 / PP=1 / EP=2), mock data, 200 steps.
• Runs: bf16 vs. fp8 indexer scoring, both with dsa_indexer_loss_coeff=0 on the cuDNN backend, differing only by dsa_indexer_scoring_precision.
• Checked:
  • the fused hook is actually taken at runtime (not the split-topk fallback): run_fused_dsa_attention returns a result and FusedIndexerSparseAttnFunc.forward runs on every rank — confirmed with temporary instrumentation;
  • scoring_precision is threaded all the way into the fused forward (shows bf16 / fp8 respectively);
  • the fused indexer top-k selection changes under fp8 vs. bf16 on identical inputs (different top-k index checksums), confirming the E4M3 quantization is effective on this path;
  • both runs train stably over 200 steps (loss decreases from ~12.2, no NaN / divergence).
• Note: the fp8 fused indexer kernel requires the indexer head count to be one of {32, 64} (bf16 additionally allows 16); dsa_indexer_n_heads=16 runs in bf16 but asserts under fp8 (qhead_per_kv_head must be one of (32, 64)).

End-to-end (RL)

Undergoing...

Issue tracking

Linked issue: N/A — incremental DSA indexer precision option.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• 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.

[RayenTian]

/review

[RayenTian]

/ok to test ba04da6

[xuantengh]

/ok to test ba04da6

[xuantengh]

/review

[svcnvidia-nemo-ci]

Static review completed for all five changed files, including the full-fusion and split indexer callers. No tests were executed.

1. The fully fused attention path ignores FP8 scoring

Failure: With dsa_kernel_backend="cudnn", dsa_indexer_scoring_precision="fp8", zero indexer loss, and supported absorbed-MLA shapes, DSAttention.forward attempts full fusion before the newly updated split top-k hook. That route reaches FusedIndexerSparseAttnFunc.forward, which calls _indexer_topk_bshd without scoring_precision. It therefore scores in BF16 and returns before the FP8-aware split path runs. The advertised precision switch silently has no effect on this supported route.

Correction: Either thread the precision through the full-fusion wrapper and autograd contract, including its backward return arity, or deliberately use the already precision-aware split selector for FP8. The smallest correction is this insertion immediately after _assert_supported_indexer_scoring(use_relu) in the cuDNN full-fusion hook:

if getattr(config, "dsa_indexer_scoring_precision", "bf16") == "fp8":
    return None

This lets DSAttention continue through the split FP8 selector while retaining fused sparse attention afterward. It also preserves compatibility with existing test config stubs that lack the new field.

Preservation and acceptance: Add a proposed DSAttention integration regression using attention_backend=auto, no index sharing/KPool, causal absorbed MLA with value width 512, sequence length 128, and top-k 64. Spy on scoring and assert that FP8 configuration reaches FP8 scoring, including with rope_fp32=True; exercise both plain causal and packed inputs. The BF16 configuration must still use full fusion. Retain split-attention backward handling of invalid slots and supplied lengths: the existing sanitization regression must continue to pass. Suggested code requires author verification.

2. Global-row scoring does not honor BF16 precision for FP32 carriers

Failure: With dsa_indexer_rope_fp32=True and BF16 scoring, _indexer_forward rounds q/k to BF16 at the scoring boundary. However, the global-row helper leaves FP32 carriers unrounded. Thus explicit-varlen or chunked scoring uses a different operand precision. A minimal helper-only example with one head/dimension, q=k=1.001, weight 1, ratio 1, scale 1, and one valid key returns approximately 1.002001; BF16 scoring must return 1.0, because both operands round to 1.

Correction: Replace the current elif scoring_precision != "bf16" branch in that helper with:

elif scoring_precision == "bf16":
    q_chunk_bshd = q_chunk_bshd.to(torch.bfloat16)
    if k_bshd.dtype != torch.bfloat16:
        k_bshd = k_bshd.to(torch.bfloat16)
        k_bdk = None
else:
    raise ValueError(
        f"Unsupported DSA indexer scoring precision {scoring_precision!r}; "
        "expected 'bf16' or 'fp8'."
    )

Invalidating k_bdk for an FP32 key is necessary: the caller currently caches the unrounded key. The invariant is BF16-rounded scoring operands regardless of preprocessing carrier dtype or chunking.

Preservation and acceptance: Proposed tests should assert 1.0 for the example both without a cache and with a cache built from the original FP32 key. Already-BF16 inputs must retain their reusable cache and unchanged scores; FP8 must retain its E4M3 branch. Preserve global causal offsets and distinct rows with repeated varlen ends: the existing repeated-ends regression must still produce indices [[1,0],[0,1]], lengths [2,2], and scores [[2,1],[-1,-2]]. Suggested code requires author verification.

3. Add the required registered bit-exact replay coverage

Failure: This adds a new cuDNN FP8 kernel dispatch, but no determinism test or manifest update accompanies it. The added scoring tests compare implementations once or allow 98% top-k overlap; neither checks repeated byte-identical results. This is required by the kernel testing contract. Independently, the changed DSA preprocessing source is already registered to a determinism test that this PR does not update, which violates the coverage gate's test-update check.

Correction and acceptance: Add SM90+/dependency-guarded replay cases to the registered determinism test module, and register the cuDNN scoring source in the manifest. Using the same (1,128,32,128) q and (1,128,1,128) k setup as the added scoring tests, the scoring replay core can be:

for precision in ("bf16", "fp8"):
    assert_replays_bit_exact(
        lambda qq, kk, ww: kern._indexer_forward(
            qq, kk, ww, ratio=1, sm_scale=1.0, scoring_precision=precision
        ),
        (q.float(), k.float(), w),
        replays=3,
        backward=False,
        contention=True,
        what=f"DSA indexer {precision}",
    )

Also replay top-k outputs and the new FP32 preprocessing route, asserting that q/k remain FP32 while weights remain BF16. Retain BF16/default replay cases and the existing zero-centered-gamma behavior. Expected outcomes are byte-identical replay outputs and a passing manifest/coverage check—not merely approximate reference agreement. This is proposed coverage, not evidence of executed tests.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[RayenTian]

/ok to test a7dcf09

[copy-pr-bot]

/ok to test a7dcf09

@RayenTian, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[RayenTian]

/ok to test a7dcf09

[RayenTian]

/ok to test e16c82a

[RayenTian]

/ok to test 3f7e796

[svcnvidia-nemo-ci]

LGTM for this light follow-up. Verified the earlier findings against the current source:

• Full fusion now forwards scoring precision through the wrapper and autograd contract to indexer selection, with matching backward return arity.
• Global-row BF16 scoring now rounds FP32 carriers and invalidates the unrounded key cache; the added regression covers cached and uncached inputs.
• The cuDNN scoring dispatch now has BF16/FP8 bit-exact replay coverage and a corresponding manifest entry.

No unresolved prior findings or eligible new regressions identified. Review was limited to the prior findings, supplied fix evidence, and directly affected paths; the previously reviewed tree was unavailable for an exact inter-review comparison. Static inspection only—tests were not executed.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a findings-only review. Request a full review with /review reset=true.