[HaochenYuan]

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

What does this PR do?

Under CUDA graphs the fused compact DSA indexer (--dsa-kernel-backend cudnn, --dsa-indexer-precision mxfp8) keeps a persistent workspace per static geometry (MXFP8 q/k quantization destinations, packed scales, candidate offsets) that must be prepared outside capture. Every CSA layer owned its own copy: about 7 GiB per rank on DSv4 at 16K tokens per rank.

• All CSA layers built from one TransformerConfig now share one workspace per geometry and balanced-indexer slot (_CompactIndexerWorkspaceStore, keyed on the config). The layers of a graph run serially on one stream and nothing in the workspace outlives the indexer dispatch that fills it, so sharing only removes duplicates.
• New config field dsa_compact_indexer_workspace_sharing (default True; --no-dsa-compact-indexer-workspace-sharing restores per-layer workspaces), golden-config entry, doc paragraph in cuda_graph.md.
• Tests: kernel-level workspace tests (test_csa_compact_indexer_workspace_sharing.py) and a contract test in the CSA test class that a workspace prepared by one layer serves the capture of another.

Validation: GB200, 4 ranks: CSA test files 118 passed. DSv4-flash 1M on 128 GPUs with chunk-wise graphs: last-stage NVML memory 173.8 → 149.4 GiB together with linear CE (#5258 validation).

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

[buptzyb]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review: share the fused compact DSA indexer workspace across CSA layers

Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 2 (both posted inline)

Risk: low. Sharing one workspace across layers looks safe. I traced the workspace from the point it is written to every point it is read:

• Quantization buffers: the MXFP8 quantization destinations and packed scales are overwritten on every dispatch.
• Per-dispatch refresh: cand_batch_offsets and the MXFP8 scale offsets (cu_seqlens_*_scale_padded) are rewritten in place on every dispatch (fused_sparse_attention.py L1875-L1877, L1925-L1930).
• Scratch and outputs: candidate scratch and the compact outputs are allocated per dispatch, and a check confirms the kernel wrote into those per-dispatch buffers (L540-L566, L1974-L1990).
• Backward: the sparse-attention backward saves q, kv, topk_idxs, out and lse, never the workspace.
• Ordering: with the local CUDA graph implementation, each graph's warm-up runs right before its own capture (cuda_graphs.py L930-L980). Changing who owns the "active" workspace therefore does not change which workspace a capture sees.

So a later layer overwriting the shared buffers is the same reuse that eager mode already does.

The new config field is read in _compact_indexer_workspace_store, so it is not dead code. Its --no-/--disable- CLI flags are generated automatically by ArgumentGroupFactory (argument_utils.py L199-L211, arguments.py L2861). The golden-config entry is updated.

Suggestions (inline):

1. BSHD capture path: it still accepts only the active workspace, while the THD path now scans every warmed-up workspace. The full-recompute geometry alternation described in the new THD comment applies to BSHD unchanged. This already existed before the PR and is not caused by sharing. Mirroring the THD scan is a short change, and the patch is in the inline comment.
2. Stale comments in _get_thd_compact_indexer_workspace: two unchanged comments now contradict the code. The docstring says capture reuses only the active workspace. Another comment says each attention module owns its storage independently, which is no longer true with the default sharing.

Limitations:

• The skill's style-guide.md is not in the base tree (/context/base/style-guide.md and /context/base/skills/style-guide.md are both missing), so no style-guide audit was done.
• I could not search the whole repo for code that assigns to the now read-only _*_compact_indexer_workspaces attributes. I spot-checked cuda_graphs.py, cp_balanced_indexer.py and the DSv4 CP CUDA graph tests and found no such assignments.
• None of the suggested code has been run.

Review generated by claude

[svcnvidia-nemo-ci]

[SUGGESTION Consistency] The BSHD capture path still accepts only the active workspace, while THD now scans every warmed-up workspace.

Failure. The new THD comment (L2359-L2362) says that under full activation recompute, the no-grad forward and its recompute alternate between two geometries. The BSHD path alternates the same way:

• No-grad forward: forward goes to _forward_fused_indexer_inference, which asks for a workspace with return_softmax=False (L2607).
• Recompute with grad: it goes to _forward_fused_indexer_training, which uses return_softmax=indexer_loss_coeff > 0 and sparse_loss (L2674-L2680).

BSHDCompactIndexerWorkspace.matches_shape compares return_softmax. However, the capture branch (L2284-L2292) only tries self._active_bshd_compact_indexer_workspace.

Trigger: warm up with return_softmax=False, then warm up with True, then capture the False geometry. The capture raises ValueError: BSHD compact CUDA graph capture requires an eagerly prepared compact workspace..., even though a matching workspace is already in _bshd_compact_indexer_workspaces. With sharing on (the default), the active slot now holds the last warm-up done by any CSA layer on the device. The base code already checked only the active workspace, so this is not a regression from sharing. It is the same failure the PR fixes for THD, left in place for BSHD.

Correction. Replace the capturing branch of _get_bshd_compact_indexer_workspace with the same candidate scan the THD path uses:

        capturing = torch.cuda.is_current_stream_capturing()
        workspace = self._active_bshd_compact_indexer_workspace
        if capturing:
            # As in the THD path: the last warm-up may have prepared the other recompute
            # geometry. Matching is shape-only (no host synchronization).
            candidates = [workspace] if workspace is not None else []
            candidates += [w for w in self._bshd_compact_indexer_workspaces if w is not workspace]
            for candidate in candidates:
                if candidate.matches_shape(**match_kwargs):
                    self._active_bshd_compact_indexer_workspace = candidate
                    return candidate
            raise ValueError(
                "BSHD compact CUDA graph capture requires an eagerly prepared compact "
                "workspace for the active static geometry. Run eager warmup before capture."
            )

Invariant: during capture, return any eagerly prepared workspace whose static geometry matches. Never prepare one during capture, because preparing needs a host sync.

Preservation. The existing test test_bshd_compact_workspace_prepared_in_warmup_and_reused_in_capture (test_attention_variant_csa.py L900-L953) should still pass. Warm-up sees an empty list and prepares once; capture has exactly one candidate. So workspace.matches_shape.assert_called_once() and _bshd_compact_indexer_workspaces == [workspace] still hold. A capture with no matching warm-up still raises ValueError.

Acceptance (proposed test, not run).

• Setup: set dsa_compact_indexer_workspace_sharing=False and patch bshd_compact_indexer_available to return True.
• Patch prepare_bshd_compact_indexer_workspace with side_effect=[ws_f, ws_t]. These are MagicMocks whose matches_shape returns kwargs["return_softmax"] is False and is True respectively.
• Patch is_current_stream_capturing with side_effect=[False, False, True].
• Steps: call with return_softmax=False, then True, then capture with False.
• Expected: the third call returns ws_f and prepare was called twice. Today the third call raises ValueError.

[svcnvidia-nemo-ci]

[SUGGESTION Comments] Two unchanged comments in this function now contradict the new behavior:

• Docstring (L2331-L2334): it says "Capture reuses the active workspace". Capture now tries the active workspace first and then every warmed-up workspace in the slot (L2363-L2385).
• Comment at L2344-L2346: it says "...just as each attention module owns storage independently of the other layers". With dsa_compact_indexer_workspace_sharing=True (the new default), modules no longer own their storage independently. Only the separation between head and tail slots remains.

A wrong comment here will mislead the next person debugging a capture failure. Suggested replacement for both:

        """Return persistent compact storage for a warmed-up THD graph geometry.

        Workspace sizing synchronizes with the host, so it is performed only
        during eager CUDA-graph warmup. Capture reuses the active workspace or any
        other warmed-up workspace of the same slot whose static geometry matches;
        when none matches (warmup skipped or static shapes changed), capture must
        fail clearly. Unsupported devices/frontends retain the non-compact path.
        """
        ...
        # Head and tail have different causal offsets and may have different
        # geometry. Keep their quantization/candidate metadata in separate slots;
        # each slot is shared by every CSA layer bound to the same workspace store.

Acceptance (static check): the docstring and comment should match the capture branch at L2358-L2389 and the binding logic in _compact_indexer_workspace_store. Behavior does not change.