Fix fused MLA YaRN RoPE query kernel corrupting upstream tensor in-place

[ShauryaaSharma]

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

What does this PR do ?

Stop the fused MLA RoPE query kernel (rotary_fwd_q_kernel) from modifying its input in-place, which silently corrupts the activation TransformerEngine saved for the weight gradient and produces incorrect dW.

Problem

When apply_rope_fusion=True for MLA (Multi-Latent Attention), rotary_fwd_q_kernel modifies the query tensor Q in-place. The buffer passed as Q shares underlying storage with the output of linear_q_up_proj, because the caller does q = linear_q_up_proj(q_compressed) then q = q.view(...) — a zero-copy reshape.

TransformerEngine saves that linear-layer output for weight-gradient computation (dW = input^T @ grad_output). The in-place Triton write silently overwrites those saved bytes before backward runs, producing incorrect dW and therefore corrupted model weights after the optimizer step.

The result is a silent correctness bug — wrong gradients that degrade training rather than failing loudly. (The iteration-2 NaN reported in #5317 has a separate root cause: an out-of-bounds cos/sin read when rotary_percent < 1 shrinks the cache below qk_pos_emb_head_dim. This PR therefore relates to #5317 rather than closing it.)

Solution

Clone q inside ApplyMLARotaryEmbQ.forward() before the Triton kernel call. The clone is contiguous, satisfies stride(-1) == 1, and has identical shape and strides — no callers or public interfaces change.

The KV path (rotary_fwd_kv_kernel) already allocates fresh o_key and o_value output buffers and is not affected.

Changes

• megatron/core/fusions/fused_mla_yarn_rope_apply.py: add q = q.clone() before kernel dispatch in ApplyMLARotaryEmbQ.forward(); update docstrings to reflect that the input tensor is no longer modified.
• tests/unit_tests/fusions/test_mla_yarn_rope_apply.py: add TestFusedApplyMLARopeNoInputMutation — verifies the input tensor is byte-for-byte unchanged after fused_apply_mla_rope_for_q returns, parameterized over sbhd and thd formats.

Issue tracking

Linked issue: Related to #5317

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• 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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[Phlip79]

/ok to test 59acdf2

[ShauryaaSharma]

Hey Team, following up on this PR, which has been ready for review for about a week. All CI checks are passing (70/70), /ok to test was approved by @Phlip79, and the Final Review label is set. The only remaining step is a merge from NVIDIA/core-adlr or NVIDIA/core-nemo. Would appreciate it if someone could take a look when capacity allows. Thank you.

[guihong-nv]

/ok to test 5641ee3

[ShauryaaSharma]

Hey Team, following up on this PR, which has been ready for review for about two week. All CI checks are passing (70/70), /ok to test was approved by @Phlip79 and @guihong-nv , and the Final Review label is set. The only remaining step is a merge from code owners approval. Would appreciate it if someone could take a look when capacity allows. Thank you.

[guihong-nv]

/ok to test b1f04ac

[guihong-nv]

@ShauryaaSharma Thanks for the fix and the regression test. I took a closer look at the claimed TE aliasing path, and I think we need to validate the premise before making this out-of-place.

In the in-tree MLA caller, q is the fresh output of linear_q_up_proj/linear_q_proj, followed by a view. The pinned TE implementation appears to save the linear input (or normalized GEMM input for LayerNormLinear) for wgrad—not the linear output q; delayed wgrad also stores input plus grad-output. So I don’t yet see how the in-place RoPE write can corrupt that projection’s dW.

This makes q.clone() a full allocation/copy on the fused hot path and changes the documented in-place behavior. The added test verifies the new non-mutating contract on a synthetic alias, but it does not exercise a TE projection or compare its weight.grad.

Could you add a minimal end-to-end reproducer that compares the Q projection wgrad with the fused and reference paths (including the relevant delayed-wgrad/TE configuration), or point to the specific consumer that retains the unrotated Q buffer? Otherwise, I think we should avoid the clone or retarget the fix.

[ShauryaaSharma]

@guihong-nv Thanks for pushing on this, you're right, and I want to correct the record rather than defend the original claim.

I pulled transformer_engine==2.17.0 and read transformer_engine/pytorch/module/linear.py directly. The wgrad GEMM is:

def wgrad_gemm(x, dy, ...):
    """Perform wgrad GEMM: dw = dy^T * x"""
    dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs)

called as wgrad_gemm(inputmat_total, grad_output), where inputmat is saved in forward() as inputmat = inp, i.e. the linear's input (q_compressed in the MLA caller), not its output. The delayed-wgrad path (linear_q_up_proj.backward_dw(), which this repo uses) is the same: wgrad_store.put([inputmat_total, grad_output], wgrad_gemm), again input + grad_output, never the output tensor.

q (the output of linear_q_up_proj, and what the fused RoPE kernel mutates in-place) is never referenced anywhere in _Linear.backward. It also isn't the same tensor as inputmat , different GEMM, different shape, so there's no aliasing path into dW here. The premise in the PR description is wrong; I don't have a reproducer showing wgrad divergence, and based on this I don't expect one to show anything.

Given that, I don't think the q.clone() is justified as-is, it's a real allocation/copy on the fused hot path with no demonstrated correctness benefit. I'd rather not carry it on a false premise. Let me either:

1. Drop the clone and close/narrow this PR, since the actual iteration-2 NaN root cause (cos/sin OOB read when rotary_percent < 1) is being handled separately in Add cos/sin width guard to fused MLA RoPE kernels #5497, or
2. If you know of a real consumer that retains a reference to the unrotated q (activation recompute, CUDA-graph buffer reuse, an MTP head, etc.), point me at it and I'll retarget the fix there instead of TE's wgrad path.

Appreciate you catching this before it merged on a wrong justification.

[guihong-nv]

Thanks for validating this and correcting the record. I agree with the conclusion: I do not know of an in-tree consumer that retains the unrotated projection output, so I’d take option 1—drop the clone and its associated API/test changes, then close or narrow this PR. The cos/sin-width guard in #5497 is the appropriate follow-up for the reported NaN path.