[alignment-megatron] PR7A: SGLang Ulysses CP attention and RoPE

[maocheng23]

Summary

Implements PR 7A — the Megatron-side half of Ulysses (cp_comm_type="a2a") context-parallel support for the SGLang backend, pairing with Miles PR 11A. Stacked on PR 7 (#26).

Non-Ulysses paths are unchanged: when cp_comm_type is not "a2a" (or cp_size <= 1), SGLangCoreAttention falls back to DotProductAttention and RoPE continues to use the existing cp_size/cp_rank from the CP group.

Contract satisfied

Contract H (Ulysses CP-aware training data, logprob, and loss) on the Megatron scoring path:

• Under Ulysses, the full sequence is kept on every rank and sharding happens only over attention heads (a2a over the head axis).
• SGLang-backend attention now has a Ulysses implementation (SGLangFlashAttention) that mirrors SGLang's inference layout, so scoring and rollout agree when Miles routes true-on-policy through this backend.
• RoPE is applied to the full unsplit sequence under Ulysses so position encoding matches the rollout side.

Files changed

megatron/core/extensions/sglang.py

• SGLangFlashAttention: FA3 varlen attention for Ulysses. Slices heads across the CP group on input, calls flash_attn_varlen_func, all-gathers heads on output. Handles packed (thd) and unpacked inputs, BF16 cast, GQA repeat-interleave, and FA3 kwarg-compat via inspect.signature.
• SGLangCoreAttention: thin dispatcher — picks SGLangFlashAttention iff config.context_parallel_size > 1 and cp_comm_type == "a2a", otherwise DotProductAttention. Preserves current_max_attn_logits plumbing that Megatron expects.
• SGLangSpecProvider.core_attention(): now returns SGLangCoreAttention. The dispatcher defaults back to DotProductAttention outside Ulysses, so this is a no-op for non-Ulysses runs.

megatron/core/models/common/embeddings/rope_utils.py

• _apply_rotary_pos_emb_thd and apply_rotary_pos_emb take a new ulysses_cp: bool flag (defaults False).
• When ulysses_cp=True, RoPE uses cp_size=1 / cp_rank=0, applying freqs to the full unsplit sequence — matching SGLang's sequence layout.

megatron/core/transformer/attention.py

• _is_ulysses_cp(config) helper reads cp_comm_type from config, handling both str and list[str] forms and the cp_size <= 1 cases.
• Threads ulysses_cp into both Q and K apply_rotary_pos_emb call sites in the split_qkv path.

Tests

tests/unit_tests/extension/test_sglang_ulysses_cp.py

• test_is_ulysses_cp_detects_a2a_mode — covers cp_size=1 + a2a (false), cp_size=2 + p2p (false), cp_size=2 + a2a (true), and the list[str] form ["a2a", "p2p"] (true).
• test_apply_rotary_pos_emb_ulysses_matches_unsplit_sequence — Ulysses RoPE output is identical to an unsplit (cp_size=1) RoPE pass.
• test_sglang_core_attention_dispatches_ulysses_backend — SGLangCoreAttention picks the Ulysses impl under cp_size=2 + a2a.
• test_sglang_core_attention_falls_back_outside_ulysses — fallback to DotProductAttention under cp_size=2 + p2p.

Compatibility

• Non-Ulysses (cp_comm_type != "a2a" or cp_size <= 1): SGLangCoreAttention routes to DotProductAttention, and RoPE defaults to ulysses_cp=False (preserving existing cp_size/cp_rank splitting). No behavior change.
• Use-sglang off: SGLangSpecProvider is not selected, so none of these paths execute.
• FA3 dependency: flash_attn_varlen_func is only imported/used inside SGLangFlashAttention.forward. Non-Ulysses runs do not require FA3.

Known performance impact

• CP=1 / standard CP / use_sglang off: none (dispatcher resolves to existing paths; RoPE defaults preserve existing cp math).
• Ulysses CP: adds an all-gather on attention output and a head-slice on input. Expected — this is how Ulysses is implemented in SGLang, and is what makes scoring numerically match rollout.

Pre-commit

• black: pass.
• isort: pass (PR 7A files only).
• pylint: the file inherits preexisting missing-function-docstring warnings from PR 6's SGLangSpecProvider; no new line-too-long or other regressions are introduced by this PR. Rating moved from 9.85 → 9.89.

Dependencies

• Pairs with Miles PR 11A (#1032) for end-to-end Ulysses CP correctness on the true-on-policy scoring path.
• Logically stacked on PR 7 ([alignment-megatron] PR7: Deterministic TP Reduction And Logits Contract #26). Per the project's existing convention, base is miles-main with the PR diff containing PR 6 + PR 7 + PR 7A commits; reviewers should focus on the latest commit.

Test plan

• pytest tests/unit_tests/extension/test_sglang_ulysses_cp.py -v on CPU (no GPU required for the three included tests)
• Non-Ulysses regression: existing tests/unit_tests/extension/test_sglang_extension.py passes unchanged
• GPU E2E: dense Qwen3 Ulysses CP exact-zero with Miles PR 11A once Milestone 1 E2E is ready (tracked under PR 17 in miles_migration.md)

🤖 Generated with Claude Code