[HaochenYuan]

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

What does this PR do?

Under --deterministic-mode, a DSv4 run was not reproducible even against itself: the loss diverges from iteration 2 (grad norm from iteration 1). All sources are kernel-level and independent of CUDA graphs. This PR fixes them, gated on deterministic mode only; the default path is unchanged.

Source 1: cuDNN compact indexer Top-K returns the selected ids in a run-dependent slot order, and its compact
softmax is accumulated in that order (FlashMLA and the indexer backward consume both in slot order)
Fix: _canonicalize_compact_topk: sort ids per row (padding last), recompute the softmax in that order.

Source 2: cuDNN sparse-attention backward accumulates dK/dV with atomics
Fix: pass deterministic=True + a retained scratch workspace on all three sparse_attention_backward_wrapper call
sites; raise if the installed cudnn-frontend / GPU / head count cannot provide the deterministic kernel.

Source 3: cuDNN indexer backward accumulates dK with atomics
Fix: deterministic_sparse_indexer_grads_wk: atomic-free dK (dW on SM90) via torch's deterministic index_add (fallback: stable sort + prefix sum), padding slots spread over dummy rows; computes the same clipped-log KL gradient as cuDNN (predict · Σ kept_target − kept_target, targets clamped at e⁻¹⁰⁰) with per-sample key ids, and overrides cuDNN's dK on the fused and CP call sites (indexing and gradient-formula fixes by @Connor-XY, from #7706)

Source 4: TE fused grouped-MLP backward accumulates the routing-prob gradient (dprob) with atomics (newer TE
raises for this)
Fix: feed a unit scale to the fused op and apply the probs on the FC2 output (exact without FC2 bias; raises
otherwise).

Source 5: TE fused router aux loss (fused_moe_aux_loss) uses float atomicAdd
Fix: --moe-router-fusion is rejected in deterministic mode via ARG_VALUES_REQUIRED_FOR_DETERMINISM (explicit,
no silent kernel swap)

Relation to #7169: that PR aligns the id order between the no-grad and grad forward paths and stable-sorts the standalone Top-K fallback. It does not touch the compact cuDNN Top-K path used in training (slot and softmax order still vary run to run) nor any backward kernel; with #7169 alone the DSv4 proxy is still not reproducible. This PR builds on it.

• Unit tests: test_csa_deterministic_indexer.py (GPU): deterministic dK/dW vs an fp64 reference across B=1/2, fp32/bf16, padding, cross-sample chunks and both reduction paths, bit-identical on replay; on SM100 a direct comparison against cuDNN's indexer_backward_wrapper (rel. L2 < 1e-2 on rows whose targets sum to 0.5). test_indexer_topk_thd_compact_dispatch updated to expect the canonicalized softmax.
• Also validated at CP=2 (PP1×EP2×CP2, with and without --dsa-cp-balance-indexer): two runs bitwise identical on every logged scalar for 10 iterations. Not covered: CP>2, multi-node, mixed-length packed sequences.

Performance (iteration time vs the non-deterministic run, same kernels otherwise)

• --deterministic-mode with this PR: +14.8%, almost entirely PyTorch's torch.utils.deterministic.fill_uninitialized_memory, a default side effect of use_deterministic_algorithms(True) that adds one fill kernel per allocation (~40k launches/step here).
• Setting fill_uninitialized_memory = False has no effect on DSv4 reproducibility (bitwise identical losses with it on or off) and brings the total cost down to +6.1%. It is not changed in this PR because it is a global switch and was only validated on DSv4; users can set it in their launcher.

Notes

• Requires nvidia-cudnn-frontend ≥ 1.28 and SM100 for the deterministic sparse-attention backward; otherwise deterministic mode fails fast with a clear error.
• Numerics under deterministic mode differ slightly from the default path (fixed reduction orders; probs applied after FC2 instead of before); this is expected and only affects deterministic mode.
• Overlaps with [Dev] [DO NOT MERGE] GLM-5.2 integration stack and full-model benchmark tracker #6944 (deterministic DSA backward in dsa_cudnn_kernels.py, which dsv4_hybrid does not use).

⚠️ 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
• 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.

[Connor-XY]

Hi @HaochenYuan, two things we hit while stacking #7706 on this PR:

1. The deterministic dK/dW path uses predict - target, but cuDNN's indexer backward uses the clipped KL gradient predict * sum(kept_target) - kept_target. CSA targets don't sum to 1 (sink and local keys take mass), so deterministic mode changes the gradient, not only the reduction order. The fix is in Fix batch indexing and KL gradient in deterministic CSA indexer backward #7706 if you want to fold it in.
2. test_indexer_topk_thd_compact_dispatch fails on this branch and passes on dev. The conflict with dev is a single hunk where dev added _align_topk_width; keeping both resolves it.

If you rebase onto dev, feel free to fold #7706's commits in and I'll close #7706.

[HaochenYuan]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review: PR #7441

Findings: CRITICAL: 0 · IMPORTANT: 3 · SUGGESTION: 0
Overall risk: Medium. Through megatron/training, --deterministic-mode also turns on torch's global deterministic switch, and on that path the gating is consistent. The main gap: of the three CSA autograd functions, only one honours the Megatron Core TransformerConfig.deterministic_mode flag.

Most impactful

1. Deterministic gating is inconsistent. See fused_sparse_attention.py L1872, L3650 and L3903.
   • CSASparseAttnFunc (Path A: layers without an indexer) and FusedCSAIndexerSparseAttnFromTopkFunc (the CP path) only check torch's global switch.
   • So a run that sets only deterministic_mode=True in the config still uses the atomic dK/dV kernels on those layers, plus the atomic indexer dK on CP. Nothing reports it.
2. The workspace cache does not only grow, as its comment says. It keeps one buffer per distinct cuDNN workspace size and never frees any of them (L1585).
3. Deterministic mode with NCCL-EP zero-copy. In deterministic mode the routing probs are applied after FC2. Combine then receives a new tensor instead of the shared symmetric-memory buffer that zero-copy relies on (experts.py L855).

Checked, no issues found

• _deterministic_sparse_indexer_grads_wk:
  • per-sample key offsets for B>1, including chunks that cross a sample boundary;
  • the valid-slot mask and kept_target clipping at exp(-100);
  • the score gradient predict·Σkept_target − kept_target;
  • normalisation loss_coeff/(B·S_q)·sm_scale·grad_loss;
  • padding slots spread over dummy rows that [:total_keys] drops;
  • both reduction paths (index_add_ and stable sort + prefix sum).
• _canonicalize_compact_topk: padding sorts last, logits are permuted with the same perm, and rows that are all padding become 0 rather than NaN.
• experts.py: applying the probs on the FC2 output gives the same result as the fused op when FC2 has no bias, and the code raises when it has one. The determinism MoE test presets set add_bias_linear=False (configs.py L63), so the new guard does not break them.
• determinism.py: moe_router_fusion: False follows the existing check-only pattern.
• Unused-identifier check: every new identifier (ctx.deterministic, requested, compute_grad_w, det_w/det_k, post_probs) is read at runtime.

Scope notes

• This is the first bot review on this PR. I covered all six changed files.
• To check the findings I also read csa.py (the Path A caller), token_dispatcher.py (the NCCL-EP zero-copy contract) and the determinism test presets.
• style-guide.md is not in the base tree.
• The testing guide mentions a kernel manifest at tests/unit_tests/determinism/kernels/manifest.py. It does not exist in this tree, so I did not flag it.
• All conclusions come from reading the source; I ran no tests.

Review generated by claude

[FDecaYed]

I re-reviewed d3b32b3 and independently confirmed the three open findings. Requesting changes before approval:

1. Propagate the explicit deterministic configuration through every CSA path. CSASparseAttnFunc (no-indexer) and FusedCSAIndexerSparseAttnFromTopkFunc (caller-supplied top-k / CP) currently fall back to the global PyTorch switch, unlike the integrated indexer path. With TransformerConfig(deterministic_mode=True) and the global switch off, these paths can still use atomic backward kernels. Please pass/save the flag through the callers and autograd contexts, preserve existing apply-call arities, and test both paths with the global switch disabled. Existing thread: #7441 (comment)

2. Make scratch-workspace retention bounded and concurrency-safe. (device, required_bytes) permanently retains a separate allocation for every distinct size and does not reuse a larger allocation for a smaller request. Please use an explicit reuse/growth policy, preserve any pointers captured by live CUDA graphs, and account for overlapping CUDA streams rather than sharing a mutable device-global scratch buffer without synchronization. Cover shrink/reuse, growth, graph lifetime, and stream isolation. Existing thread: #7441 (comment)

3. Preserve or explicitly reject the NCCL-EP zero-copy contract. In deterministic mode the fused FC2 writes an unscaled result to output_buffer, but post-FC2 probability scaling returns a different allocation. The returned result is therefore no longer the registered symmetric output buffer. Unless that combination is implemented and validated end-to-end, please reject it explicitly before launching the fused op and add a regression test. I am not assuming how the external TE implementation handles the mismatched pointer. Existing thread: #7441 (comment)

CI follow-up: I inspected the exact-head failing/cancelled GDN bucket and its per-rank artifact: https://github.com/NVIDIA/Megatron-LM/actions/runs/38258534867/job/114861040388. Rank 0 reports TestFusedPreGatedDeltaRule::test_fused_and_unfused_pre_gated_delta_rule_match FAILED, while ranks 1-7 pass it and advance. Rank 0 subsequently times out in the teardown barrier; cancellation prevents the original assertion details from being printed. This is not enough evidence to classify the primary failure as a transient NCCL issue. Please reproduce the individual test and preserve the first assertion details before teardown; do not merely loosen tolerance or retry the whole matrix.

This is not the same test as the recent #7336 fixes: those adjusted deterministic mHC functional recipes (--gdn-kernel-backend torch), their GB200 goldens, and a separate fused-RoPE inference fixture. There is related GDN normalization work earlier in #7336, but these artifacts do not establish a shared root cause. Please reconcile the normalization/reference contracts and re-run this bucket on the updated exact head.

[HaochenYuan]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict follow-up review: PR #7441

Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 0
Overall risk: Low. The current head fixes all three findings from the previous bot review (review 5492800782 at d3b32b3). Nothing changed since then introduces a regression above the prior impact ceiling (moderate). Strict review passed with no significant issues found. LGTM.

Prior findings: fixed

1. TransformerConfig.deterministic_mode was ignored on Path A and the CP path. Fixed.
   • Path A: CSASparseAttnFunc now takes trailing deterministic / cuda_graph arguments and stores them on ctx (L1841-L1874). Its backward passes requested=ctx.deterministic (L1902-L1913). csa_sparse_attn forwards both flags (L2008-L2009). _forward_fused_no_indexer passes self.config.deterministic_mode (csa.py L2471-L2480).
   • CP path: FusedCSAIndexerSparseAttnFromTopkFunc now calls _deterministic_requested(deterministic) for the indexer dK (L3699) and passes requested= to the sparse backward (L3954-L3965).
   • Arity: I counted the positional arguments of the CP apply call (csa.py L3712-L3727). deterministic and cuda_graph land at forward slots 29 and 30. The 31-entry gradients tuple is still truncated by num_forward_inputs. The fused-indexer function's backward now returns 30 entries, matching its 30 forward inputs.
   • Test: the new test_csa_sparse_attn_honours_config_flag_without_torch_switch covers Path A with torch's global switch off. The CP fix was checked by reading the source; no unit test exercises it directly.
2. The workspace cache kept one buffer per distinct size. Fixed.
   • The cache now holds one buffer per device and reuses it whenever it is large enough, returning workspace[:required_bytes] (L1595-L1604).
   • A replaced buffer is kept in _DETERMINISTIC_SPARSE_BWD_RETIRED_WORKSPACES only when cuda_graph_impl != "none". cuda_graph_impl is an existing TransformerConfig field.
   • test_sparse_bwd_kwargs_select_deterministic_kernel_and_bound_workspace covers these cases: same size (same data_ptr), smaller request (same data_ptr), growth without graphs (nothing retired), and growth with graphs (one retired buffer).
3. Deterministic mode with NCCL-EP zero-copy. Fixed.
   • _fused_forward now raises when output_buffer is not None under deterministic_mode. The check runs before ops(...) is called (experts.py L855-L863).
   • Deterministic mode without zero-copy, and zero-copy without deterministic mode, take the same paths as before.
   • test_fused_forward_deterministic_mode_applies_probs_on_fc2_output now asserts that RuntimeError matches zero_copy.

Follow-up scope

• I compared the head with the previously reviewed tree (d3b32b3) in the areas changed since then:
  • cuda_graph / deterministic threading through every fused CSA call site (training, inference, THD and CP);
  • the new workspace retention policy;
  • the zero-copy guard;
  • the updated tests.
• On the inference call sites, the new deterministic flag only reaches ctx and is read only in backward, so forward behaviour is unchanged. The default path (deterministic_mode=False and torch's switch off) still gets {} from _deterministic_sparse_bwd_kwargs.
• The CI concern another reviewer raised (TestFusedPreGatedDeltaRule, the GDN bucket) involves code this PR does not change, so it is outside this review's scope.
• All conclusions come from reading the source; I ran no tests.

Review generated by claude

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

Review generated by claude

[FDecaYed]

Follow-up review of ee85072: the explicit deterministic flag now reaches Path A and FromTopk/CP, and the deterministic + NCCL-EP zero-copy combination is rejected before the fused op. I checked those changes independently with isolated CPU autograd/contract probes; this is not GPU validation.

The workspace item is not fully addressed yet: the cache still lacks stream ownership and can discard an allocation previously used by a graph when the next caller is non-graph. Details and the missing regression sequences are inline. Keeping changes requested for that remaining item.

The new exact-head CICD run 38385527145 is active, including the GDN fusion bucket. I have not retried it, and the previous-head GDN result is not a result for this SHA. Request NVSkills CI failed before creating any jobs; that startup failure provides no numerical-test evidence.

[FDecaYed]

Reviewed 6b8676e with the agreed support boundary: Megatron's existing ordered compute schedule, with mixed graph/non-graph deterministic CSA configurations on one device explicitly rejected. The guard executes before workspace reuse/growth, preserves graph-mode warmup/capture and retention, and has focused regression coverage. GB200 selection and the two-stream test input-readiness fence are fixed.

No remaining blocker from my review. Validation: 37 CPU diagnostics pass; the four new guard cases reject the previous implementation; 11 GB200 cases are selected by collection; mandatory formatter/lint checks pass. Actual CUDA/cuDNN/FlashMLA validation is pending in exact-head CICD 38689052085, so this approval does not assert new-head GPU CI has already passed or bypass required checks.

[HaochenYuan]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict follow-up review: PR #7441

Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 0
Overall risk: Low. I compared this head with the head of my last review (ee85072). Nothing that changed since then adds a regression above the prior impact ceiling (moderate). All earlier bot findings are still fixed. Strict review passed — no significant issues found. LGTM.

What changed since ee85072

Only two things changed:

• the workspace helper _get_deterministic_sparse_attention_workspace, with its ownership notes at L1521-L1556;
• the workspace tests in test_csa_deterministic_indexer.py.

The csa.py call sites, experts.py and determinism.py are the same as at ee85072.

Earlier findings: still fixed

1. Config flag on Path A and the CP path.
   • CSASparseAttnFunc.backward passes requested=ctx.deterministic (L1955-L1966).
   • FusedCSAIndexerSparseAttnFromTopkFunc uses _deterministic_requested(deterministic) for the indexer dK (L3752) and passes requested= to the sparse backward (L4007-L4018).
2. Bounded workspace. There is still one buffer per device. It is reused whenever it is large enough, and the helper returns an exact-size slice (L1635-L1656).
3. Deterministic mode with NCCL-EP zero-copy. The combination is still rejected before the fused op runs (experts.py L855-L863).

The other reviewer's workspace-ownership thread

I checked both of FDecaYed's sequences against the current helper.

• Two streams sharing the scratch.
  • A hand-out on a stream that does not own the buffer, outside capture, now calls stream.wait_stream(owner) and record_stream, then takes ownership (L1649-L1654).
  • Releasing a non-graph buffer after growth is therefore covered by the allocator's stream tracking.
• Graph warmup at 4096 bytes, capture, then a non-graph caller asking for 8192.
  • That call now raises at the graph_mode check (L1630-L1634) before the entry is replaced, so the graph's allocation is kept.
  • Growth in graph mode keeps the old buffer whenever graph_owned is set (L1636-L1638). That flag is set during capture or by the cuda_graph config (L1655), so retention no longer depends on who makes the request.
• Same cuda_graph value everywhere. Every training call site in csa.py takes it from the same self.config.cuda_graph_impl != "none" expression, for example csa.py L3725-L3726 and L3757-L3758. One model config therefore cannot trip the new check.
• Capture. During capture no handover is issued, and the stream that owns the buffer does not change. This matches the support boundary agreed in the thread and what test_deterministic_sparse_bwd_workspace_survives_graph_replay_after_growth asserts.
• Tests. The new mixed-mode test checks:
  • that both directions are rejected, for both reuse and growth;
  • that no reallocation happens and the entry's buffer, owner stream and graph flags are unchanged.

Scope and limitations

• Policy. This is a follow-up review with the default policy: inspect a subset of files, and publish only findings above the prior impact. The pinned pr-review skill does not override that.
• Comparison. I compared against the previously reviewed tree at /context/previous-review/.
• Style guide. style-guide.md does not exist in the base tree.
• What I did not check. All conclusions come from reading the source; I ran no tests. I did not check the CUDA paths (record_stream, capture, replay) or the real FlashMLA/cuDNN two-stream test by execution. Per the PR thread, GPU validation for this head is pending in CICD 38689052085.

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

Review generated by claude

[FDecaYed]

Re-reviewed the test-only update at 5a20d5f and investigated why the preceding merge-group left the queue. The queue removal was explicitly failed_checks at 13:55:47Z, before this push, not a consequence of Codecov or the head update.

The two functional-test blockers below are still present because their production code/config/goldens are unchanged in this head. I also found an input-dependent assertion in the new coverage test. My earlier workspace/graph-mode concern remains resolved under the agreed Megatron schedule boundary; I am not reopening it. The separate merge-queue notification token failure is infrastructure, not a requested code change. Please address these focused test/config/baseline items before the next landing attempt.

[HaochenYuan]

Re-reviewed the test-only update at 5a20d5f and investigated why the preceding merge-group left the queue. The queue removal was explicitly failed_checks at 13:55:47Z, before this push, not a consequence of Codecov or the head update.

The two functional-test blockers below are still present because their production code/config/goldens are unchanged in this head. I also found an input-dependent assertion in the new coverage test. My earlier workspace/graph-mode concern remains resolved under the agreed Megatron schedule boundary; I am not reopening it. The separate merge-queue notification token failure is infrastructure, not a requested code change. Please address these focused test/config/baseline items before the next landing attempt.

Thanks, all three addressed in a28c359:

• determinism.py: dropped the moe_router_fusion requirement instead of narrowing it. The scoped-CUDA-graph recipes set no --moe-aux-loss-coeff, so every _apply_*_aux_loss returns early and TE's fused_moe_aux_loss never runs there. Even with a non-zero coefficient, only the logged aux-loss value is affected: the forward reduces the scalar with a cross-block float atomicAdd, but FusedAuxLoss.backward reads Const_buf + tokens_per_expert and the autoscaler feeds a constant grad, so grad_probs is the closed form coeff·E/(topk·N²)·tokens_per_expert[e]. Measured over 200 repeats (65536×128): 15 distinct bit patterns for the scalar, grad_probs bitwise identical every time. The fused top-k/score kernels are one-warp-per-token with integer atomics only. Documented in docs/user-guide/deterministic-training.md; determinism.py is now identical to dev.

• experts.py: you were right that the golden drift was BF16 rounding from moving the prob scaling, and refreshing the golden would have been the wrong fix. TE only fuses the CuTe grouped MLP (the path whose dactivation epilogue accumulates dprob with atomics, where we observed the non-determinism) for MXFP8/NVFP4 recipes; the BF16 op-fuser path runs the separate ScaledSwiGLU op whose scale gradient is already reproducible. The post-FC2 placement is now gated on fp8 ∧ mxfp8 / fp4 ∧ nvfp4, so the BF16 op-fuser test keeps its existing golden. Unit test covers both branches.

• test_csa_fused_sparse_attention.py: seeded the FromTopk test and kept q_indexer/k_indexer positive so the ReLU is active for every selected key; 0/200 seeds fail (23/200 with the original random-sign inputs).

[FDecaYed]

Re-reviewed e48ce3c; the three follow-up findings are addressed.

• The blanket router-fusion restriction was too broad. The scoped-graph recipe has zero auxiliary-loss coefficient, and I checked the pinned TE backward plus its CUDA kernel: gradients use the saved coefficient and token counts, not the reduced aux-loss scalar. MCore's autoscaler supplies a value-independent gradient. Restoring dev behavior with the explicit logging caveat is appropriate.
• BF16 now follows the original probability-scaling path; the post-FC2 workaround is limited to the quantized recipes. The exact-method CPU branch test passes, and the old implementation fails its new BF16 check. No golden rewrite is needed for that BF16 regression.
• The FromTopk test now constructs positive Q/K instead of relying only on a seed. My independent exact-helper CPU probe has 0/200 failures with the fix versus 23/200 before it, with global determinism both off and on.

These CPU/source checks do not claim native TE/CUDA end-to-end validation. My three threads are resolved; the earlier workspace/graph-mode thread remains resolved under the accepted Megatron schedule boundary. No remaining blocking review finding. Exact-head CI 38886359330 is still running; required checks and normal merge protections still apply.