[dev] moe(perf): Refactor CP layout organization for Qwen3.5-style hybrid attention model

[yuzhongw-nvidia]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

PR for main: #6515

This PR adds fixed context-parallel (CP) layout support for long-context hybrid models that want to run most layers in cp_partition_mode="contiguous" while still keeping regular attention/GQA in the zigzag layout it requires.

The main goals are:

• Move CP layout conversion to the GQA/regular-attention boundary for Qwen3.5-style schedules. In the target GQA:GDN = 1:3 pattern, GDN chunkwise CP can consume contiguous layout directly, so only the less frequent GQA layers need to convert to zigzag and back. This reduces layout conversion frequency by roughly two thirds.
• Prebuild THD CP layout routes once per microbatch. THD layout conversion can then reuse precomputed split/routing metadata instead of rebuilding the route on the conversion hot path.

E2E perf gain (A proxy model of Qwen3.5-397B):

• Prebuild THD CP layout routes: 1.20x speedup
• Globally using contiguous layout: 1.075x speedup
• Overall: 1.29x speedup

Motivation

The previous fixed CP layout path effectively kept the model in zigzag layout because regular attention requires zigzag. That is correct for attention, but it is not optimal for Qwen3.5-style hybrid stacks where most layers are GDN layers.

For chunkwise CP, GDN can run in contiguous layout. This PR lets the model keep the broader stack in contiguous layout, then converts only when entering regular attention/GQA and converts the attention output back afterward. With a GQA:GDN = 1:3 schedule, this shifts layout shuffles away from the dense GDN path and places them on the sparser attention boundary.

For THD packed sequences, the route for a given microbatch is deterministic. This PR moves route construction into batch preparation and reuses the result for later zigzag <-> contiguous conversions.

Major Changes

1. CP layout conversion package

• Splits megatron/core/context_parallel_layout.py into a package:
  • conversion.py: SBHD/THD tensor conversion and CpPartitionModeConverter.
  • routes.py: THD route construction, encoding/decoding, lookup, and prebuild.
  • metadata.py: helpers for packed sequence layout metadata.
  • policy.py: stage-entry layout policy helpers.
  • __init__.py: compatibility exports for existing imports.
• Adds convert_module_input_tensors_cp_partition_mode() for module-boundary conversion:
  • converts module inputs to the required concrete layout;
  • returns a converter for mapping module outputs back to the caller layout;
  • rejects unsupported dense attention masks/bias and cross-attention key/value states.

2. Attention/GDN layout boundary

• Regular attention converts its input to zigzag when the caller layout is different, then converts the output back to the caller layout.
• GDN chunkwise CP can run directly in contiguous layout.
• GDN headwise CP keeps the existing zigzag requirement.
• This makes cp_partition_mode="contiguous" useful for Qwen3.5-style schedules: most GDN layers avoid conversion, while GQA/regular-attention layers handle their own layout boundary.

3. THD route prebuild

• Adds prebuild_thd_cp_partition_routes() to compute both THD conversion directions during batch construction.
• Updates GPT and Hybrid get_batch paths to annotate the stage-entry layout and prebuild THD routes before model forward.
• Updates route lookup so conversion can consume the prebuilt route. Missing routes still fall back to lazy construction with a warning for compatibility, but the intended path is prebuild during batch preparation.

4. Checks and unsupported cases

This PR adds two layers of validation.

First, TransformerConfig rejects unsupported configuration combinations early:

• contiguous CP currently requires experimental_attention_variant to be gated_delta_net or dsv4_hybrid.
• dsv4_hybrid with CP requires contiguous.
• gated_delta_net with linear_cp_mode="headwise" cannot use contiguous.
• multi_latent_attention outside dsv4_hybrid cannot use contiguous.

Second, runtime layout is carried through PackedSeqParams.cp_partition_mode, and attention variants validate the layout they receive:

• Regular attention requires zigzag.
• GDN chunkwise CP requires contiguous.
• GDN headwise CP requires zigzag.
• MTP validates its rolling path and rejects non-packed/SBHD contiguous layout until that path is implemented.

These checks are meant to fail early when tensor order would otherwise silently diverge from the module's layout assumption.

