[dev] feat(moe): Add KDA Low-Rank Projections and Extend MLA Support

[yuzhongw-nvidia]

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

What does this PR do?

• Add low-rank F-decay and output-gate projections to KDA while preserving the legacy fused projection path.
• Use fused LayerNorm column-parallel projections for MLA Q/KV up-projections in hybrid specs.
• Add per-layer NOPE support to MLA, including forward-path coverage.

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.

[yuzhongw-nvidia]

/claude review

[claude]

Light review. One likely bug (posted inline) plus two smaller notes; the rest looks solid — the legacy-fused path is preserved bit-for-bit and the FLOPs formula reduces to the old one when both ranks are None (I checked: hidden*(2qk+v+nh) + hidden*qk + hidden*v + hidden*v == hidden*(3qk+2v+nh) + hidden*v).

1. tp_comm_buffer_name="fc1" on the up-projections — see the inline comment on kda.py:148. "fc1" keeps TP-comm-overlap enabled with a UB buffer sized for hidden_size, but f_b_proj/g_b_proj consume kda_*_lora_rank-wide inputs. MLA uses non-recognized names for exactly this reason.

2. FP8 alignment is only asserted for the fused projection. common.py:235 checks in_proj_dim % fp8_align_size == 0, but in the low-rank layouts the F/gate GEMM widths (kda_f_lora_rank, kda_gate_lora_rank) are new FP8 GEMM dimensions that go unchecked. A user setting kda_f_lora_rank=100 with --fp8 gets no early error. Worth a validation in transformer_config.py next to the existing > 0 checks:

if self.fp8:
    fp8_align_size = get_fp8_align_size(self.fp8_recipe)
    for field_name in ("kda_f_lora_rank", "kda_gate_lora_rank"):
        rank = getattr(self, field_name)
        if rank is not None and rank % fp8_align_size != 0:
            raise ValueError(
                f"KDA requires {field_name} to be a multiple of {fp8_align_size} "
                f"under FP8, got {rank}."
            )

3. No _get_in_proj_dim override test for the checkpoint-compat claim. test_kda_forward_backward asserts in_proj_split_names/in_proj_dim per layout, which covers the shape. But the stated goal — "existing KDA checkpoints load without migration" — isn't directly exercised: nothing saves a legacy-fused checkpoint and loads it back. test_parallel_kda_correctness does save/load, but each parametrization uses the same rank config on both sides, so a layout regression would pass. Optional, but a single save-legacy/load-legacy roundtrip would pin the compatibility guarantee that motivates the use_legacy_fused_projections branch.

Nit: transformer_config.py:1761 has a stray f-string split ("KDA requires kda_gate_lora_rank > 0, " f"got {...}") — the adjacent kda_f_lora_rank message puts the whole thing in one f-string. Cosmetic only.

[yuzhongw-nvidia]

/claude review

[yuzhongw-nvidia]

/ok to test 6f9e6bd

[claude]

Light review. The KDA low-rank work is well-structured — the _get_in_proj_dim hook, the legacy-fused checkpoint preservation, and the parametrized test coverage across all four rank combinations are all solid. Three things flagged inline, plus notes below.

Blocking

Circular import (transformer_config.py:14). fp8_utils already imports transformer_config at module level, so the new top-level from megatron.core.fp8_utils import get_fp8_align_size closes the cycle. Installation Test and both build-and-test-wheels jobs are currently red, which is consistent with this. Deferring the import into __post_init__ matches how the rest of the tree handles it.

Worth a look

tp_comm_buffer_name collision (kda.py:118, kda.py:182) — f_proj and g_proj reuse "fc1", already claimed by in_proj.

Unconditional LayerNorm in hybrid MLA specs (hybrid_layer_specs.py) — the TEColumnParallelLinear → TELayerNormColumnParallelLinear swap isn't gated on qk_layernorm, and q_layernorm/kv_layernorm stay IdentityOp, so this adds a norm rather than fusing an existing one. Details inline.

Test coverage

no_rope for MLA has a forward-shape test, but it doesn't assert the no-RoPE path is numerically distinct. The use_rope=False branch at multi_latent_attention.py:1002-1021 concatenates q_pos_emb / k_pos_emb unrotated, so a cheap regression guard would be to build a second layer with no_rope_freq=[0, 0], copy weights across, and assert the outputs differ. As written, the test would still pass if the branch accidentally applied RoPE.

Also, the low-rank projections are only exercised by unit tests. hybrid_mr_mcore_te_tp1_pp1_ep8_cp2_kda_gated_mla_1N8G still runs with the legacy fused layout, so the new path gets no functional/ckpt-resume coverage. Given the two distinct checkpoint layouts this PR introduces, a --kda-f-lora-rank variant of that recipe would be valuable.

Minor

Two unrelated parenthesis-stripping edits in training.py (lines ~4945, ~5243) are outside the PR's stated scope — docs/developer/contribute.md asks contributors not to touch formatting in code unrelated to the change.

