[feat] Add SP and PP kernels for qwen_moe true on policy

[maocheng23]

Stacked on top of #30.

Summary

Adds the Megatron-LM training-side support for Qwen3-30B-A3B MoE SP+PP true-on-policy parity.

This PR is one of three coupled PRs that should be reviewed and landed together because they extend the qwen3_moe_true_on_policy_v1 contract to allow Megatron sequence parallel and pipeline parallel training while preserving exact SGLang rollout parity.

Companion PRs:

• SGLang: [feat] Add Qwen3 MoE true-on-policy SP/PP support sgl-project/sglang#26317
• Miles: [feat] Add SP and PP support for qwen_moe true on policy miles#1088

Main Changes

• Enable the true-on-policy MoE path under sequence parallel and pipeline parallel while preserving the narrow qwen3_moe_true_on_policy_v1 contract gate.
• Add SP-aware tensor-parallel reduce-scatter/all-gather behavior needed by the Qwen3 MoE train/logprob path.
• Adjust MoE token dispatch/shared expert/local-masked expert handling so SP inputs do not use the incompatible padding-compaction path.
• Carry final residuals correctly through the PP schedule and transformer block path used by true-on-policy recompute/training.
• Update RMSNorm/batch-invariant residual handling and the shared schema flag for Megatron SP support.
• Keep debug-only helpers and unused defensive paths out of the final branch.

Validation

Cleaned 8-GPU ion7 real-workload run:

• Run id: moe_sppp_pr_tp2_pp2_ep2_onpolicy_real3_cleaned_noci_nosave_260524_ion7
• Ray job: raysubmit_a9NCtiu6K5dRVCsc
• Topology: Megatron TP=2, PP=2, EP=2, ETP=1, sequence parallel enabled; SGLang 4 rollout engines, each TP=2, EP=2.
• Result: Ray job succeeded, all 8 GPUs returned to 0 MiB.
• Step 1: train/train_rollout_logprob_abs_diff=0.0, train/train_rollout_kl=0.0, grad_norm=0.0377418305, weight version 1.0, mixed version 0.0.
• Step 2: train/train_rollout_logprob_abs_diff=0.0, train/train_rollout_kl=0.0, grad_norm=0.0389085777, weight version 2.0, mixed version 0.0.
• Timing: cleaned step times 485.2571s and 462.8948s, about 2.6-2.9% slower than the previous on-policy no-debug timing run and in the same overhead band versus off-policy.

Focused checks after cleanup:

• git diff --check
• Miles fast checks: 53 passed
• Megatron targeted extension/MoE checks: 3 passed
• Megatron 8-rank tensor-parallel mapping check passed on all ranks
• Full SP+PP 8-GPU E2E exact-logprob rerun on ion7

Local record:
recovery/qwen3_moe_sppp_clean/journal/2026-05-23-sppp-e2e.md

[maocheng23]

Updated this PR with cleanup commit cb0c31f95 (Clean up true-on-policy SPPP plumbing).

What changed:

• Added miles_megatron_plugins/true_on_policy/residual_carrier.py as the single owner for SGLang hidden/residual PP carrier packing, unpacking, output detection, and shape extension.
• Moved common true-on-policy runtime decisions into contracts.py helpers.
• Cleaned call sites in GPT model, pipeline schedules, transformer block/layer, MoE layer/dispatcher, row-parallel reductions.
• Added focused unit coverage for policy helpers and residual carrier behavior.

Validation on ion7 (miles-maocheng-e2e):

• git diff --check passed.
• python3 -m compileall on edited files passed.
• Targeted pytest passed: 6 passed, 38 warnings.
• Normal TP2/PP2/EP2 E2E, SGLang TP2/EP2, SGLANG_MEM_FRACTION_STATIC=0.90:
  • on-policy short: logprob abs diff 0.0 on steps 1 and 2.
  • off-policy short: logprob abs diff 0.0189507, 0.0135899.
  • on-policy ns4/512: logprob abs diff 0.0.
  • off-policy ns4/512: logprob abs diff 0.0168828.

Note: grad_norm remained 0.0 in both on/off-policy smoke runs because the sampled DAPO math rollouts were truncated and reward/advantage were 0. I do not treat that as meaningful gradient-quality evidence; the correctness signal here is exact on-policy rollout/train logprob parity versus non-zero off-policy drift.

Rejected path: rollout TP4 / SGLang EP2 payload expansion failed during SGLang weight update with CUDA invalid argument, so that temporary Miles experiment was removed and is not part of this PR update.