Limitations / Future Work

• Dense attention masks, attention bias, and cross-attention key/value states are still rejected for CP layout conversion.
• Non-packed/SBHD contiguous MTP is intentionally not supported yet.
• Full-iteration CUDA graph is not supported for THD CP layout conversion today because torch-native NCCL A2A only supports fixed splits. Supporting this path will require GIN, and we plan to enable it through DeepEP / HybridEP dispatch.
• TP x CP conversion can be further optimized to use a single A2A communication. The current implementation still needs two communication steps.
• Auto conversion follow-up: <https://github.com/NVIDIA/Megatron-LM/pull/5664|[dev] moe(perf): Refactor CP layout - auto layout conversion by yuzhongw-nvidia · Pull Request [dev] moe(perf): Refactor CP layout - auto layout conversion #5664 · NVIDIA/Megatron-LM>

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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

[yuzhongw-nvidia]

/claude strict-review

[claude]

Strict Review Summary — PR #6387

Findings by severity: CRITICAL: 6 · IMPORTANT: 6 · SUGGESTION: 9 (21 total)

Overall risk: HIGH — do not merge in current state

The refactor itself is well-structured (splitting context_parallel_layout.py into a package, hoisting conversion to the attention boundary, prebuilding THD routes), and the direction is sound. But the branch as pushed does not run.

---

Most impactful findings

1. cp_stage_entry_partition_mode does not exist anywhere in megatron/. (CRITICAL)

grep -rn "cp_stage_entry_partition_mode" megatron/ returns zero hits, yet it is read unconditionally at pretrain_gpt.py:438 and pretrain_hybrid.py:362:

decoder = get_attr_wrapped_model(model, "decoder")
cp_partition_mode = decoder.cp_stage_entry_partition_mode

This is not gated on CP size, TP, or attention variant. Every GPT and every hybrid pretraining/SFT run on this branch raises AttributeError before the first microbatch. The same name is passed as a GPTModel(...) kwarg in two test files, where it also raises TypeError — GPTModel.__init__ has no such parameter and no **kwargs.

2. New test coverage cannot have been executed. (CRITICAL)

tests/unit_tests/ssm/test_gated_delta_net.py:17 imports get_experimental_attention_variant_stage_input_cp_partition_mode, which does not exist — a module-level ImportError that takes down all GDN tests in that file, including pre-existing ones. And tests/unit_tests/models/test_hybrid_moe_model.py:65 sets GOLDEN_CONFIG["cp_partition_mode"] = None, which contradicts the live TransformerConfig default of "zigzag" (transformer_config.py:312) — a default that __post_init__:1556 also explicitly requires be non-None. test_constructor fails with [CHANGED DEFAULTS].

Taken together, (1) and (2) suggest a partially-landed change: it looks like a TransformerConfig.cp_partition_mode = None default and a decoder-level stage-entry attribute were both intended, but only the PackedSeqParams half arrived.

3. Silent numerical corruption on the flash-decode path. (CRITICAL)

Attention.forward has exactly two exits. Line 1747 applies back_to_input_converter; the flash-decode early return at line 1547 does not. Under a "contiguous" stage-entry mode the output is returned in the "zigzag" layout while packed_seq_params claims "zigzag" — so downstream modules compute on mis-permuted tokens rather than asserting. Affects speculative-decoding / MTP verification with CP>1.

4. The route prebuild does not actually remove the syncs it was meant to remove. (IMPORTANT Performance)

decode_thd_cp_partition_route performs three blocking device-to-host copies (routes.py:309 plus two calls into routes.py:278), and conversion.py:440 invokes it on every THD conversion of every tensor. The route tensor is built on device at prebuild time and then immediately copied back to host on each use. Every .cpu() is a cudaStreamSynchronize immediately preceding an all-to-all, which is precisely where overlap matters. The header and split sizes are pure host-side functions of (cu_seqlens, cp_size, cp_rank, source, target) and are already available as a Python list inside _encode_thd_cp_partition_route — keep them there.

5. Backward-compatibility break with no deprecation window. (IMPORTANT Compatibility)

