Add cos/sin width guard to fused MLA RoPE kernels

[ShauryaaSharma]

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

What does this PR do ?

Add a width guard to the fused MLA RoPE kernels so an undersized cos/sin cache fails with a clear error instead of silently reading out of bounds.

The fused kernels (ApplyMLARotaryEmbQ / ApplyMLARotaryEmbKV) read emb_dim cos/sin values per token and assume the cache is emb_dim wide. When the rotary cache is built with rotary_percent < 1, the base RotaryEmbedding shrinks it to int(emb_dim * rotary_percent) , narrower than emb_dim , and the kernel reads past the buffer, picks up uninitialized GPU memory, and produces NaN gradients. This is reachable on main: MultiLatentAttention forwards rotary_percent into the base RotaryEmbedding for rope_type="rope" (multi_latent_attention.py:184), and the fused branch (:870) calls the kernel regardless of rope type with no width check today (fused_mla_yarn_rope_apply.py:251-255). The guard asserts cos.shape[-1] == emb_dim (and sin) in both forward paths, turning a silent ~1e14 out-of-bounds read into an immediate, readable failure.

Issue tracking

Linked issue: Related to #5317

This guard makes the failure mode in #5317 loud instead of silent; the root-cause fix for that issue is rotary_percent = 1.0 in the Megatron-Bridge DSv4 config mapping. Independent of #5412, which fixes a separate in-place aliasing bug.

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.

[ShauryaaSharma]

Hey team, its been 2 weeks since the PR is open, Kindly approve the rest of the workflow so that if any errors occur and code change is required, i can proceed with the same. Also, the code has been reviews and approved by: @hxbai
Thank you

[guihong-nv]

/ok to test 3ede266

[guihong-nv]

One blocking safety issue: the new width guard must remain active in optimized Python; details inline.

[ShauryaaSharma]

Fixed this now. In 430e64fb7: both the Q and KV checks are now a plain if/raise ValueError instead of assert, so they run regardless of optimization mode.

if cos.shape[-1] != emb_dim or sin.shape[-1] != emb_dim:
    raise ValueError(
        f"cos/sin last dim must equal emb_dim={emb_dim} "
        f"(got cos={cos.shape[-1]}, sin={sin.shape[-1]}); a narrower rotary cache "
        f"(e.g. rotary_percent < 1) makes the fused MLA kernel read past the buffer."
    )

Also updated the two unit tests to expect ValueError instead of AssertionError so they still pass.

Before pushing I double-checked the behavior directly (mismatched width raises, matching width passes through fine, and the asymmetric case where only one of cos/sin is narrow still raises) under both normal and -O execution — confirmed the old assert version really did let a narrow cache through silently under -O, and the new version catches it in both modes.