[buptzyb]

Summary

Enable attention_output_gate in absorbed MLA, including the default DSA path, with headwise and elementwise gates applied before the output projection.

Main counterpart: #8008.

Changes

• Add gate projections to the GPT and hybrid DSA specifications, including gate gradients during recomputation and deferred weight-gradient execution.
• Share FP32 sigmoid output gating across Attention, MLA, and KDA while preserving each caller's casting and compilation policies. Unify headwise layout handling while retaining numerical compatibility with scalar attention.
• Add independent numerical and distributed regressions for packed sequences, TP/SP/CP, selective recomputation, deferred gradients, disabled gating, and checkpoint shard metadata.
• Skip the optional cuDNN parity case when CUDA, SM90+ hardware, or its kernel dependencies are unavailable; supported environments still exercise the real fused path.
• Update existing DSA module-spec and disabled-gate test fixtures to cover the added gate projection while retaining context-parallel restoration and recomputation checks.
• Isolate compiled gate parity cases from shared compiler state and require both the actual and independent reference paths to compile, retaining bit-exact output and gradient checks.

Test plan

• The complete Transformer CI recipe bucket finished successfully on 8 H100 GPUs through the canonical runner with unchanged filters. Ranks 0–3 each passed 2,313 cases and skipped 218; ranks 4–7 each passed 2,312 and skipped 219. Every rank passed all 14 compiled gate, 50 output-gate, and 32 gated-DSA cases without skips. The experimental phase selected no tests under the existing markers and exited normally. The runner returned zero and completed coverage collection.
• After isolating compiler state, all 8 H100 ranks passed the complete 50-case output-gate file without skips or failures. Its 14 compiled parity cases require both production and independent reference paths to compile and retain zero-tolerance output and gradient checks.
• An unchanged-production GPU reproduction exhausted Dynamo's default cache and reproduced the CI case's exact mismatch count and maximum difference. Resetting the cache restored compiled execution on both sides and bit-exact outputs and gradients.
• Existing regression files rerun after synchronizing the fixtures: all 8 H100 ranks passed 33 DSA integration cases and 54 attention module-spec cases, with no skips or failures. Both previously failing cases explicitly passed.
• In a cached CI-derived image without FlashMLA, the cuDNN case skipped on all 8 H100 ranks with the dependency reason and a successful runner exit.
• After the availability update, the cuDNN parity case passed on all 8 H100 ranks, including the original numerical, gradient, and fused-kernel assertions.
• Full suite before the availability guard, using the canonical unit-test runner on 8 H100 GPUs: every rank passed 100 cases with no skips (50 output-gate, 18 existing headwise-attention, and 32 gated-DSA cases).
• Isolated gate outputs and VJPs match independent frozen primitives with zero tolerance. DSA parity includes the real fused cuDNN path.
• CHECK_ONLY=true BASE_REF=dev bash tools/autoformat.sh passed (Black, isort, Pylint, and Ruff).

The full GPU suite ran in the pinned CI-derived container on dev b4407209. Production source is unchanged on this PR's dev f11f03405 base. This is not validation of a fresh install from the current dev lockfile. The PR contains one signed commit.

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

[buptzyb]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review summary

Findings: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 0. No tests were run; everything below comes from reading the source.

Most impactful finding

[IMPORTANT Compatibility] The new test test_hy4_attention_shape_cudnn_gated_dsa_matches_native always builds a config with dsa_kernel_backend="cudnn". It does not first check that the GPU and kernel dependencies are present, which every other real-kernel cuDNN DSA test does.

• The config check at transformer_config.py L1992-L1993 raises a ValueError when flash_mla or cuDNN-frontend DSA is missing.
• The LTS image installs neither of them.
• The tests/unit_tests/transformer/**/*.py bucket runs in the lts environment.
• Result: the LTS unit-test lane always errors on this test instead of skipping it.

Details and the exact fix are in the inline comment. Proposed replacement for lines 471-472 of that test file (written to match the existing guards; not run):

def test_hy4_attention_shape_cudnn_gated_dsa_matches_native():
    if torch.cuda.get_device_capability()[0] < 9:
        pytest.skip("cuDNN fused DSA path requires SM90+")
    missing = _missing_cudnn_dsa_kernel_dependencies()
    if missing:
        pytest.skip(f"cuDNN fused DSA dependencies are unavailable: {', '.join(missing)}")
    _run_native_parity("elementwise", production_shape=True, kernel_backend="cudnn")

This also needs from megatron.core.utils import _missing_cudnn_dsa_kernel_dependencies added to the test's imports.

What was checked and found sound

Shared gate helper (attention_output_gate.py). Every caller keeps its previous precision and compile behaviour:

• Attention, scalar per-head gate: still eager, still casts the scale to the input dtype before multiplying.
• Attention _apply_output_gate: still compiled with jit_fuser, still multiplies by the FP32 scale and casts afterwards.
• MLA per-head gate: still compiled, casts before multiplying.
• MLA per-element gate: still eager, casts before multiplying.
• KDA: compiled, leaves the product in its natural dtype, and the caller applies the final .to(x_dtype).
• Shape checks match what the old inline MLA checks enforced.

Absorbed MLA / DSA gate (absorbed_mla.py):

• Gate input: the post-LayerNorm hidden_states. The spec sets fuse_input_layernorm: False, and gating combined with mla_down_proj_fusion is already rejected in the config.
• Projection: a column-parallel linear produces the TP-local head slice and gathers the sequence-parallel input, which matches the existing assert core_attn_out.shape[0] == hidden_states.shape[0]*TP.
• Placement: applied after the V up-projection and the recompute discard, before linear_proj.
• Deferred weight gradients: included in backward_dw, which the EP-overlap scheduler reaches through layer.self_attention.backward_dw.
• LayerNorm recompute: included in set_for_recompute_input_layernorm.

Checkpoints: the key linear_gate.weight and its row sharding are the same in the training AbsorbedMLASelfAttention spec and the inference MLASelfAttention DSA spec in hybrid_layer_specs.py.

Unused-identifier check (mandatory): the new linear_gate field, gate_projection_size, and the helper's granularity and cast_mode arguments all have real uses at runtime.

Overall risk

• Production code: low risk. The gate math is behaviour-preserving for existing callers, and the new absorbed-MLA path matches ordinary MLA's TP/SP/CP, recompute and deferred-gradient handling.
• Remaining risk: the CI-environment issue above.

Limitation: the style-guide.md that the skill references is not in the base tree, so style was not checked beyond the configured linters.

Review generated by claude

[buptzyb]

/ok to test b9c36e0

[yuzhongw-nvidia]

LGTM. Thanks.

[hxbai]

LGTM