Explicit safe MLA RoPE primitives and consumers

[FDecaYed]

Explicit safe MLA RoPE primitives and consumers

Summary

Expose explicit in-place and out-of-place MLA RoPE APIs, including inverse
rotation, KV-split output, and packed-position handling safe for graph capture
and fused backward.

Scope and non-goals

• Support SBHD and THD position mapping and packed maximum-sequence handling.
• Use out-of-place inverse RoPE where saved tensors must remain unchanged.
• Preserve legacy-call safety without depending on Fix fused MLA YaRN RoPE query kernel corrupting upstream tensor in-place #5412.
• Do not add CSA/DSA algorithms, DSv4 orchestration, or model construction.

Provenance

This local recut is reconstructed from the frozen main baseline
bb5647a9bdd0, not a replay. It references #5795 and the complete API from
#5944, incorporating fixes represented by #5018 and #5526. #5412 is an
independent compatibility gate, not a dependency. Credit goes to the original
contributors: @hxbai in #5018/#5795, @kunlunl in #5526, and
@ShauryaaSharma in #5412.

Dependencies

None beyond the frozen baseline used for this local recut. This is a reusable
prerequisite outside the nine-PR DSv4-specific series.

Tests

• In-place/out-of-place forward and inverse parity.
• SBHD/THD mapping, KV-split output, and packed maximum-sequence behavior.
• CUDA-graph-safe metadata handling and fused-backward saved-tensor safety.

Publication

Publish in the first prerequisite round on then-current main. This is not a
stacked-review PR.

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

[FDecaYed]

/ok to test 047a52c

[FDecaYed]

/ok to test fbbc501

[guihong-nv]

Strict review of 7d2cac16.

Findings: CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 0.

The asymmetric packed-frequency fix is now applied consistently at the three affected call sites and has focused regression coverage. The remaining blocker is the packed fused dispatch: the newly added inverse/output-layout semantics are silently dropped for THD inputs. Please route unsupported modes through the unfused implementation (or extend the fused backend) and add packed-fusion coverage.

I did not find another blocking issue in the current diff. GitHub reports codecov/patch as successful, but no Actions workflow run is associated with this head.

[guihong-nv]

[IMPORTANT Correctness] This new argument is not honored for packed inputs when apply_rope_fusion=True: the THD branch at lines 372–381 returns directly from TE with only interleaved, so inverse=True still applies forward RoPE; mla_output_remove_interleaving is likewise dropped. The SBHD branch already falls back when TE cannot honor inverse; please add an equivalent THD guard/fallback (including the MLA interleaving/mscale constraints as applicable), or pass supported flags through. A regression with cu_seqlens, apply_rope_fusion=True, and inverse=True would catch this.

[FDecaYed]

/ok to test 8a99b73

[asolergi-nv]

All this RoPE-Yarn-Fused-Unfused-MLA-THD-BSHD is getting more and more complex over time...