[yuzhongw-nvidia]

/ok to test 66c6205

[claude]

Light review — the KDA low-rank plumbing and MLA NoPE changes look sound overall, and unit coverage for the new rank combinations is thorough. No new direct parallel_state.get_*_group() usage (KDA uses self.pg_collection.tp, MLA uses self.pg_collection.cp). Two things:

1. Contradictory NoPE handling for cached MLA latents — flagged inline at multi_latent_attention.py:935. The new assert self.use_rope or not self.cache_mla_latents makes the if self.use_rope: guard inside qkv_up_proj_and_rope_apply_for_cached_latent_kv dead code. Worth resolving one way or the other.

2. Spec wiring for the 6 new submodules is untested. get_gated_delta_net_module_spec now wires f_proj/f_a_proj/f_b_proj/g_proj/g_a_proj/g_b_proj, but tests/unit_tests/models/test_experimental_attention_variant_module_specs.py::test_kda_uses_direct_projection_submodules still only asserts in_proj and beta_proj. A silent regression there (e.g. f_a_proj accidentally wired to column_parallel_linear instead of linear) would only surface as a TP shape error at runtime. Suggested additional assertions in that test: f_proj, f_b_proj, g_proj, g_b_proj should equal _FakeColumnParallelLinear, and f_a_proj, g_a_proj should equal _FakeLinear.

Relatedly, backend.linear() is used here but is not declared on the BackendSpecProvider Protocol in megatron/core/models/backends.py and is absent from LocalSpecProvider. That is safe today only because _get_backend_spec_provider asserts transformer_impl == "transformer_engine"; adding it to the Protocol would make the contract explicit.

[yuzhongw-nvidia]

/claude review

[claude]

Light review. The three features are cleanly separated and the parametrized test matrices are good — test_parallel_kda_correctness covering all four low-rank combinations across TP/SP/CP with a dist-ckpt round trip is exactly the right coverage for the new projections, and the FLOPs formula is legacy-equivalent (the refactored projection_flops expands to the old hidden_size * (in_proj_dim + num_qk_heads) + hidden_size * v_dim when both ranks are None).

Two comments left inline:

1. common.py — hard-coding act_fn = silu in the shared base also changes GatedDeltaNet, and only KDA gets an assertion for it.
2. kda.py — the new FP8 rank-alignment guard (and the rank <= 0 guards in kda_layer_flops) are untested.

Both are minor; no blocking issues found.

One thing I checked and it's fine: f_a_proj/g_a_proj use parallel_mode="duplicated" without passing tp_group, which satisfies TELinear's assert tp_group is None and routes their dist-ckpt through the replicated TELinear.sharded_state_dict path.

[yuzhongw-nvidia]

/claude review

[claude]

Light review — the production changes look correct to me. te_checkpoint usage in gdn.py/kda.py matches the established FP8/FP4 selective-recompute pattern in transformer_layer.py/moe_layer.py (signature and active_pg_collection.tp are right), the legacy vs. low-rank in_proj split names/sections stay consistent with common.py's sharded_state_dict width assertion, the FLOPs plumbing into kda_layer_flops/hybrid_flops is correct, and no new global parallel_state.get_*_group() reads are introduced.

Two test-coverage nits:

1. test_gdn_recompute_mxfp8_backward_has_main_grad asserts nothing (tests/unit_tests/test_muon_decouple_fp8_param_gather.py)

The name and docstring promise a main_grad check, but the body ends at self._run_steps(args, model, optimizer, 1). Since _run_steps only collects grads where p.main_grad is not None, a regression that drops main_grad on recomputed GDN/KDA params would pass silently. Something like:

    self._run_steps(args, model, optimizer, 1)

    gdn_params = [
        (name, param)
        for module in model
        for name, param in module.named_parameters()
        if "mixer" in name and param.requires_grad
    ]
    assert gdn_params, "no GDN/KDA parameters found in the model"
    missing = [name for name, param in gdn_params if getattr(param, "main_grad", None) is None]
    assert not missing, f"GDN/KDA parameters missing main_grad after backward: {missing}"

(adjust the name filter to whatever actually matches the GDN/KDA submodule for this hybrid pattern).

2. The six new KDA submodule wirings have no spec-level assertion

TestGetGatedDeltaNetModuleSpec::test_kda_uses_direct_projection_submodules in tests/unit_tests/models/test_experimental_attention_variant_module_specs.py still only checks in_proj/beta_proj, so a regression in the new f_proj/f_a_proj/f_b_proj/g_proj/g_a_proj/g_b_proj wiring (or the linear() vs column_parallel_linear() choice for the duplicated *_a_proj path) would not be caught there. Worth extending that assertion, and mirroring it for hybrid_layer_specs.py.

[yuzhongw-nvidia]

/claude review

[yuzhongw-nvidia]

/claude review