[port] Fuse standard RoPE packing for MLA and DSA (upstream #6343)

[HJSang]

Update 2026-09-23 — ready for review. Rebased onto current miles-main (96fe16ecb); it applied cleanly. The commit content is unchanged. This PR is the base of the #99 → #100 → #101 stack.

Stack and merge order: #97 (independent) · #98 → #99 → #100 → #101. Each PR is based on miles-main and contains its parent PR's commits. Review only the top commit of each PR (the Commits tab lists them). After a parent is squash-merged, rebase the child onto miles-main and drop the parent's commit. The order follows upstream's own merge order (NVIDIA#6058 on 2026-09-14 before NVIDIA#7169 on 2026-09-17), so #101 now applies upstream's final NVIDIA#7169 as-is.

Test plan: tracked in the comment on #99.

Test results (2026-09-23) — ✅ no regressions

Test	Base	This PR
MLA / RoPE / DSA set: test_mla_yarn_rope_apply, test_absorbed_mla, test_attention_variant_dsa, test_multi_latent_attention (4 ranks)	175 passed / 93 failed	212 passed / 106 failed

None of the 106 failures is caused by this PR. The 15 that are not on the base are all in tests this PR adds, and they fail for reasons already present on miles-main:

• 14× ValueError: spec sets linear_q_up_proj=TEColumnParallelLinear and q_layernorm=…. This comes from the QK-norm resolver ported in [port] Megatron-Core APIs needed by Megatron-Bridge main (#4452 MLA resolver, gdn alias, #5865 FSDP names) #102, and the same error already breaks 80 existing test_absorbed_mla cases on the base. It should be tracked separately against [port] Megatron-Core APIs needed by Megatron-Bridge main (#4452 MLA resolver, gdn alias, #5865 FSDP names) #102.
• 1× TE "No dot product attention backend". This is the new [rope-True] case of a CP-THD precision test whose 3 sibling cases fail the same way on the base.

This PR's kernels are also exercised by the #99 / #101 E2E runs (RoPE fusion is off in Miles recipes, so there is no numerics change).

Environment: 2× new 4× B200 devboxes (b200-verda-k8s), radixark/miles:dev rebuilt after miles#3296 (cuDNN Frontend 1.28.0 with the compact wrapper, TE 2.17, torch 2.13), Miles @ a6b72d31f, base miles-main@96fe16ecb. Unit tests ran under torchrun --nproc_per_node=4, and every failure was compared against the base run the same way. The Miles E2E recipe is DeepSeek-V4-Flash-FP8-4layer, GSM8K, deterministic mode, 3 rollouts, --dsv4-impl megatron --dsa-kernel-backend cudnn, with a per-token old-vs-train log-prob dump (~3.14M response tokens per arm). Rollouts were identical across arms. Full details: #99 (comment)

---

Ports NVIDIA#6343 (53f4976901, merged to upstream dev 2026-08-26, author Robin Zhang) onto miles-main. Cherry-picked with -x; auto-merged with no conflicts.

What it does

• Fuses absorbed-MLA non-positional packing, RoPE, and the contiguous output copy into one Triton kernel (fused_mla_rope_concat) with a fused backward. Upstream measured the long-sequence packing region going from 85 kernels to 2 and from 0.848 ms to 0.099 ms per attention call.
• Enables the fused MLA RoPE path for standard RoPE, which was previously YaRN-only. The two transformer_config.py asserts that forbade RoPE fusion for MLA with standard RoPE and for DSA are removed; a rotary_percent < 1 guard falls back to the unfused path with a one-time warning.
• Extends the in-place MLA RoPE primitive to rotate leading channels so the DSA indexer works directly on its packed [PE, NOPE] projection.

Upstream A/B on 8x B200, GLM-5.2 proxy: step time -2.11% (THD 4K, CP1) and -2.80% (THD 32K, CP8); self-attention GPU span -10.6% / -11.5%; memory unchanged.

Interaction with radixark changes

fused_mla_yarn_rope_apply.py is also modified by radixark commit ab0327a0b (MLA RoPE Triton kernel head indexing and v_dim=0 support, #5). The merge kept every line of that fix; verified by grepping each added line of ab0327a0b in the merged file.

Applicability to Miles: opt-in, not automatic

Miles' scripts/models/deepseek-v4-flash.py passes --no-rope-fusion (since NVIDIA#1910), and the true-on-policy contract disables RoPE fusion for the Megatron backend to keep trainer and SGLang numerics aligned. With apply_rope_fusion=False, should_use_fused_mla_rope() returns False and none of the new kernels run. So this port changes nothing for current Miles recipes. To realize the speedup, a follow-up in Miles would drop --no-rope-fusion for DSv4 runs that do not require bitwise trainer/rollout alignment, and that change should be A/B'd for numerics before adoption.

The port is still worth landing now: it keeps miles-main aligned with upstream ahead of the dev bump, and NVIDIA#7064 (next in this series) builds on the RoPE primitives introduced here.

Validation

4x B200 devbox, radixark/miles:dev, Megatron miles-main @ 73b54618f + this commit, single-process pytest:

File	Result
fusions/test_mla_yarn_rope_apply.py	28 passed
experimental_attention_variant/test_absorbed_mla.py	1 passed, 1 failed (pre-existing)
experimental_attention_variant/test_attention_variant_dsa.py	59 passed, 1 failed (pre-existing)
transformer/test_multi_latent_attention.py	11 passed, 1 skipped, 1 failed (pre-existing)
experimental_attention_variant/test_attention_delay_wgrad.py	2 passed

The three failures reproduce identically on unpatched miles-main: test_absorbed_mla::test_functionality[...tp_cp_sp1] and TestFusedDSAIndexerLossGradientTP::test_fused_indexer_loss_gradient_tp_consistency require a multi-process TP launch (world_size not divisible by model_size in a single process; upstream also marks the latter flaky), and TestParallelMLAAttention::test_gpu_forward_thd_padded[yarn] fails a packed_seq_params assertion on this base. Not caused by this change.

Series

Second of two ports of DeepSeek V4 commits from upstream dev that miles-main lacks. Companion: #97 (upstream NVIDIA#7079). Further DSv4 commits (NVIDIA#7064 fused THD ops, NVIDIA#5992 compact indexer) are left for the full dev bump.
🤖 Generated with Claude Code