[pushkar-sarvam]

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

What does this PR do?

Keeps the MLA linear_q_proj that the layer spec asks for when q_lora_rank is None and qk_layernorm is on, instead of always replacing it with a fused norm+linear.

Since #4452, QKNormConfigResolver._resolve_mla_qk_layernorm swaps linear_q_proj for backend.column_parallel_layer_norm_linear() whenever there is no Q-LoRA, although the GPT layer specs pass backend.column_parallel_linear(). That adds an RMSNorm on the Q path (linear_q_proj.layer_norm_weight) that MLA models without Q-LoRA do not have; in DeepSeek-V2-Lite and Sarvam 105B the query is q_proj(hidden_states). Up to core_r0.18.0 the projection was plain. With the fused default, Megatron-Bridge cannot import such Hugging Face checkpoints (the norm has no source weight), and models trained from scratch get a Q norm that their inference implementations lack.

A spec that selects the fused norm+linear for linear_q_proj still gets the Q norm, as the message in _raise_unused_q_norm already suggests. The KV path is unchanged.

Testing: tests/unit_tests/transformer/test_mla_qk_norm_config.py fails on main and passes with this change. tests/unit_tests/transformer/test_multi_latent_attention.py gives the same results with and without this change on one GPU. A small Sarvam MLA model imported through Megatron-Bridge with a plain linear_q_proj matches the Hugging Face logits (relative L2 difference 5e-4 in FP32).

Issue tracking

Linked issue: related to NVIDIA-NeMo/Megatron-Bridge#5261. Megatron-Bridge works around this for DeepSeek with an MLASelfAttention subclass (NVIDIA-NeMo/Megatron-Bridge#5262); other MLA models without Q-LoRA, such as Moonlight and Sarvam 105B, still get the extra Q norm.

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

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

[pushkar-sarvam]

Hi @guihong-nv @janEbert, could one of you /ok to test and take a look when you get a chance? Thanks!

[pushkar-sarvam]

hey @guihong-nv @janEbert, just bringing this to the top - can you have a look at the PR?

[janEbert]

I think this is quite difficult. On the one hand, it's correct that Q also gets a norm when a setting called qk_layernorm is set by the user. On the other, it changes the existing default and is an extension that's not specified in the paper.

I would argue that it's more correct to give a Q norm when it's asked for (qk_layernorm=True) rather than just ignoring the setting like it used to (the behavior this PR would revert back to). Subjectively, I'd say it's correct that we add norm support for q_proj and that Megatron-Bridge works around this for HF support.

Unless there's a mathematical reason that having q_proj with a norm doesn't make sense?

[janEbert]

Feel free to let me know if you disagree. I'll close this for now so we don't accidentally merge.