PackedSeqParams.cp_partition_mode default "zigzag" becomes None. External callers that construct PackedSeqParams by hand (NeMo, third-party SFT harnesses) previously got a working default and now hit new hard ValueErrors in attention.py:1389, multi_token_prediction._roll_tensor_packed_seq, and data_schedule.py:643. Meanwhile TransformerConfig.cp_partition_mode still defaults to "zigzag" — the same logical setting now has opposite defaults on the two objects.

6. GDN init and RNG stream changed inside a CP-layout PR. (IMPORTANT Compatibility)

The new kaiming_uniform_ / bias init in reset_parameters plus the torch.random.fork_rng wrapper around the nn.Conv1d constructor change both the initial weights and the RNG stream consumed — so every subsequent Megatron-tracked init in the model shifts, not just conv1d. The fan_in math is correct; the problem is that existing seeds no longer reproduce and GDN golden values will mismatch. This deserves its own commit and a release note.

---

Also flagged

• pretrain_hybrid.py:382 — the new FLOPs fallback uses packed_seq_params.cu_seqlens_q, which _build_packed_seq_params_for_batch:97 sets to the padded boundaries; the deleted comment explained exactly why unpadded is required. Inflates reported MFU on the sequence_packing_scheduler path.
• megatron/core/utils.py:2462 — loop-invariant CPU-to-GPU index copy inside the per-key loop.
• megatron/core/transformer/attention.py:317 — self.cp_comm_type has no read path (mandatory-check finding).
• context_parallel_layout/policy.py:54-63 — get_stage_entry_partition_mode lets the packed value win and then asserts it against the expected value, so for every caller in the repo the assert compares a value against itself. The central layout guard of this PR currently validates nothing.
• routes.py:509 — blanket except ValueError around build_thd_cp_partition_route swallows genuine cu_seqlens validation errors into route = None, which then surfaces as a misleading "missing precomputed routes" FutureWarning from a different layer.
• context_parallel_layout/__init__.py:18 — CpPartitionMode is defined in the package __init__ and imported back by submodules whose import order is the only thing keeping it working. A routine uv run isort (required by CLAUDE.md after any import edit) will reorder this into an ImportError that breaks the whole megatron.core import graph.
• tests/unit_tests/ssm/test_gdn_moe_cp_loss_parity.py:49 — _DIAGNOSTIC_REPEATS = 5 expands to ~40 heavy 4096-token EP4 MoE runs per CI invocation; line 63 asserts on a warning substring ("TransformerConfig.cp_partition_mode is deprecated and ignored") that no code emits.
• Plus naming/simplification items on recompute.py:97, data_schedule.py:643, pretrain_gpt.py:215 (three near-duplicate blocks), gpt_model.py:720, and conversion.py:511.

---

Process

The PR head commit is literally titled DEBUG-ONLY REMOVE BEFORE MERGE: test: add non-auto GDN MoE CP loss parity coverage, which is consistent with the state above — this reads as a work-in-progress push rather than a merge candidate.

Suggested path forward

1. Land the missing cp_stage_entry_partition_mode definition (or derive the mode from config.cp_partition_mode) and decide the cp_partition_mode default story on TransformerConfig and PackedSeqParams together.
2. Run the affected unit tests locally before re-pushing — the four CRITICAL findings above are all things a single green run would have caught:
   uv run python -m torch.distributed.run --nproc-per-node 8 -m pytest -q tests/unit_tests/ssm/test_gated_delta_net.py tests/unit_tests/models/test_hybrid_moe_model.py tests/unit_tests/test_context_parallel_layout.py tests/unit_tests/transformer/test_attention.py
3. Fix the flash-decode back-conversion and the per-conversion host syncs.
4. Split the GDN init/RNG change out, and regenerate GDN/hybrid functional golden values (Run functional tests label, then download_golden_values.py --source github --pipeline-id <run-id>).
5. Drop _DIAGNOSTIC_REPEATS to 1 and remove the debug-only commit.

Happy to re-review once the entry points execute.

[yuzhongw-nvidia]

/ok to test 63faa41

[yuzhongw-nvidia]

/ok to test 3bdaba2