[Dev] Fuse standard RoPE packing for MLA and DSA

[buptzyb]

Summary

• enable the fused MLA RoPE path for standard RoPE while preserving Yarn, unfused, and inference behavior
• fuse absorbed-MLA non-positional packing, RoPE, and contiguous output into one Triton kernel with a fused backward
• extend the shared in-place MLA RoPE primitive to rotate leading channels, allowing the DSA indexer to operate directly on its packed [PE, NOPE] projection output

Performance

Formal 8x B200 GLM-5.2 proxy A/B results, using 8 warmup and 20 measured iterations:

Workload	Step time	Peak device memory
THD 4K, dynamic effective CP1	-2.11%	<0.2% change
THD 32K, CP8	-2.80%	<0.2% change

In the targeted absorbed-MLA packing region, the long-sequence path is reduced from 85 kernels to 2 kernels, and GPU kernel time falls from 0.848 ms to 0.099 ms per attention call. Whole self-attention GPU span falls by 10.6% for THD 4K and 11.5% for THD 32K.

The DSA indexer in-place integration is included for a cleaner packed-tensor path and shared primitive ownership. It is correctness- and trace-validated, but no additional end-to-end speedup is claimed for that follow-up beyond the MLA A/B results above.

Testing

• standard and fused MLA RoPE forward/backward parity across SBHD, THD, CP1/CP>1, unequal packed segments, padding, BF16, and FP32
• MLASelfAttention and AbsorbedMLA output, hidden-gradient, and parameter-gradient parity
• DSA indexer fused/unfused parity for SBHD and THD; 8 primitive and 4 module-level focused cases passed on GB200
• GLM-5.2 MXFP8/cuDNN DSA proxy training and single-iteration Nsight Systems 2026.4.1 capture on 16x GB200
• Megatron autoformat checks: Black, isort, pylint, and ruff

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

[buptzyb]

/ok to test e354f1f