From f57ad6f12d8d1eb9c147feb3d05ca82218d2082d Mon Sep 17 00:00:00 2001
From: Yan Bai <bayan@nvidia.com>
Date: Wed, 24 Jun 2026 07:51:20 -0700
Subject: [PATCH] [dev] Megatron Lite (1/4): shared primitive core

Signed-off-by: Yan Bai <bayan@nvidia.com>
---
 .../megatron/lite/model/protocol_utils.py     | 129 ++++++
 .../lite/model/qwen3_5/lite/checkpoint.py     |  17 +-
 .../megatron/lite/model/qwen3_5/lite/model.py |  49 +-
 .../lite/megatron/lite/primitive/bundle.py    |   2 +-
 .../megatron/lite/primitive/ckpt/__init__.py  |   9 +-
 .../lite/megatron/lite/primitive/ckpt/dcp.py  |  51 ++-
 .../megatron/lite/primitive/ckpt/distckpt.py  |  91 ++--
 .../lite/megatron/lite/primitive/data.py      |   6 +-
 .../lite/primitive/kernels/__init__.py        |   6 +
 .../megatron/lite/primitive/kernels/jit.py    |  27 ++
 .../megatron/lite/primitive/kernels/swiglu.py | 141 ++++++
 .../lite/primitive/modules/__init__.py        |   2 +
 .../lite/primitive/modules/dispatcher.py      |  13 +-
 .../lite/primitive/modules/experts.py         |  81 +---
 .../lite/primitive/modules/gated_delta_net.py | 296 +++++++++---
 .../megatron/lite/primitive/modules/gqa.py    |  71 +--
 .../megatron/lite/primitive/modules/mlp.py    |  23 +
 .../megatron/lite/primitive/modules/moe.py    |   8 +-
 .../megatron/lite/primitive/modules/mrope.py  |   4 +-
 .../megatron/lite/primitive/modules/mtp.py    |  21 +-
 .../megatron/lite/primitive/modules/router.py |  21 +-
 .../lite/primitive/optimizers/__init__.py     |   2 +-
 .../lite/primitive/optimizers/fsdp2/adamw.py  |  37 +-
 .../primitive/optimizers/fsdp2/grad_clip.py   |  22 +-
 .../lite/primitive/optimizers/fsdp2/state.py  |  12 +-
 .../lite/primitive/parallel/__init__.py       |  17 +-
 .../megatron/lite/primitive/parallel/cp.py    | 215 ++++++++-
 .../megatron/lite/primitive/parallel/mhc.py   |  52 +++
 .../lite/primitive/parallel/pipeline.py       | 117 ++---
 .../megatron/lite/primitive/parallel/pp.py    |  96 ++--
 .../megatron/lite/primitive/parallel/state.py |   6 +
 .../megatron/lite/primitive/parallel/thd.py   | 336 ++++++++++----
 .../megatron/lite/primitive/train_step.py     |  22 +-
 .../primitive/{utils.py => utils/__init__.py} |   0
 .../lite/megatron/lite/primitive/utils/moe.py | 422 ++++++++++++++++++
 .../lite/primitive/utils/packed_seq.py        |  61 +++
 .../megatron/lite/primitive/utils/rope.py     | 144 ++++++
 .../megatron/lite/primitive/utils/rotary.py   | 299 +++++++++++++
 .../lite/runtime/contracts/__init__.py        |   3 +
 .../megatron/lite/runtime/contracts/config.py |   4 +
 .../megatron/lite/runtime/contracts/data.py   |  17 +-
 .../megatron/lite/runtime/contracts/loss.py   |  52 +++
 .../model/test_qwen_lite_forward_smoke.py     |  16 +-
 .../test_parallel_topologies_smoke.py         |  54 ++-
 .../unit/primitive/test_attention_moe_unit.py |  11 +-
 ...st_parallel_dimensions_independent_unit.py |  14 +-
 .../unit/primitive/test_parallel_unit.py      | 302 ++++++++++++-
 47 files changed, 2841 insertions(+), 560 deletions(-)
 create mode 100644 experimental/lite/megatron/lite/model/protocol_utils.py
 create mode 100644 experimental/lite/megatron/lite/primitive/kernels/__init__.py
 create mode 100644 experimental/lite/megatron/lite/primitive/kernels/jit.py
 create mode 100644 experimental/lite/megatron/lite/primitive/kernels/swiglu.py
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/mlp.py
 create mode 100644 experimental/lite/megatron/lite/primitive/parallel/mhc.py
 rename experimental/lite/megatron/lite/primitive/{utils.py => utils/__init__.py} (100%)
 create mode 100644 experimental/lite/megatron/lite/primitive/utils/moe.py
 create mode 100644 experimental/lite/megatron/lite/primitive/utils/packed_seq.py
 create mode 100644 experimental/lite/megatron/lite/primitive/utils/rope.py
 create mode 100644 experimental/lite/megatron/lite/primitive/utils/rotary.py
 create mode 100644 experimental/lite/megatron/lite/runtime/contracts/loss.py

diff --git a/experimental/lite/megatron/lite/model/protocol_utils.py b/experimental/lite/megatron/lite/model/protocol_utils.py
new file mode 100644
index 00000000000..f1c4b53d3d5
--- /dev/null
+++ b/experimental/lite/megatron/lite/model/protocol_utils.py
@@ -0,0 +1,129 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Helpers shared by model protocol forward steps.
+
+The verl/runtime layers hand each protocol a raw, model-agnostic ``PackedBatch``
+(true per-sequence lengths, no padding, no ``PackedSeqParams``). Each model owns
+its pack/unpack pair: ``pack_thd_forward_kwargs`` pads + CP-splits the batch into
+model forward kwargs, and ``unpack_thd_forward_output`` reverses a model output
+back to jagged true-length form. THD models share the zigzag-CP pair below;
+models with a different CP layout (e.g. DeepSeek-V4 contiguous DSA) provide their
+own pair.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from megatron.lite.primitive.parallel import ParallelState
+from megatron.lite.primitive.parallel.thd import (
+    pack_nested_thd,
+    parallel_state_from_model,
+    prepare_packed_thd_kwargs_for_context_parallel,
+    thd_pack_meta,
+    unpack_thd_to_nested,
+)
+from megatron.lite.primitive.utils.packed_seq import PackedSeqParams
+from megatron.lite.runtime.contracts.data import PackedBatch
+from megatron.lite.runtime.contracts.loss import get_loss_context
+
+
+def _parallel_state(model) -> ParallelState:
+    return parallel_state_from_model(model) or ParallelState()
+
+
+def nested_from_packed(tensor: torch.Tensor | None, seq_lens: torch.Tensor):
+    """Split a 1-D packed (true, unpadded) tensor back into a jagged nested tensor."""
+    if tensor is None:
+        return None
+    if tensor.dim() == 2 and tensor.size(0) == 1:
+        tensor = tensor.squeeze(0)
+    if tensor.dim() != 1:
+        raise ValueError(f"PackedBatch tensor must be 1-D, got {tuple(tensor.shape)}.")
+    pieces = []
+    offset = 0
+    for length_t in seq_lens:
+        length = int(length_t.item())
+        pieces.append(tensor.narrow(0, offset, length))
+        offset += length
+    if offset != tensor.numel():
+        raise ValueError(f"PackedBatch sizes sum to {offset}, tensor has {tensor.numel()} tokens.")
+    return torch.nested.as_nested_tensor(pieces, layout=torch.jagged)
+
+
+def pack_thd_forward_kwargs(model, batch: PackedBatch) -> dict[str, Any]:
+    """Pad + zigzag-CP-split a raw THD batch into model forward kwargs.
+
+    Pads each sequence to the TE/zigzag alignment, then CP-splits tokens,
+    labels, masks and position ids through the shared THD primitive — the same
+    layout the model was validated against, now produced inside the protocol
+    rather than the connector.
+    """
+    ps = _parallel_state(model)
+    seq_lens = batch.seq_lens
+    packed = pack_nested_thd(
+        nested_from_packed(batch.input_ids, seq_lens),
+        tp_size=ps.tp_size,
+        cp_size=ps.cp_size,
+        cp_rank=ps.cp_rank,
+        cp_group=ps.cp_group if ps.cp_size > 1 else None,
+        split_cp=False,
+        labels=nested_from_packed(batch.labels, seq_lens),
+        roll_labels=batch.labels is not None,
+        loss_mask=nested_from_packed(batch.loss_mask, seq_lens),
+        roll_loss_mask=batch.loss_mask is not None,
+    )
+    max_seqlen = int(packed.padded_lengths.max().item()) if packed.padded_lengths.numel() else 0
+    # pack_nested_thd already returns [1, T] token rows; do not unsqueeze again.
+    kwargs: dict[str, Any] = {
+        "input_ids": packed.input_ids,
+        "labels": packed.labels,
+        "loss_mask": packed.loss_mask,
+        "position_ids": packed.position_ids,
+        "packed_seq_params": PackedSeqParams.from_cu_seqlens(
+            packed.cu_seqlens_padded, max_seqlen=max_seqlen
+        ),
+    }
+    prepare_packed_thd_kwargs_for_context_parallel(model, kwargs)
+    return kwargs
+
+
+def unpack_thd_forward_output(model, batch: PackedBatch, output: torch.Tensor) -> torch.Tensor:
+    """Reverse a zigzag-CP THD model output back to jagged true-length form."""
+    ps = _parallel_state(model)
+    meta = thd_pack_meta(
+        batch.seq_lens,
+        tp_size=ps.tp_size,
+        cp_size=ps.cp_size,
+        cp_group=ps.cp_group if ps.cp_size > 1 else None,
+    )
+    return unpack_thd_to_nested(output, meta, contiguous=False)
+
+
+def add_loss_context_kwargs(kwargs: dict[str, Any], *, include_return_log_probs: bool = False) -> None:
+    loss_context = get_loss_context()
+    if loss_context is None:
+        return
+    kwargs["temperature"] = loss_context.temperature
+    kwargs["calculate_entropy"] = loss_context.calculate_entropy
+    if include_return_log_probs:
+        kwargs["return_log_probs"] = loss_context.return_log_probs
+
+
+def add_cross_entropy_fusion(kwargs: dict[str, Any], model) -> None:
+    kwargs["use_fused_kernels"] = bool(getattr(model, "cross_entropy_fusion", False))
+
+
+def set_cross_entropy_fusion(chunks: list, enabled: bool) -> None:
+    for chunk in chunks:
+        chunk.cross_entropy_fusion = bool(enabled)
+
+
+__all__ = [
+    "add_cross_entropy_fusion",
+    "add_loss_context_kwargs",
+    "nested_from_packed",
+    "pack_thd_forward_kwargs",
+    "set_cross_entropy_fusion",
+    "unpack_thd_forward_output",
+]
diff --git a/experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py b/experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py
index 3f55303ab68..ec6efdf8890 100644
--- a/experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py
+++ b/experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py
@@ -35,9 +35,13 @@ def PLACEMENT_FN(param_name: str) -> list:
         return [Replicate(), Replicate(), Replicate(), Shard(0)]
     if "qkv" in param_name and "layer_norm" not in param_name:
         return [Replicate(), Replicate(), Replicate(), Shard(0)]
-    if ("proj" in param_name or "o_proj" in param_name) and (
-        "full_attn" in param_name or "linear_attn" in param_name
+    if (
+        ("proj" in param_name or "o_proj" in param_name)
+        and ("full_attn" in param_name or "linear_attn" in param_name)
+        and "layer_norm" not in param_name
     ):
+        # Row-parallel output proj weight: TP-shard on dim 1. Exclude layer_norm_weight (1-D, replicated
+        # under TP) which otherwise matches here ("in_proj" contains "proj") and gets an invalid Shard(1).
         return [Replicate(), Replicate(), Replicate(), Shard(1)]
     if "gate_up" in param_name and "shared" in param_name:
         return [Replicate(), Replicate(), Replicate(), Shard(0)]
@@ -760,10 +764,19 @@ def export_hf_weights(
     )
 
 
+def save_hf_weights(
+    model: nn.Module | list[nn.Module], path: str, config: Qwen35Config, ps: ParallelState
+) -> None:
+    from megatron.lite.primitive.ckpt.hf_weights import save_hf_weights as _save
+
+    _save(model, path, Qwen35WeightSpec(config), ps, vocab_size=config.vocab_size)
+
+
 __all__ = [
     "EXPERT_CLASSIFIER",
     "PLACEMENT_FN",
     "Qwen35WeightSpec",
     "export_hf_weights",
     "load_hf_weights",
+    "save_hf_weights",
 ]
diff --git a/experimental/lite/megatron/lite/model/qwen3_5/lite/model.py b/experimental/lite/megatron/lite/model/qwen3_5/lite/model.py
index 9b8f36d6048..76e069e8b6a 100644
--- a/experimental/lite/megatron/lite/model/qwen3_5/lite/model.py
+++ b/experimental/lite/megatron/lite/model/qwen3_5/lite/model.py
@@ -2,9 +2,7 @@
 """Qwen3.5 lite native model.
 
 This implementation keeps the lightweight qwen3_moe/lite composition style
-and does not wrap Megatron-Core layer modules. It still reuses Megatron Lite
-parallel/TE primitives and small Megatron atomic RoPE helpers where those are
-already used by other native Megatron Lite modules.
+and uses Megatron Lite parallel, TE, RoPE, and MoE primitives directly.
 """
 
 from __future__ import annotations
@@ -16,15 +14,20 @@
 import torch.nn as nn
 import transformer_engine.pytorch as te
 
-from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
 from megatron.lite.model.qwen3_5.config import Qwen35Config
+from megatron.lite.primitive.kernels.swiglu import bias_swiglu_impl
 from megatron.lite.primitive.modules.dispatcher import TokenDispatcher
-from megatron.lite.primitive.modules.experts import Experts
+from megatron.lite.primitive.modules.experts import Experts, swiglu_with_probs
 from megatron.lite.primitive.modules.gated_delta_net import GatedDeltaNet
 from megatron.lite.primitive.modules.gqa import GQAttention as FullAttention
 from megatron.lite.primitive.modules.gqa import split_grouped_qkvg as _split_grouped_qkvg
 from megatron.lite.primitive.modules.mrope import MultimodalRotaryEmbedding as Qwen35MRoPE
-from megatron.lite.primitive.modules.mtp import MTPBlock, MTPDecoderLayer, MTPLossAutoScaler
+from megatron.lite.primitive.modules.mtp import (
+    MTPBlock,
+    MTPDecoderLayer,
+    MTPLossAutoScaler,
+    roll_mtp_tensor_left,
+)
 from megatron.lite.primitive.modules.router import TopKRouter
 from megatron.lite.primitive.ops.cross_entropy import vocab_parallel_cross_entropy
 from megatron.lite.primitive.ops.linear_cross_entropy import linear_cross_entropy
@@ -37,7 +40,6 @@
     VocabParallelOutput,
     build_pipeline_chunk_layout,
     gather_from_sequence_parallel,
-    roll_packed_thd_left,
     scatter_to_sequence_parallel,
 )
 from megatron.lite.primitive.utils import build_fp8_recipe
@@ -66,7 +68,7 @@ def _collect_sp_grad_params(model: nn.Module) -> list[nn.Parameter]:
 
 
 def _swiglu(x: torch.Tensor) -> torch.Tensor:
-    return bias_swiglu_impl(x, bias=None)
+    return swiglu_with_probs(x, probs=None)
 
 
 def _qwen_mrope_section(config: Qwen35Config) -> list[int]:
@@ -231,6 +233,7 @@ def __init__(
         moe_act_recompute: bool = False,
         use_thd: bool = False,
         deterministic: bool = False,
+        gdn_cp_mode: str = "fla_allgather",
     ):
         super().__init__()
         self.layer_idx = layer_idx
@@ -264,6 +267,7 @@ def __init__(
                 rms_norm_eps=config.rms_norm_eps,
                 ps=ps,
                 deterministic=deterministic,
+                cp_mode=gdn_cp_mode,
             )
         self.mlp_norm = te.RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps, zero_centered_gamma=True
@@ -332,6 +336,7 @@ def __init__(
         mtp_enable_train: bool = False,
         mtp_detach_encoder: bool = False,
         mount_vision_model: bool = False,
+        gdn_cp_mode: str = "fla_allgather",
     ):
         super().__init__()
         del attention_backend_override
@@ -373,6 +378,7 @@ def __init__(
                     moe_act_recompute=moe_act_recompute,
                     use_thd=use_thd,
                     deterministic=getattr(train_config, "deterministic", False),
+                    gdn_cp_mode=gdn_cp_mode,
                 )
                 for idx in self.layer_indices
             ]
@@ -410,6 +416,7 @@ def make_mtp_layer(layer_idx: int) -> MTPDecoderLayer:
                         moe_act_recompute=moe_act_recompute,
                         use_thd=use_thd,
                         deterministic=getattr(train_config, "deterministic", False),
+                        gdn_cp_mode=gdn_cp_mode,
                     ),
                     detach_encoder=mtp_detach_encoder,
                 )
@@ -552,10 +559,10 @@ def _apply_mtp_loss(
         mtp_loss_mask = loss_mask.clone()
         mtp_loss_values = []
         for mtp_hidden in mtp_hidden_states:
-            mtp_labels, _ = roll_packed_thd_left(
+            mtp_labels, _ = roll_mtp_tensor_left(
                 mtp_labels, packed_seq_params=packed_seq_params, dims=-1
             )
-            mtp_loss_mask, num_tokens = roll_packed_thd_left(
+            mtp_loss_mask, num_tokens = roll_mtp_tensor_left(
                 mtp_loss_mask, packed_seq_params=packed_seq_params, dims=-1
             )
             labels_sb = mtp_labels.transpose(0, 1).contiguous()
@@ -658,7 +665,7 @@ def _hook_vision_params_avg_grad_across_tp(module: nn.Module) -> None:
         param.average_gradients_across_tp_domain = True  # type: ignore[assignment]
 
 
-def _build_native_vision_model(hf_path: str) -> nn.Module:
+def _build_native_vision_model(hf_path: str) -> nn.Module | None:
     if not hf_path:
         raise ValueError("mount_vision_model requires hf_path.")
     try:
@@ -669,12 +676,22 @@ def _build_native_vision_model(hf_path: str) -> nn.Module:
     hf_config = AutoConfig.from_pretrained(hf_path, trust_remote_code=True)
     vision_config = getattr(hf_config, "vision_config", None)
     if vision_config is None:
-        raise RuntimeError("HF config does not expose vision_config; cannot build vision_model.")
-    hf_vision_cls = _resolve_hf_vision_cls(hf_config, hf_path)
-    if hasattr(hf_vision_cls, "_from_config"):
-        vision = hf_vision_cls._from_config(vision_config)
+        # Text-only checkpoint (no vision tower): nothing to mount -> graceful skip.
+        return None
+    auto_map = getattr(hf_config, "auto_map", None) or {}
+    if auto_map:
+        # Remote-code checkpoint: resolve the vision class from its dynamic module.
+        hf_vision_cls = _resolve_hf_vision_cls(hf_config, hf_path)
+        if hasattr(hf_vision_cls, "_from_config"):
+            vision = hf_vision_cls._from_config(vision_config)
+        else:
+            vision = hf_vision_cls(vision_config)
     else:
-        vision = hf_vision_cls(vision_config)
+        # Native-transformers checkpoint (auto_map=None, e.g. Qwen3.5-35B-A3B): build the vision tower
+        # directly from its vision_config via the native AutoModel registry (no remote code).
+        from transformers import AutoModel
+
+        vision = AutoModel.from_config(vision_config)
     _hook_fp32_rotary_emb(vision)
     _hook_vision_params_avg_grad_across_tp(vision)
     return vision.to(torch.bfloat16)
diff --git a/experimental/lite/megatron/lite/primitive/bundle.py b/experimental/lite/megatron/lite/primitive/bundle.py
index e74603d060e..439968c1c6f 100644
--- a/experimental/lite/megatron/lite/primitive/bundle.py
+++ b/experimental/lite/megatron/lite/primitive/bundle.py
@@ -24,6 +24,6 @@ class ModelBundle:
     parallel_state: ParallelState
     optimizer: Any | None = None
     finalize_grads: Callable[[], None] | None = None
-    forward_step: Callable[[nn.Module, dict], dict] | None = None
+    forward_step: Callable[..., dict] | None = None
     # extra metadata (expert_classifier, model_cfg, etc.)
     extras: dict[str, Any] = field(default_factory=dict)
diff --git a/experimental/lite/megatron/lite/primitive/ckpt/__init__.py b/experimental/lite/megatron/lite/primitive/ckpt/__init__.py
index 14ff5b4d884..4c2842a1220 100644
--- a/experimental/lite/megatron/lite/primitive/ckpt/__init__.py
+++ b/experimental/lite/megatron/lite/primitive/ckpt/__init__.py
@@ -2,7 +2,6 @@
 """Checkpoint helpers."""
 
 from megatron.lite.primitive.ckpt.dcp import load_training_checkpoint, save_training_checkpoint
-from megatron.lite.primitive.ckpt.distckpt import attach_model_sharded_state_dict
 from megatron.lite.primitive.ckpt.hf_weights import HFWeights
 
 __all__ = [
@@ -11,3 +10,11 @@
     "load_training_checkpoint",
     "save_training_checkpoint",
 ]
+
+
+def __getattr__(name: str):
+    if name == "attach_model_sharded_state_dict":
+        from megatron.lite.primitive.ckpt.distckpt import attach_model_sharded_state_dict
+
+        return attach_model_sharded_state_dict
+    raise AttributeError(name)
diff --git a/experimental/lite/megatron/lite/primitive/ckpt/dcp.py b/experimental/lite/megatron/lite/primitive/ckpt/dcp.py
index cb5fb1ebc2c..f53fe705877 100644
--- a/experimental/lite/megatron/lite/primitive/ckpt/dcp.py
+++ b/experimental/lite/megatron/lite/primitive/ckpt/dcp.py
@@ -58,15 +58,15 @@ def save_training_checkpoint(
     if not use_dcp:
         _save_local_training_checkpoint(model, optimizer, step, path, save_rng=save_rng)
         return
-    if _supports_distopt_distckpt(model, optimizer):
+    if _supports_dist_opt_distckpt(model, optimizer):
         ckpt_path = os.path.join(path, f"step_{step}")
         os.makedirs(ckpt_path, exist_ok=True)
-        _save_distopt_checkpoint(
+        _save_dist_opt_checkpoint(
             model, optimizer, step, ckpt_path, save_model=save_model, save_optimizer=save_optimizer
         )
         if save_rng:
             _save_rng_sidecar(ckpt_path)
-        log_rank0(f"Saved distopt checkpoint at step {step} to {ckpt_path}")
+        log_rank0(f"Saved dist_opt checkpoint at step {step} to {ckpt_path}")
         return
     if config is None or ps is None:
         raise ValueError("DCP checkpointing requires config and ParallelState.")
@@ -74,12 +74,17 @@ def save_training_checkpoint(
         raise TypeError("DCP checkpointing currently expects a single nn.Module.")
     dense_mesh, expert_mesh = _build_meshes(config)
     state_dict: dict = {"step": step}
+    # Pipeline stages own DIFFERENT parameters but their local layers re-index
+    # to 0..N, so without a per-stage prefix the DCP FQNs collide across pp ranks
+    # (stage0 layer0 and stage1 layer1 both -> "model.0.layers.0..."), corrupting
+    # the round-trip. Mirror distckpt's pp-aware keying: disjoint keyspace per stage.
+    model_prefix = f"model_pp{ps.pp_rank}" if ps.pp_size > 1 else "model"
 
     if save_model:
         for name, param in model.named_parameters():
             placements = get_placements(name)
             mesh = expert_mesh if is_expert(name) else dense_mesh
-            state_dict[f"model.{name}"] = _dcp_tensor_from_param(param, mesh, placements)
+            state_dict[f"{model_prefix}.{name}"] = _dcp_tensor_from_param(param, mesh, placements)
 
     ckpt_path = os.path.join(path, f"step_{step}")
     os.makedirs(ckpt_path, exist_ok=True)
@@ -118,13 +123,13 @@ def load_training_checkpoint(
             load_parameter_state_update_legacy_format=load_parameter_state_update_legacy_format,
         )
     ckpt_path = _resolve_step_checkpoint_path(path)
-    if _supports_distopt_distckpt(model, optimizer):
-        step = _load_distopt_checkpoint(
+    if _supports_dist_opt_distckpt(model, optimizer):
+        step = _load_dist_opt_checkpoint(
             model, optimizer, ckpt_path, load_model=load_model, load_optimizer=load_optimizer
         )
         if load_rng:
             _load_rng_sidecar(ckpt_path)
-        log_rank0(f"Loaded distopt checkpoint from {path} at step {step}")
+        log_rank0(f"Loaded dist_opt checkpoint from {path} at step {step}")
         return step
     if config is None or ps is None:
         raise ValueError("DCP checkpointing requires config and ParallelState.")
@@ -133,18 +138,23 @@ def load_training_checkpoint(
     dense_mesh, expert_mesh = _build_meshes(config)
 
     state_dict: dict = {"step": 0}
+    # Same pp-aware keying as save (see save_training_checkpoint): per-stage
+    # disjoint keyspace so pp ranks don't read each other's colliding FQNs.
+    model_prefix = f"model_pp{ps.pp_rank}" if ps.pp_size > 1 else "model"
 
     if load_model:
         for name, param in model.named_parameters():
             placements = get_placements(name)
             mesh = expert_mesh if is_expert(name) else dense_mesh
-            state_dict[f"model.{name}"] = _empty_dcp_tensor_like_param(param, mesh, placements)
+            state_dict[f"{model_prefix}.{name}"] = _empty_dcp_tensor_like_param(
+                param, mesh, placements
+            )
 
     dcp.load(state_dict, checkpoint_id=ckpt_path)
 
     if load_model:
         for name, param in model.named_parameters():
-            key = f"model.{name}"
+            key = f"{model_prefix}.{name}"
             if key in state_dict:
                 t = state_dict[key]
                 with torch.no_grad():
@@ -172,13 +182,18 @@ def _resolve_step_checkpoint_path(path: str) -> str:
     return path
 
 
-def _supports_distopt_distckpt(model: nn.Module | Iterable[nn.Module], optimizer) -> bool:
-    from megatron.lite.primitive.ckpt.distckpt import supports_distopt_distckpt
+def _supports_dist_opt_distckpt(model: nn.Module | Iterable[nn.Module], optimizer) -> bool:
+    try:
+        from megatron.lite.primitive.ckpt.distckpt import supports_dist_opt_distckpt
+    except ModuleNotFoundError as exc:
+        if exc.name != "megatron.core":
+            raise
+        return False
 
-    return supports_distopt_distckpt(model, optimizer)
+    return supports_dist_opt_distckpt(model, optimizer)
 
 
-def _save_distopt_checkpoint(
+def _save_dist_opt_checkpoint(
     model: nn.Module | Iterable[nn.Module],
     optimizer,
     step: int,
@@ -187,14 +202,14 @@ def _save_distopt_checkpoint(
     save_model: bool,
     save_optimizer: bool,
 ) -> None:
-    from megatron.lite.primitive.ckpt.distckpt import save_distopt_checkpoint
+    from megatron.lite.primitive.ckpt.distckpt import save_dist_opt_checkpoint
 
-    save_distopt_checkpoint(
+    save_dist_opt_checkpoint(
         model, optimizer, step, path, save_model=save_model, save_optimizer=save_optimizer
     )
 
 
-def _load_distopt_checkpoint(
+def _load_dist_opt_checkpoint(
     model: nn.Module | Iterable[nn.Module],
     optimizer,
     path: str,
@@ -202,9 +217,9 @@ def _load_distopt_checkpoint(
     load_model: bool,
     load_optimizer: bool,
 ) -> int:
-    from megatron.lite.primitive.ckpt.distckpt import load_distopt_checkpoint
+    from megatron.lite.primitive.ckpt.distckpt import load_dist_opt_checkpoint
 
-    return load_distopt_checkpoint(
+    return load_dist_opt_checkpoint(
         model, optimizer, path, load_model=load_model, load_optimizer=load_optimizer
     )
 
diff --git a/experimental/lite/megatron/lite/primitive/ckpt/distckpt.py b/experimental/lite/megatron/lite/primitive/ckpt/distckpt.py
index e50769dba68..480d9ae2fed 100644
--- a/experimental/lite/megatron/lite/primitive/ckpt/distckpt.py
+++ b/experimental/lite/megatron/lite/primitive/ckpt/distckpt.py
@@ -1,5 +1,5 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-"""Megatron Core distributed checkpoint bridge for MLite distopt."""
+"""Megatron Core distributed checkpoint bridge for MLite dist_opt."""
 
 from __future__ import annotations
 
@@ -36,29 +36,29 @@ def attach_model_sharded_state_dict(
     get_placements: PlacementFn = default_placement_fn,
     is_expert: ExpertClassifierFn = default_expert_classifier,
 ) -> None:
-    """Attach an MLite-local mcore sharded_state_dict method to distopt chunks."""
+    """Attach an MLite-local mcore sharded_state_dict method to dist_opt chunks."""
 
     for chunk in model_chunks:
         chunk.sharded_state_dict = MethodType(  # type: ignore[method-assign]
             _build_bound_sharded_state_dict(ps, get_placements, is_expert), chunk
         )
-        chunk._mlite_distopt_sharded_state_dict = True  # type: ignore[attr-defined]
-        chunk._mlite_distopt_parallel_state = ps  # type: ignore[attr-defined]
+        chunk._mlite_dist_opt_sharded_state_dict = True  # type: ignore[attr-defined]
+        chunk._mlite_dist_opt_parallel_state = ps  # type: ignore[attr-defined]
 
 
-def supports_distopt_distckpt(model: nn.Module | Iterable[nn.Module], optimizer: Any) -> bool:
+def supports_dist_opt_distckpt(model: nn.Module | Iterable[nn.Module], optimizer: Any) -> bool:
     """Return whether this model/optimizer pair can use mcore dist_checkpointing."""
 
     if optimizer is not None and not callable(getattr(optimizer, "sharded_state_dict", None)):
         return False
     return all(
-        bool(getattr(chunk, "_mlite_distopt_sharded_state_dict", False))
+        bool(getattr(chunk, "_mlite_dist_opt_sharded_state_dict", False))
         and callable(getattr(chunk, "sharded_state_dict", None))
         for chunk in _model_chunks(model)
     )
 
 
-def save_distopt_checkpoint(
+def save_dist_opt_checkpoint(
     model: nn.Module | Iterable[nn.Module],
     optimizer: Any,
     step: int,
@@ -91,7 +91,7 @@ def save_distopt_checkpoint(
     )
 
 
-def load_distopt_checkpoint(
+def load_dist_opt_checkpoint(
     model: nn.Module | Iterable[nn.Module],
     optimizer: Any,
     checkpoint_dir: str,
@@ -113,7 +113,22 @@ def load_distopt_checkpoint(
             )
         finally:
             _restore_state_dict_patches(patches)
-    state_dict = dist_checkpointing.load(load_sd, checkpoint_dir, validate_access_integrity=False)
+    # torch>=2.6 flips torch.load's weights_only default to True, which rejects the trusted dist_opt
+    # common state (mcore's load_common torch.loads optimizer/scheduler classes like AdamW). We are
+    # loading our OWN checkpoint -> force weights_only=False for the duration of the load.
+    _orig_torch_load = torch.load
+
+    def _trusted_torch_load(*args, **kwargs):
+        kwargs.setdefault("weights_only", False)
+        return _orig_torch_load(*args, **kwargs)
+
+    torch.load = _trusted_torch_load
+    try:
+        state_dict = dist_checkpointing.load(
+            load_sd, checkpoint_dir, validate_access_integrity=False
+        )
+    finally:
+        torch.load = _orig_torch_load
     if load_model:
         _load_model_state_dict(model, state_dict)
     if load_optimizer and optimizer is not None and "optimizer" in state_dict:
@@ -155,40 +170,40 @@ def _patch_empty_native_optimizer_state_dicts(
     optimizer: Any, *, fallback_step: int
 ) -> list[tuple[Any, Any]]:
     patches: list[tuple[Any, Any]] = []
-    for distopt in _iter_distributed_optimizers(optimizer):
-        inner = getattr(distopt, "optimizer", None)
+    for dist_opt in _iter_distributed_optimizers(optimizer):
+        inner = getattr(dist_opt, "optimizer", None)
         state = getattr(inner, "state", None)
         if not isinstance(state, MutableMapping) or state:
             continue
-        original_state_dict = distopt.state_dict
+        original_state_dict = dist_opt.state_dict
 
         def patched_state_dict(
-            original_state_dict=original_state_dict, distopt=distopt, fallback_step=fallback_step
+            original_state_dict=original_state_dict, dist_opt=dist_opt, fallback_step=fallback_step
         ):
             try:
                 return original_state_dict()
             except AssertionError:
-                return _empty_native_optimizer_state_dict(distopt, fallback_step)
+                return _empty_native_optimizer_state_dict(dist_opt, fallback_step)
 
-        distopt.state_dict = patched_state_dict  # type: ignore[method-assign]
-        patches.append((distopt, original_state_dict))
+        dist_opt.state_dict = patched_state_dict  # type: ignore[method-assign]
+        patches.append((dist_opt, original_state_dict))
     return patches
 
 
 def _restore_state_dict_patches(patches: list[tuple[Any, Any]]) -> None:
-    for distopt, original_state_dict in patches:
-        distopt.state_dict = original_state_dict  # type: ignore[method-assign]
+    for dist_opt, original_state_dict in patches:
+        dist_opt.state_dict = original_state_dict  # type: ignore[method-assign]
 
 
 def _patch_native_optimizer_step_load(optimizer: Any) -> list[tuple[Any, Any]]:
     patches: list[tuple[Any, Any]] = []
-    for distopt in _iter_distributed_optimizers(optimizer):
-        original_set_state = distopt._set_main_param_and_optimizer_states
+    for dist_opt in _iter_distributed_optimizers(optimizer):
+        original_set_state = dist_opt._set_main_param_and_optimizer_states
 
         def patched_set_state(
-            model_param, tensors, distopt=distopt, original_set_state=original_set_state
+            model_param, tensors, dist_opt=dist_opt, original_set_state=original_set_state
         ):
-            removed_step = _pop_optimizer_step_for_model_param(distopt, model_param, tensors)
+            removed_step = _pop_optimizer_step_for_model_param(dist_opt, model_param, tensors)
             try:
                 return original_set_state(model_param, tensors)
             finally:
@@ -196,25 +211,25 @@ def patched_set_state(
                     state, step = removed_step
                     state["step"] = step
 
-        distopt._set_main_param_and_optimizer_states = patched_set_state  # type: ignore[method-assign]
-        patches.append((distopt, original_set_state))
+        dist_opt._set_main_param_and_optimizer_states = patched_set_state  # type: ignore[method-assign]
+        patches.append((dist_opt, original_set_state))
     return patches
 
 
 def _restore_set_state_patches(patches: list[tuple[Any, Any]]) -> None:
-    for distopt, original_set_state in patches:
-        distopt._set_main_param_and_optimizer_states = original_set_state  # type: ignore[method-assign]
+    for dist_opt, original_set_state in patches:
+        dist_opt._set_main_param_and_optimizer_states = original_set_state  # type: ignore[method-assign]
 
 
 def _pop_optimizer_step_for_model_param(
-    distopt: Any, model_param, tensors: dict[str, Any]
+    dist_opt: Any, model_param, tensors: dict[str, Any]
 ) -> tuple[MutableMapping, Any] | None:
     if "step" in tensors:
         return None
     try:
-        group_index, group_order = distopt.model_param_group_index_map[model_param]
-        main_param = distopt.optimizer.param_groups[group_index]["params"][group_order]
-        state = distopt.optimizer.state[main_param]
+        group_index, group_order = dist_opt.model_param_group_index_map[model_param]
+        main_param = dist_opt.optimizer.param_groups[group_index]["params"][group_order]
+        state = dist_opt.optimizer.state[main_param]
     except (KeyError, IndexError, TypeError):
         return None
     if not isinstance(state, MutableMapping) or "step" not in state:
@@ -269,8 +284,8 @@ def _safe_inner_optimizer(obj: Any) -> Any | None:
     return getattr(obj, "optimizer", None)
 
 
-def _empty_native_optimizer_state_dict(distopt: Any, fallback_step: int) -> dict[str, Any]:
-    inner_state_dict = distopt.optimizer.state_dict()
+def _empty_native_optimizer_state_dict(dist_opt: Any, fallback_step: int) -> dict[str, Any]:
+    inner_state_dict = dist_opt.optimizer.state_dict()
     optimizer_state = {
         key: ([group.copy() for group in value] if key == "param_groups" else value)
         for key, value in inner_state_dict.items()
@@ -280,7 +295,7 @@ def _empty_native_optimizer_state_dict(distopt: Any, fallback_step: int) -> dict
         param_group.pop("params", None)
         param_group["step"] = int(fallback_step)
     state_dict: dict[str, Any] = {"optimizer": optimizer_state}
-    grad_scaler = getattr(distopt, "grad_scaler", None)
+    grad_scaler = getattr(dist_opt, "grad_scaler", None)
     if grad_scaler:
         state_dict["grad_scaler"] = grad_scaler.state_dict()
     return state_dict
@@ -513,11 +528,11 @@ def _load_model_state_dict(
 
 
 def _chunk_parallel_state(chunk: nn.Module) -> ParallelState | None:
-    ps = getattr(chunk, "_mlite_distopt_parallel_state", None)
+    ps = getattr(chunk, "_mlite_dist_opt_parallel_state", None)
     if ps is not None:
         return ps
     wrapped = _wrapped_module(chunk)
-    return getattr(wrapped, "_mlite_distopt_parallel_state", None)
+    return getattr(wrapped, "_mlite_dist_opt_parallel_state", None)
 
 
 def _model_chunks(model: nn.Module | Iterable[nn.Module]) -> list[nn.Module]:
@@ -536,7 +551,7 @@ def _wrapped_module(model: nn.Module) -> nn.Module:
 
 __all__ = [
     "attach_model_sharded_state_dict",
-    "load_distopt_checkpoint",
-    "save_distopt_checkpoint",
-    "supports_distopt_distckpt",
+    "load_dist_opt_checkpoint",
+    "save_dist_opt_checkpoint",
+    "supports_dist_opt_distckpt",
 ]
diff --git a/experimental/lite/megatron/lite/primitive/data.py b/experimental/lite/megatron/lite/primitive/data.py
index 776f7ba9cc8..2f854dfceb2 100644
--- a/experimental/lite/megatron/lite/primitive/data.py
+++ b/experimental/lite/megatron/lite/primitive/data.py
@@ -7,6 +7,8 @@
 
 import torch  # pyright: ignore[reportMissingImports]
 
+from megatron.lite.primitive.utils.packed_seq import PackedSeqParams
+
 _TRUE_ENV_VALUES = {"1", "true", "yes", "on"}
 _FALSE_ENV_VALUES = {"0", "false", "no", "off"}
 
@@ -114,10 +116,6 @@ def infinite_batches_thd(
     split via zigzag striping. position_ids stay FULL because lite's RoPE
     (is_thd_format=False) auto-slices emb internally, same as BSH path.
     """
-    from megatron.core.packed_seq_params import (  # pyright: ignore[reportMissingImports]  # noqa: I001
-        PackedSeqParams,
-    )
-
     if cp_size < 1:
         raise ValueError(f"cp_size must be >= 1, got {cp_size}")
     if cp_rank < 0 or cp_rank >= cp_size:
diff --git a/experimental/lite/megatron/lite/primitive/kernels/__init__.py b/experimental/lite/megatron/lite/primitive/kernels/__init__.py
new file mode 100644
index 00000000000..39af219f47e
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/kernels/__init__.py
@@ -0,0 +1,6 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Optional kernel shims used by MLite primitives."""
+
+from __future__ import annotations
+
+__all__: list[str] = []
diff --git a/experimental/lite/megatron/lite/primitive/kernels/jit.py b/experimental/lite/megatron/lite/primitive/kernels/jit.py
new file mode 100644
index 00000000000..63e630b7a92
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/kernels/jit.py
@@ -0,0 +1,27 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Local JIT/compile decorator shim for small fused Python kernels."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+
+def noop_decorator(func):
+    return func
+
+
+def _build_jit_fuser():
+    if os.environ.get("MEGATRON_LITE_DISABLE_JIT_FUSER", "0") == "1":
+        return noop_decorator
+    compile_fn = getattr(torch, "compile", None)
+    if compile_fn is not None:
+        return compile_fn
+    return torch.jit.script
+
+
+jit_fuser = _build_jit_fuser()
+
+
+__all__ = ["jit_fuser", "noop_decorator"]
diff --git a/experimental/lite/megatron/lite/primitive/kernels/swiglu.py b/experimental/lite/megatron/lite/primitive/kernels/swiglu.py
new file mode 100644
index 00000000000..ec805a57c7f
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/kernels/swiglu.py
@@ -0,0 +1,141 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""SwiGLU fused autograd helpers for MLite."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from megatron.lite.primitive.kernels.jit import jit_fuser
+
+
+@jit_fuser
+def swiglu(y):
+    y_1, y_2 = torch.chunk(y, 2, -1)
+    return F.silu(y_1) * y_2
+
+
+@jit_fuser
+def bias_swiglu(y, bias):
+    y = y + bias
+    return swiglu(y)
+
+
+@jit_fuser
+def weighted_swiglu(y, weights):
+    dtype = y.dtype
+    result = swiglu(y) * weights
+    return result.to(dtype)
+
+
+@jit_fuser
+def swiglu_back(g, y):
+    y_1, y_2 = torch.chunk(y, 2, -1)
+    return torch.cat(
+        (g * torch.sigmoid(y_1) * (1 + y_1 * (1 - torch.sigmoid(y_1))) * y_2, g * F.silu(y_1)), -1
+    )
+
+
+@jit_fuser
+def bias_swiglu_back(g, y, bias):
+    y = y + bias
+    return swiglu_back(g, y)
+
+
+@jit_fuser
+def weighted_swiglu_back(g, y, weights):
+    input_dtype = y.dtype
+    weight_dtype = weights.dtype
+    input_grad = swiglu_back(g * weights, y)
+    weights_grad = swiglu(y) * g.to(weight_dtype)
+    weights_grad = torch.sum(weights_grad, dim=-1, keepdim=True)
+    return input_grad.to(input_dtype), weights_grad.to(weight_dtype)
+
+
+class BiasSwiGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, bias, fp8_input_store=False, cpu_offload_input=False):
+        input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        if cpu_offload_input:
+            input_for_backward.activation_offloading = True
+            if bias is not None:
+                bias.activation_offloading = True
+        ctx.save_for_backward(input_for_backward, bias)
+        ctx.ori_input_dtype = input.dtype
+        ctx.fp8_input_store = fp8_input_store
+        return bias_swiglu(input, bias)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, bias = ctx.saved_tensors
+        input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+        grad = bias_swiglu_back(grad_output, input, bias)
+        return grad, grad, None, None
+
+
+class SwiGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, fp8_input_store=False, cpu_offload_input=False):
+        input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        if cpu_offload_input:
+            input_for_backward.activation_offloading = True
+        ctx.save_for_backward(input_for_backward)
+        ctx.ori_input_dtype = input.dtype
+        ctx.fp8_input_store = fp8_input_store
+        return swiglu(input)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (input,) = ctx.saved_tensors
+        input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+        grad = swiglu_back(grad_output, input)
+        return grad, None, None
+
+
+class WeightedSwiGLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weights, fp8_input_store=False):
+        input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        ctx.save_for_backward(input_for_backward, weights)
+        ctx.ori_input_dtype = input.dtype
+        ctx.fp8_input_store = fp8_input_store
+        return weighted_swiglu(input, weights)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weights = ctx.saved_tensors
+        input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+        input_grad, weights_grad = weighted_swiglu_back(grad_output, input, weights)
+        return input_grad, weights_grad, None
+
+
+def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False):
+    original_shape = input.shape
+    assert len(original_shape) in [2, 3]
+    input = input.view(-1, original_shape[-1])
+    if bias is not None:
+        output = BiasSwiGLUFunction.apply(input, bias, fp8_input_store, cpu_offload_input)
+    else:
+        output = SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input)
+    return (
+        output
+        if len(original_shape) == 2
+        else output.view(original_shape[0], original_shape[1], -1)
+    )
+
+
+def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False):
+    original_shape = input.shape
+    assert len(original_shape) in [2, 3]
+    input = input.view(-1, original_shape[-1])
+    if bias is not None:
+        raise NotImplementedError("Bias is not supported for weighted swiglu fusion")
+    output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store)
+    return (
+        output
+        if len(original_shape) == 2
+        else output.view(original_shape[0], original_shape[1], -1)
+    )
+
+
+__all__ = ["bias_swiglu_impl", "swiglu", "weighted_bias_swiglu_impl"]
diff --git a/experimental/lite/megatron/lite/primitive/modules/__init__.py b/experimental/lite/megatron/lite/primitive/modules/__init__.py
index ff35613ced6..dab0fa137fd 100644
--- a/experimental/lite/megatron/lite/primitive/modules/__init__.py
+++ b/experimental/lite/megatron/lite/primitive/modules/__init__.py
@@ -16,6 +16,7 @@
         "MultimodalRotaryEmbedding",
     ),
     "SigmoidTopKRouter": ("megatron.lite.primitive.modules.router", "SigmoidTopKRouter"),
+    "SwiGLUMLP": ("megatron.lite.primitive.modules.mlp", "SwiGLUMLP"),
     "TokenDispatcher": ("megatron.lite.primitive.modules.dispatcher", "TokenDispatcher"),
     "TopKRouter": ("megatron.lite.primitive.modules.router", "TopKRouter"),
     "_AllToAll": ("megatron.lite.primitive.modules.moe", "_AllToAll"),
@@ -46,6 +47,7 @@ def __getattr__(name: str):
     "MoEAuxLossAutoScaler",
     "MultimodalRotaryEmbedding",
     "SigmoidTopKRouter",
+    "SwiGLUMLP",
     "split_grouped_qkvg",
     "TokenDispatcher",
     "TopKRouter",
diff --git a/experimental/lite/megatron/lite/primitive/modules/dispatcher.py b/experimental/lite/megatron/lite/primitive/modules/dispatcher.py
index e18c6ec5e26..59554b32ca1 100644
--- a/experimental/lite/megatron/lite/primitive/modules/dispatcher.py
+++ b/experimental/lite/megatron/lite/primitive/modules/dispatcher.py
@@ -8,13 +8,10 @@
 import torch  # pyright: ignore[reportMissingImports]
 import torch.distributed as dist  # pyright: ignore[reportMissingImports]
 
-from megatron.core.transformer.moe.moe_utils import (  # pyright: ignore[reportMissingImports]
-    permute,
-    unpermute,
-)
 from megatron.lite.primitive.modules.moe import _AllToAll
 from megatron.lite.primitive.parallel import ParallelState
 from megatron.lite.primitive.utils import ensure_divisible
+from megatron.lite.primitive.utils.moe import permute, unpermute
 
 try:
     import deep_ep  # pyright: ignore[reportMissingImports]
@@ -323,7 +320,13 @@ def _dispatch_alltoall(self, hidden_states, topk_scores, topk_indices):
 
         routing_map = torch.zeros(t, e, dtype=torch.bool, device=hidden_states.device)
         routing_map.scatter_(1, topk_indices, True)
-        num_out = t * topk_indices.size(1)
+        # Use the actual number of routed (token, expert) pairs from routing_map
+        # rather than t * topk: hash routing (ds4) can map a token's topk slots to
+        # DUPLICATE experts, which scatter_ dedups, so t*topk would overcount and
+        # leave permuted.size(0) != sum(input_splits) (all-to-all split mismatch).
+        # Unique-topk routers (every other model) have routing_map.sum() == t*topk,
+        # so this is a no-op for them.
+        num_out = int(routing_map.sum().item())
 
         probs_2d = torch.zeros(t, e, dtype=topk_scores.dtype, device=hidden_states.device)
         probs_2d.scatter_(1, topk_indices, topk_scores)
diff --git a/experimental/lite/megatron/lite/primitive/modules/experts.py b/experimental/lite/megatron/lite/primitive/modules/experts.py
index 2ab182b9905..e3ca4b51ed7 100644
--- a/experimental/lite/megatron/lite/primitive/modules/experts.py
+++ b/experimental/lite/megatron/lite/primitive/modules/experts.py
@@ -10,9 +10,9 @@
 import torch  # pyright: ignore[reportMissingImports]
 import torch.distributed as dist  # pyright: ignore[reportMissingImports]
 import torch.nn as nn  # pyright: ignore[reportMissingImports]
-import torch.nn.functional as F  # pyright: ignore[reportMissingImports]
 import transformer_engine.pytorch as te  # pyright: ignore[reportMissingImports]
 
+from megatron.lite.primitive.kernels.swiglu import bias_swiglu_impl, weighted_bias_swiglu_impl
 from megatron.lite.primitive.modules.lora import (
     LoraConfig,
     SharedGroupedLinearLoRA,
@@ -37,71 +37,21 @@ def _expert_nvtx_range(name: str):
         torch.cuda.nvtx.range_pop()
 
 
-@torch.compile
-def _swiglu(y):
-    y_1, y_2 = torch.chunk(y, 2, -1)
-    return F.silu(y_1) * y_2
-
-
-@torch.compile
-def _weighted_swiglu(y, weights):
-    dtype = y.dtype
-    res = _swiglu(y) * weights
-    return res.to(dtype)
-
-
-@torch.compile
-def _swiglu_back(g, y):
-    y_1, y_2 = torch.chunk(y, 2, -1)
-    return torch.cat(
-        (g * torch.sigmoid(y_1) * (1 + y_1 * (1 - torch.sigmoid(y_1))) * y_2, g * F.silu(y_1)), -1
-    )
-
-
-@torch.compile
-def _weighted_swiglu_back(g, y, weights):
-    input_dtype = y.dtype
-    w_dtype = weights.dtype
-    input_grad = _swiglu_back(g * weights, y)
-    weights_grad = _swiglu(y) * g.to(w_dtype)
-    weights_grad = torch.sum(weights_grad, dim=-1, keepdim=True)
-    return input_grad.to(input_dtype), weights_grad.to(w_dtype)
-
-
-class _WeightedSwiGLUFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, input, weights, fp8_input_store):
-        input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
-        ctx.save_for_backward(input_for_backward, weights)
-        ctx.ori_input_dtype = input.dtype
-        ctx.fp8_input_store = fp8_input_store
-        return _weighted_swiglu(input, weights)
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        input, weights = ctx.saved_tensors
-        input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
-        tmp, wgrad = _weighted_swiglu_back(grad_output, input, weights)
-        return tmp, wgrad, None
-
-
-def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False):
-    """Token-wise-weighted bias swiglu fusion (copied from MC)."""
-    ori_shape = input.shape
-    assert len(ori_shape) in [2, 3]
-    input = input.view(-1, ori_shape[-1])
-    if bias is not None:
-        raise NotImplementedError("Bias is not supported for weighted swiglu fusion")
-    output = _WeightedSwiGLUFunction.apply(input, weights, fp8_input_store)
-    return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
-
-
-def swiglu_with_probs(y: torch.Tensor, probs: torch.Tensor | None) -> torch.Tensor:
+def swiglu_with_probs(
+    y: torch.Tensor, probs: torch.Tensor | None, swiglu_limit: float = 0.0
+) -> torch.Tensor:
     """SwiGLU with optional expert probability scaling."""
+    if swiglu_limit > 0:
+        gate, up = y.chunk(2, dim=-1)
+        up = torch.clamp(up.float(), min=-swiglu_limit, max=swiglu_limit)
+        gate = torch.clamp(gate.float(), max=swiglu_limit)
+        out = torch.nn.functional.silu(gate) * up
+        if probs is not None:
+            out = out * probs
+        return out.to(dtype=y.dtype)
     if probs is not None:
         return weighted_bias_swiglu_impl(y, bias=None, weights=probs)
-    y1, y2 = torch.chunk(y, 2, -1)
-    return F.silu(y1) * y2
+    return bias_swiglu_impl(y, bias=None)
 
 
 class _AllReduceETP(torch.autograd.Function):
@@ -134,6 +84,7 @@ def __init__(
         self.fp8 = fp8
         self.moe_act_recompute = moe_act_recompute
         self.etp_group = ps.etp_group if ps.etp_size > 1 else None
+        self.swiglu_limit = float(getattr(config, "swiglu_limit", 0.0) or 0.0)
 
         self.fc1 = te.GroupedLinear(
             self.num_local_experts,
@@ -234,7 +185,7 @@ def forward(
                 fc1_out = self.fc1(x, m_splits)
                 if self.fc1_lora is not None:
                     fc1_out = fc1_out + self.fc1_lora(x, m_splits)
-                h = act_ckpt.checkpoint(swiglu_with_probs, fc1_out, probs)
+                h = act_ckpt.checkpoint(swiglu_with_probs, fc1_out, probs, self.swiglu_limit)
                 out = self.fc2(h, m_splits)
                 if self.fc2_lora is not None:
                     out = out + self.fc2_lora(h, m_splits)
@@ -243,7 +194,7 @@ def forward(
                 fc1_out = self.fc1(x, m_splits)
                 if self.fc1_lora is not None:
                     fc1_out = fc1_out + self.fc1_lora(x, m_splits)
-                h = swiglu_with_probs(fc1_out, probs)
+                h = swiglu_with_probs(fc1_out, probs, self.swiglu_limit)
                 out = self.fc2(h, m_splits)
                 if self.fc2_lora is not None:
                     out = out + self.fc2_lora(h, m_splits)
diff --git a/experimental/lite/megatron/lite/primitive/modules/gated_delta_net.py b/experimental/lite/megatron/lite/primitive/modules/gated_delta_net.py
index 87b315f8f0d..bae600ed2f0 100644
--- a/experimental/lite/megatron/lite/primitive/modules/gated_delta_net.py
+++ b/experimental/lite/megatron/lite/primitive/modules/gated_delta_net.py
@@ -4,17 +4,25 @@
 from __future__ import annotations
 
 import torch
-import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 import transformer_engine.pytorch as te
 
-from megatron.core.jit import jit_fuser
-from megatron.lite.primitive.ops.gated_delta_rule import l2norm, torch_chunk_gated_delta_rule
-from megatron.lite.primitive.parallel import ColumnParallelLinear, ParallelState, RowParallelLinear
+from megatron.lite.primitive.kernels.jit import jit_fuser
+from megatron.lite.primitive.ops.gated_delta_rule import (
+    l2norm,
+    torch_chunk_gated_delta_rule,
+)
+from megatron.lite.primitive.parallel import (
+    ColumnParallelLinear,
+    ParallelState,
+    RowParallelLinear,
+)
 from megatron.lite.primitive.parallel.cp import (
+    contiguous_to_zigzag_chunks,
     zigzag_reconstruct_from_cp_parts,
     zigzag_slice_for_cp,
+    zigzag_to_contiguous_chunks,
 )
 from megatron.lite.primitive.parallel.thd import (
     reconstruct_packed_from_cp_parts,
@@ -22,6 +30,7 @@
 )
 from megatron.lite.primitive.utils import ensure_divisible
 
+
 try:
     from fla.modules.convolution import (
         causal_conv1d as _fla_causal_conv1d,  # pyright: ignore[reportMissingImports]
@@ -34,9 +43,16 @@
 except ImportError:
     _HAS_FLA = False
 
+try:
+    from fla.ops.cp import build_cp_context as _fla_build_cp_context  # pyright: ignore[reportMissingImports]
+except ImportError:
+    _fla_build_cp_context = None
+
+_CONV_PAD_ALIGNMENT = 4096
+
 
 class GatedDeltaNet(nn.Module):
-    """Native Gated DeltaNet with dense/packed CP reconstruction."""
+    """Native Gated DeltaNet with dense/packed all-gather CP support."""
 
     def __init__(
         self,
@@ -50,10 +66,14 @@ def __init__(
         rms_norm_eps: float,
         ps: ParallelState,
         deterministic: bool = False,
+        cp_mode: str = "fla_allgather",
     ):
         super().__init__()
+        if cp_mode not in {"fla_allgather", "legacy_full_gather"}:
+            raise ValueError(f"Unsupported GatedDeltaNet CP mode: {cp_mode!r}.")
         self.ps = ps
         self.deterministic = bool(deterministic)
+        self.cp_mode = cp_mode
         self.num_k_heads = linear_num_key_heads
         self.num_v_heads = linear_num_value_heads
         self.dk = linear_key_head_dim
@@ -85,19 +105,49 @@ def __init__(
             bias=False,
             padding=linear_conv_kernel_dim - 1,
         )
-        self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads_local, dtype=torch.float32))
-        self.A_log = nn.Parameter(torch.zeros(self.num_v_heads_local, dtype=torch.float32))
+        self.dt_bias = nn.Parameter(
+            torch.ones(self.num_v_heads_local, dtype=torch.float32)
+        )
+        self.A_log = nn.Parameter(
+            torch.zeros(self.num_v_heads_local, dtype=torch.float32)
+        )
         self.norm = te.RMSNorm(self.dv, eps=rms_norm_eps, zero_centered_gamma=True)
         self.o_proj = RowParallelLinear(self.v_dim, hidden_size, ps, bias=False)
+        self._cp_context_cache: dict[
+            tuple[int, int, torch.device], tuple[torch.Tensor, object]
+        ] = {}
 
     def forward(
         self, x: torch.Tensor, position_ids: torch.Tensor | None = None, packed_seq_params=None
     ) -> torch.Tensor:
         del position_ids
+        is_packed = packed_seq_params is not None
         qkvzba = self.in_proj(x).transpose(0, 1).contiguous()
-        cp_restore = None
+        cu_seqlens = self._packed_cu_seqlens(packed_seq_params) if is_packed else None
+        cp_context = None
+        legacy_full_gather = False
         if self.ps.cp_size > 1:
-            qkvzba, cp_restore = self._gather_cp_qkvzba(qkvzba, packed_seq_params)
+            if self.ps.cp_group is None:
+                raise RuntimeError("CP>1 requires ParallelState.cp_group.")
+            if self.cp_mode == "legacy_full_gather":
+                qkvzba, cu_seqlens = self._legacy_full_gather_qkvzba(qkvzba, cu_seqlens)
+                legacy_full_gather = True
+            else:
+                if not _HAS_FLA or _fla_build_cp_context is None:
+                    raise NotImplementedError(
+                        "GatedDeltaNet all-gather CP requires FLA kernels."
+                    )
+                if not is_packed and qkvzba.shape[0] > 1:
+                    raise ValueError(
+                        "GatedDeltaNet all-gather CP with SBHD inputs currently requires "
+                        "micro_batch_size == 1. Use packed THD input or micro_batch_size=1."
+                    )
+                qkvzba = self._cp_swap_qkvzba(
+                    qkvzba,
+                    cu_seqlens if is_packed else None,
+                    to_contiguous=True,
+                )
+                cu_seqlens, cp_context = self._build_cp_context(qkvzba, cu_seqlens)
         batch, seq_len = qkvzba.shape[:2]
         query, key, value, gate, beta, alpha = self._split_proj(qkvzba)
         qkv = torch.cat(
@@ -109,17 +159,9 @@ def forward(
             dim=-1,
         )
 
-        cu_seqlens = None
-        if packed_seq_params is not None:
-            cu_seqlens = (
-                packed_seq_params.cu_seqlens_q_padded
-                if getattr(packed_seq_params, "cu_seqlens_q_padded", None) is not None
-                else packed_seq_params.cu_seqlens_q
-            )
-            if not _HAS_FLA:
-                raise NotImplementedError("GatedDeltaNet packed THD requires FLA kernels.")
-
-        qkv = self._causal_conv1d(qkv, seq_len, cu_seqlens=cu_seqlens)
+        qkv = self._causal_conv1d(
+            qkv, seq_len, cu_seqlens=cu_seqlens, cp_context=cp_context
+        )
         query, key, value, gate, beta, alpha = self._prepare_qkv(
             qkv, gate, beta, alpha, batch, seq_len
         )
@@ -133,68 +175,183 @@ def forward(
             initial_state=None,
             output_final_state=False,
             cu_seqlens=cu_seqlens,
+            cp_context=cp_context,
         )
 
-        if cp_restore is not None:
-            out = self._slice_cp_output(out, cp_restore)
-            gate = self._slice_cp_output(gate, cp_restore)
-            batch, seq_len = out.shape[:2]
         out = self._apply_gated_norm(out, gate)
-        out = out.reshape(batch, seq_len, self.v_dim_local).transpose(0, 1).contiguous()
+        out = out.reshape(batch, seq_len, self.v_dim_local)
+        if self.ps.cp_size > 1:
+            if legacy_full_gather:
+                out = self._legacy_slice_output(out, cu_seqlens)
+            else:
+                out = self._cp_swap_qkvzba(
+                    out,
+                    cu_seqlens if is_packed else None,
+                    to_contiguous=False,
+                )
+            batch, seq_len = out.shape[:2]
+        out = out.transpose(0, 1).contiguous()
         return self.o_proj(out)
 
-    def _all_gather_cp(self, tensor: torch.Tensor) -> list[torch.Tensor]:
-        if self.ps.cp_group is None:
-            raise RuntimeError("CP>1 requires ParallelState.cp_group.")
+    def _all_gather_cp_tensor(self, tensor: torch.Tensor) -> list[torch.Tensor]:
+        if self.ps.cp_size <= 1:
+            return [tensor]
         try:
             from torch.distributed.nn.functional import all_gather
 
             return list(all_gather(tensor, group=self.ps.cp_group))
         except Exception:
             parts = [torch.empty_like(tensor) for _ in range(self.ps.cp_size)]
-            dist.all_gather(parts, tensor, group=self.ps.cp_group)
+            torch.distributed.all_gather(parts, tensor, group=self.ps.cp_group)
             return parts
 
-    def _gather_cp_qkvzba(self, qkvzba: torch.Tensor, packed_seq_params):
-        parts = self._all_gather_cp(qkvzba)
-        if packed_seq_params is not None:
-            cu_seqlens = self._packed_cu_seqlens(packed_seq_params)
-            full = reconstruct_packed_from_cp_parts(
-                parts, cu_seqlens_padded=cu_seqlens, cp_size=self.ps.cp_size, dim=1
+    def _legacy_full_gather_qkvzba(
+        self,
+        qkvzba: torch.Tensor,
+        cu_seqlens: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if cu_seqlens is None:
+            parts = self._all_gather_cp_tensor(qkvzba)
+            return zigzag_reconstruct_from_cp_parts(parts, seq_dim=1), None
+        if qkvzba.shape[0] != 1:
+            raise ValueError("Packed THD GatedDeltaNet expects a single packed batch row.")
+        parts = self._all_gather_cp_tensor(qkvzba[0].contiguous())
+        full = reconstruct_packed_from_cp_parts(
+            parts,
+            cu_seqlens_padded=cu_seqlens,
+            cp_size=self.ps.cp_size,
+            dim=0,
+        )
+        return full.unsqueeze(0).contiguous(), cu_seqlens
+
+    def _legacy_slice_output(
+        self,
+        out: torch.Tensor,
+        cu_seqlens: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if cu_seqlens is None:
+            return zigzag_slice_for_cp(out, self.ps.cp_rank, self.ps.cp_size, seq_dim=1)
+        if out.shape[0] != 1:
+            raise ValueError("Packed THD GatedDeltaNet expects a single packed batch row.")
+        local = split_packed_to_cp_local(
+            out[0].contiguous(),
+            cu_seqlens_padded=cu_seqlens,
+            cp_size=self.ps.cp_size,
+            cp_rank=self.ps.cp_rank,
+            dim=0,
+        )
+        return local.unsqueeze(0).contiguous()
+
+    def _cp_swap_qkvzba(
+        self,
+        tensor: torch.Tensor,
+        cu_seqlens: torch.Tensor | None,
+        *,
+        to_contiguous: bool,
+    ) -> torch.Tensor:
+        if cu_seqlens is None:
+            swap = (
+                zigzag_to_contiguous_chunks
+                if to_contiguous
+                else contiguous_to_zigzag_chunks
             )
-            return full, ("packed", cu_seqlens)
-        full = zigzag_reconstruct_from_cp_parts(parts, seq_dim=1)
-        return full, ("dense",)
-
-    def _slice_cp_output(self, out: torch.Tensor, cp_restore) -> torch.Tensor:
-        kind = cp_restore[0]
-        if kind == "packed":
-            return split_packed_to_cp_local(
-                out,
-                cu_seqlens_padded=cp_restore[1],
-                cp_size=self.ps.cp_size,
-                cp_rank=self.ps.cp_rank,
-                dim=1,
+            return swap(tensor, self.ps.cp_group, seq_dim=1)
+        if tensor.shape[0] != 1:
+            raise ValueError(
+                "Packed THD GatedDeltaNet expects a single packed batch row."
+            )
+        local_cu_seqlens = cu_seqlens // self.ps.cp_size
+        pieces = []
+        swap = (
+            zigzag_to_contiguous_chunks
+            if to_contiguous
+            else contiguous_to_zigzag_chunks
+        )
+        for idx in range(int(local_cu_seqlens.numel()) - 1):
+            start = int(local_cu_seqlens[idx].item())
+            end = int(local_cu_seqlens[idx + 1].item())
+            if end <= start:
+                continue
+            pieces.append(swap(tensor[:, start:end, :], self.ps.cp_group, seq_dim=1))
+        if not pieces:
+            return tensor
+        return torch.cat(pieces, dim=1).contiguous()
+
+    def _build_cp_context(
+        self,
+        qkvzba: torch.Tensor,
+        cu_seqlens: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, object]:
+        if _fla_build_cp_context is None:
+            raise NotImplementedError(
+                "GatedDeltaNet all-gather CP requires FLA cp context."
+            )
+        if cu_seqlens is not None:
+            return (
+                cu_seqlens,
+                _fla_build_cp_context(
+                    cu_seqlens=cu_seqlens,
+                    group=self.ps.cp_group,
+                    conv1d_kernel_size=self.conv1d.kernel_size[0],
+                ),
             )
-        if kind == "dense":
-            return zigzag_slice_for_cp(out, self.ps.cp_rank, self.ps.cp_size, seq_dim=1)
-        raise RuntimeError(f"Unknown CP restore kind: {kind!r}")
+
+        batch, local_seq_len = qkvzba.shape[:2]
+        global_seq_len = local_seq_len * self.ps.cp_size
+        cache_key = (global_seq_len, batch, qkvzba.device)
+        cached = self._cp_context_cache.get(cache_key)
+        if cached is None:
+            dense_cu_seqlens = (
+                torch.arange(batch + 1, device=qkvzba.device, dtype=torch.long)
+                * global_seq_len
+            )
+            cached = (
+                dense_cu_seqlens,
+                _fla_build_cp_context(
+                    cu_seqlens=dense_cu_seqlens,
+                    group=self.ps.cp_group,
+                    conv1d_kernel_size=self.conv1d.kernel_size[0],
+                ),
+            )
+            self._cp_context_cache[cache_key] = cached
+        return cached
 
     def _causal_conv1d(
-        self, qkv: torch.Tensor, seq_len: int, *, cu_seqlens: torch.Tensor | None
+        self,
+        qkv: torch.Tensor,
+        seq_len: int,
+        *,
+        cu_seqlens: torch.Tensor | None,
+        cp_context,
     ) -> torch.Tensor:
-        if cu_seqlens is None:
-            qkv_t = qkv.transpose(1, 2).contiguous()
-            return F.silu(self.conv1d(qkv_t)[:, :, :seq_len].transpose(1, 2))
-        if _HAS_FLA:
+        if _HAS_FLA and (cp_context is not None or cu_seqlens is not None or not self.deterministic):
+            orig_seq_len = qkv.shape[1]
+            pad_n = 0 if cp_context is not None else (-orig_seq_len % _CONV_PAD_ALIGNMENT)
+            conv_input = qkv
+            conv_cu_seqlens = cu_seqlens
+            if pad_n > 0:
+                conv_input = F.pad(qkv, (0, 0, 0, pad_n))
+                if conv_cu_seqlens is not None:
+                    conv_cu_seqlens = conv_cu_seqlens.clone()
+                    conv_cu_seqlens[-1] += pad_n
+            kwargs = {}
+            if cp_context is not None:
+                kwargs["cp_context"] = cp_context
             qkv, _ = _fla_causal_conv1d(
-                x=qkv,
+                x=conv_input,
                 weight=self.conv1d.weight.squeeze(1),
                 bias=None,
                 activation="silu",
-                cu_seqlens=cu_seqlens,
+                cu_seqlens=conv_cu_seqlens,
+                **kwargs,
             )
+            if pad_n > 0:
+                qkv = qkv[:, :orig_seq_len, :]
             return qkv
+        if cu_seqlens is None and cp_context is None:
+            return F.silu(
+                self.conv1d(qkv.transpose(1, 2))[:, :, :seq_len].transpose(1, 2)
+            )
         raise NotImplementedError("GatedDeltaNet packed THD requires FLA causal conv.")
 
     def _gated_delta_rule(
@@ -208,8 +365,12 @@ def _gated_delta_rule(
         initial_state: torch.Tensor | None,
         output_final_state: bool,
         cu_seqlens: torch.Tensor | None,
+        cp_context,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if _HAS_FLA and not self.deterministic:
+        if _HAS_FLA and (cp_context is not None or not self.deterministic):
+            kwargs = {}
+            if cp_context is not None:
+                kwargs["cp_context"] = cp_context
             return _fla_chunk_gated_delta_rule(
                 query,
                 key,
@@ -220,6 +381,11 @@ def _gated_delta_rule(
                 output_final_state=output_final_state,
                 use_qk_l2norm_in_kernel=False,
                 cu_seqlens=cu_seqlens,
+                **kwargs,
+            )
+        if cp_context is not None:
+            raise NotImplementedError(
+                "GatedDeltaNet all-gather CP requires FLA gated delta rule."
             )
         return torch_chunk_gated_delta_rule(
             query,
@@ -239,7 +405,9 @@ def _packed_cu_seqlens(packed_seq_params) -> torch.Tensor:
             else packed_seq_params.cu_seqlens_q
         )
         if cu_seqlens is None:
-            raise ValueError("packed_seq_params must carry cu_seqlens_q for CP GatedDeltaNet.")
+            raise ValueError(
+                "packed_seq_params must carry cu_seqlens_q for CP GatedDeltaNet."
+            )
         return cu_seqlens
 
     def _split_proj(self, qkvzba: torch.Tensor):
@@ -264,9 +432,13 @@ def _split_proj(self, qkvzba: torch.Tensor):
             a.reshape(batch, seq_len, self.num_v_heads_local),
         )
 
-    def _prepare_qkv(self, qkv: torch.Tensor, gate, beta, alpha, batch: int, seq_len: int):
+    def _prepare_qkv(
+        self, qkv: torch.Tensor, gate, beta, alpha, batch: int, seq_len: int
+    ):
         query_key, value = qkv.split([2 * self.qk_dim_local, self.v_dim_local], dim=-1)
-        query_key = query_key.reshape(batch, seq_len, 2 * self.num_k_heads_local, self.dk)
+        query_key = query_key.reshape(
+            batch, seq_len, 2 * self.num_k_heads_local, self.dk
+        )
         value = value.reshape(batch, seq_len, self.num_v_heads_local, self.dv)
         query, key = query_key.split(self.num_k_heads_local, dim=2)
         query = self._l2norm(query.contiguous())
diff --git a/experimental/lite/megatron/lite/primitive/modules/gqa.py b/experimental/lite/megatron/lite/primitive/modules/gqa.py
index 3618379591e..7bb10d36001 100644
--- a/experimental/lite/megatron/lite/primitive/modules/gqa.py
+++ b/experimental/lite/megatron/lite/primitive/modules/gqa.py
@@ -1,42 +1,27 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-"""Grouped Query Attention + Rotary Embedding (MC-atoms).
+"""Grouped Query Attention + Rotary Embedding.
 
 Model-agnostic: takes explicit params instead of model-specific config.
 Supports sequence parallel, context parallel, and THD (packed sequences).
-
-RoPE internals call Megatron-Core's atomic
-`megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding`
-and `rope_utils._apply_rotary_pos_emb_bshd / _apply_rotary_pos_emb_thd` — this
-matches Megatron-Core's unfused rotate-half path because
-``config.apply_rope_fusion`` defaults to ``False``. See
-`docs/gqa_mc_atoms_plan.md` section 4 (Option A).
 """
 
 from __future__ import annotations
 
-import inspect
-
 import torch
 import torch.nn as nn
 import transformer_engine.pytorch as te
 
-from megatron.core.models.common.embeddings.rope_utils import (  # pyright: ignore[reportMissingImports]
-    _apply_rotary_pos_emb_bshd,
-    _apply_rotary_pos_emb_thd,
-)
-from megatron.core.models.common.embeddings.rotary_pos_embedding import (
-    RotaryEmbedding as MCoreRotaryEmbedding,  # pyright: ignore[reportMissingImports]
-)
 from megatron.lite.primitive.modules.gqa_utils import split_grouped_qkvg
 from megatron.lite.primitive.modules.lora import LinearLoRA, LoraConfig, normalize_lora_config
 from megatron.lite.primitive.modules.mrope import MultimodalRotaryEmbedding
 from megatron.lite.primitive.parallel import ColumnParallelLinear, ParallelState, RowParallelLinear
 from megatron.lite.primitive.utils import ensure_divisible
+from megatron.lite.primitive.utils.rope import _apply_rotary_pos_emb_bshd, _apply_rotary_pos_emb_thd
+from megatron.lite.primitive.utils.rotary import RotaryEmbedding
 
 # Whitelist of MC PackedSeqParams fields accepted by TE DotProductAttention.forward().
 # MC-only fields (local_cp_size, cp_group, total_tokens, seq_idx) are excluded.
-# Mirror MC TEDotProductAttention.kept_packed_seq_params pattern
-# (Megatron-LM/megatron/core/extensions/transformer_engine.py:1501-1593).
+# Mirror the TE DotProductAttention packed-sequence argument whitelist.
 _KEPT_PSP_FIELDS = (
     "qkv_format",
     "cu_seqlens_q",
@@ -48,16 +33,6 @@
 )
 
 
-def _callable_accepts_kwarg(fn, kwarg: str) -> bool:
-    try:
-        parameters = inspect.signature(fn).parameters.values()
-    except (TypeError, ValueError):
-        return False
-    return any(
-        param.kind is inspect.Parameter.VAR_KEYWORD or param.name == kwarg for param in parameters
-    )
-
-
 class GQAttention(nn.Module):
     """Grouped Query Attention with TE DotProductAttention.
 
@@ -160,10 +135,7 @@ def __init__(
             )
 
         if self._mrope_section is None:
-            # MC's RotaryEmbedding is atomic (flat kwargs, no TransformerConfig).
-            # cp_group is read from self.cp_group inside forward() when not passed
-            # — no manual CP shard needed on our side.
-            self.rotary = MCoreRotaryEmbedding(
+            self.rotary = RotaryEmbedding(
                 kv_channels=head_dim,
                 rotary_percent=rotary_percent,
                 rotary_interleaved=False,
@@ -178,10 +150,6 @@ def __init__(
                 rotary_base=rope_theta,
                 cp_group=ps.cp_group if ps.cp_size > 1 else None,
             )
-        self._rotary_accepts_packed_seq = self._mrope_section is None and _callable_accepts_kwarg(
-            self.rotary.forward, "packed_seq"
-        )
-
         cp_kwargs = {}
         if ps.cp_size > 1:
             if GQAttention._cp_stream is None:
@@ -217,20 +185,16 @@ def forward(
         q = self.q_norm(q)
         k = self.k_norm(k)
 
-        # RoPE — unfused bshd/thd to match MC's default apply_rope_fusion=False.
-        # MC's RotaryEmbedding.forward takes only `max_seq_len` + optional
-        # `offset`; position_ids is NOT consumed here (MC handles position via
-        # offset for inference / mRoPE via a separate class).
+        # RoPE uses the local unfused rotate-half helpers.
         if self._use_fp32_rope:
             orig_dtype = q.dtype
             q, k = q.float(), k.float()
         if self._mrope_section is not None:
             if position_ids is None:
                 raise ValueError("MRoPE attention requires position_ids.")
-            # For MRoPE the packed THD path applies RoPE directly through the
-            # bshd helper, so the rotary module must slice freqs for this CP
-            # rank before q/k are rotated.
-            freqs = self.rotary(position_ids, self._mrope_section, packed_seq=False)
+            # Packed THD position_ids are already CP-local after protocol.forward
+            # prepares the VERL batch; avoid slicing MRoPE frequencies twice.
+            freqs = self.rotary(position_ids, self._mrope_section, packed_seq=is_thd)
             if is_thd:
                 q = _apply_rotary_pos_emb_bshd(q[:, None], freqs).squeeze(1)
                 k = _apply_rotary_pos_emb_bshd(k[:, None], freqs).squeeze(1)
@@ -244,20 +208,9 @@ def forward(
                 seq_len_for_rope = int(packed_seq_params.cu_seqlens_q[-1])
             else:
                 seq_len_for_rope = int(max(max_q, max_kv))
-            # Match MC RotaryEmbedding.get_rotary_seq_len for packed THD: the
-            # rotary length is the max per-sequence padded length, not total
-            # packed tokens. Using total tokens makes rope_utils switch to
-            # offset mapping, so later packed sequences do not restart at pos 0.
-            #
-            # MC contract (gpt_model.py:380-381): THD path passes packed_seq=True so the
-            # rotary skips its internal cp-slice; _apply_rotary_pos_emb_thd does the
-            # cp-zigzag slice itself via _get_thd_freqs_on_this_cp_rank. Older MC
-            # runtimes do not expose this kwarg; callers without context
-            # parallelism can use the legacy call shape.
-            if self._rotary_accepts_packed_seq:
-                freqs = self.rotary(seq_len_for_rope, packed_seq=True)
-            else:
-                freqs = self.rotary(seq_len_for_rope)
+            # Packed THD uses max per-sequence padded length, not total packed tokens.
+            # The THD apply helper handles CP-zigzag frequency slicing per sequence.
+            freqs = self.rotary(seq_len_for_rope, packed_seq=True)
             q = _apply_rotary_pos_emb_thd(
                 q,
                 packed_seq_params.cu_seqlens_q,
diff --git a/experimental/lite/megatron/lite/primitive/modules/mlp.py b/experimental/lite/megatron/lite/primitive/modules/mlp.py
new file mode 100644
index 00000000000..9d3977d6901
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/mlp.py
@@ -0,0 +1,23 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+import torch
+import torch.nn as nn
+
+from megatron.lite.primitive.modules.experts import swiglu_with_probs
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        *,
+        swiglu_limit: float = 0.0,
+    ):
+        super().__init__()
+        self.gate_up = nn.Linear(hidden_size, 2 * intermediate_size, bias=False)
+        self.down = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.swiglu_limit = float(swiglu_limit or 0.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = swiglu_with_probs(self.gate_up(x), None, self.swiglu_limit)
+        return self.down(y.to(dtype=x.dtype))
diff --git a/experimental/lite/megatron/lite/primitive/modules/moe.py b/experimental/lite/megatron/lite/primitive/modules/moe.py
index c7d294195c5..f80b328dbdb 100644
--- a/experimental/lite/megatron/lite/primitive/modules/moe.py
+++ b/experimental/lite/megatron/lite/primitive/modules/moe.py
@@ -7,11 +7,9 @@
 The qwen3_moe version is used as the canonical form (adds docstring and named
 intermediate variable for clarity).
 
-Note: this is megatron.lite's own MoEAuxLossAutoScaler, kept deliberately separate
-from MC's `megatron.core.transformer.moe.moe_utils.MoEAuxLossAutoScaler`.
-`runtime/backends/mlite/runtime.py` calls `set_loss_scale` on this class to
-apply the 1/num_microbatches aux-loss gradient scale. Megatron-Core MoE modules
-use MC's class directly when they are imported by MC internally.
+Note: this is megatron.lite's own MoEAuxLossAutoScaler. The native MLite
+runtime calls `set_loss_scale` on this class to apply the
+1/num_microbatches aux-loss gradient scale.
 """
 
 from __future__ import annotations
diff --git a/experimental/lite/megatron/lite/primitive/modules/mrope.py b/experimental/lite/megatron/lite/primitive/modules/mrope.py
index 71718d87adc..0efde0006b2 100644
--- a/experimental/lite/megatron/lite/primitive/modules/mrope.py
+++ b/experimental/lite/megatron/lite/primitive/modules/mrope.py
@@ -7,9 +7,7 @@
 import torch.distributed as dist
 import torch.nn as nn
 
-from megatron.core.models.common.embeddings.rope_utils import (  # pyright: ignore[reportMissingImports]
-    get_pos_emb_on_this_cp_rank,
-)
+from megatron.lite.primitive.utils.rope import get_pos_emb_on_this_cp_rank
 
 __all__ = ["MultimodalRotaryEmbedding"]
 
diff --git a/experimental/lite/megatron/lite/primitive/modules/mtp.py b/experimental/lite/megatron/lite/primitive/modules/mtp.py
index 299e143618c..919f027849b 100644
--- a/experimental/lite/megatron/lite/primitive/modules/mtp.py
+++ b/experimental/lite/megatron/lite/primitive/modules/mtp.py
@@ -17,7 +17,22 @@
     scatter_to_sequence_parallel,
 )
 
-__all__ = ["MTPBlock", "MTPDecoderLayer", "MTPLossAutoScaler"]
+__all__ = ["MTPBlock", "MTPDecoderLayer", "MTPLossAutoScaler", "roll_mtp_tensor_left"]
+
+
+def roll_mtp_tensor_left(
+    tensor: torch.Tensor, *, packed_seq_params=None, dims: int = -1
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Roll MTP inputs/labels one token left for dense or packed THD batches."""
+    if packed_seq_params is not None:
+        return roll_packed_thd_left(tensor, packed_seq_params=packed_seq_params, dims=dims)
+
+    dim = dims if dims >= 0 else tensor.dim() + dims
+    rolled = torch.roll(tensor, shifts=-1, dims=dim)
+    index = [slice(None)] * tensor.dim()
+    index[dim] = slice(-1, None)
+    rolled[tuple(index)] = 0
+    return rolled, rolled.sum()
 
 
 class MTPLossAutoScaler(torch.autograd.Function):
@@ -76,9 +91,9 @@ def forward(
         attention_position_ids = (
             rotary_position_ids if rotary_position_ids is not None else position_ids
         )
-        input_ids, _ = roll_packed_thd_left(input_ids, packed_seq_params=packed_seq_params, dims=-1)
+        input_ids, _ = roll_mtp_tensor_left(input_ids, packed_seq_params=packed_seq_params, dims=-1)
         if position_ids is not None:
-            position_ids, _ = roll_packed_thd_left(
+            position_ids, _ = roll_mtp_tensor_left(
                 position_ids, packed_seq_params=packed_seq_params, dims=-1
             )
         decoder_input = scatter_to_sequence_parallel(self.embedding(input_ids), self.ps)
diff --git a/experimental/lite/megatron/lite/primitive/modules/router.py b/experimental/lite/megatron/lite/primitive/modules/router.py
index 821407a2a7b..66ebbb2a7cf 100644
--- a/experimental/lite/megatron/lite/primitive/modules/router.py
+++ b/experimental/lite/megatron/lite/primitive/modules/router.py
@@ -1,11 +1,5 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-"""MoE router implementations: TopKRouter (softmax) and SigmoidTopKRouter.
-
-Internals call the atomic free functions in Megatron-Core's
-`megatron.core.transformer.moe.moe_utils` (plan `docs/moe_mc_wrap_plan.md`
-D3/D4). The outer classes keep the flat-kwargs + `ParallelState` constructor
-style of megatron.lite primitives — no `TransformerConfig`, no mpu globals.
-"""
+"""MoE router implementations: TopKRouter (softmax) and SigmoidTopKRouter."""
 
 from __future__ import annotations
 
@@ -15,13 +9,13 @@
 import torch.distributed as dist  # pyright: ignore[reportMissingImports]
 import torch.nn as nn  # pyright: ignore[reportMissingImports]
 
-from megatron.core.transformer.moe.moe_utils import (  # pyright: ignore[reportMissingImports]
+from megatron.lite.primitive.modules.moe import MoEAuxLossAutoScaler
+from megatron.lite.primitive.utils.moe import (
     compute_routing_scores_for_aux_loss,
     router_gating_linear,
     switch_load_balancing_loss_func,
     topk_routing_with_score_function,
 )
-from megatron.lite.primitive.modules.moe import MoEAuxLossAutoScaler
 
 if TYPE_CHECKING:
     from megatron.lite.primitive.parallel import ParallelState
@@ -129,7 +123,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
 
 
 class SigmoidTopKRouter(nn.Module):
-    """Sigmoid-based TopK router for DeepSeek V3."""
+    """Sigmoid-family TopK router for DeepSeek-style MoE."""
 
     def __init__(
         self,
@@ -149,8 +143,9 @@ def __init__(
             )
         self.topk = config.num_experts_per_tok
         self.num_experts = config.n_routed_experts
-        self.aux_loss_coeff = config.aux_loss_alpha
+        self.aux_loss_coeff = getattr(config, "aux_loss_alpha", 0.0)
         self.scaling_factor = config.routed_scaling_factor
+        self.score_function = getattr(config, "scoring_func", "sigmoid")
         self.router_bias_rate = router_bias_rate
         self.compute_aux_loss = compute_aux_loss
         self.use_pre_softmax = use_pre_softmax
@@ -172,7 +167,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         probs_dense, routing_map = topk_routing_with_score_function(
             logits,
             self.topk,
-            score_function="sigmoid",
+            score_function=self.score_function,
             expert_bias=self.expert_bias.to(logits.dtype),
             scaling_factor=(self.scaling_factor or None),
             fused=self.moe_router_fusion,
@@ -184,7 +179,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
 
         if self.compute_aux_loss and self.training and torch.is_grad_enabled():
             _, aux_scores = compute_routing_scores_for_aux_loss(
-                logits, self.topk, score_function="sigmoid", fused=self.moe_router_fusion
+                logits, self.topk, score_function=self.score_function, fused=self.moe_router_fusion
             )
             tokens_per_expert = routing_map.sum(dim=0).to(torch.int64)
             total_num_tokens = num_tokens
diff --git a/experimental/lite/megatron/lite/primitive/optimizers/__init__.py b/experimental/lite/megatron/lite/primitive/optimizers/__init__.py
index 5bbfdd7cd00..669dd8eae9d 100644
--- a/experimental/lite/megatron/lite/primitive/optimizers/__init__.py
+++ b/experimental/lite/megatron/lite/primitive/optimizers/__init__.py
@@ -6,7 +6,7 @@
 import importlib
 
 BACKENDS = {
-    "mc": "megatron.lite.primitive.optimizers.megatron_wrap",
+    "dist_opt": "megatron.lite.primitive.optimizers.megatron_wrap",
     "fsdp2": "megatron.lite.primitive.optimizers.fsdp2",
 }
 
diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/adamw.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/adamw.py
index 84eb029b530..f7669770f7d 100644
--- a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/adamw.py
+++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/adamw.py
@@ -65,21 +65,20 @@ def copy_local_tensor_to_param_(param: nn.Parameter, local_tensor: torch.Tensor)
         param.detach().copy_(local_tensor.to(device=param.device, dtype=param.dtype))
         return
 
+    # Copy straight into the param's local shard. Reconstructing via
+    # DTensor.from_local mis-sizes an unevenly-sharded param (it infers global =
+    # local * mesh, e.g. a (3,) param over 8 ranks -> 0 or 8), so copy local->local
+    # (master is init'd from this same local shard, so shapes match).
     local_param = to_local_tensor(param)
-    local_value = local_tensor.to(device=local_param.device, dtype=local_param.dtype)
-    from torch.distributed.tensor import DTensor
-
-    param.detach().copy_(DTensor.from_local(local_value, param.device_mesh, param.placements))
+    local_param.copy_(local_tensor.to(device=local_param.device, dtype=local_param.dtype))
 
 
 def all_reduce_grad_(grad: torch.Tensor, *, group: dist.ProcessGroup) -> None:
+    # ``to_local_tensor`` returns the DTensor's local shard storage, so the
+    # in-place all-reduce updates the grad directly -- no DTensor.from_local
+    # round-trip (which mis-sizes unevenly-sharded grads, e.g. (3,) over 8 ranks).
     local_grad = to_local_tensor(grad)
     dist.all_reduce(local_grad, op=dist.ReduceOp.SUM, group=group)
-    if local_grad is grad:
-        return
-    from torch.distributed.tensor import DTensor
-
-    grad.copy_(DTensor.from_local(local_grad, grad.device_mesh, grad.placements))
 
 
 class ChainedOptimizer:
@@ -368,10 +367,7 @@ def maybe_build_te_fused_adam_optimizer(
 ) -> Any | None:
     if not get_bool_opt(opt, "fsdp2_use_te_fused_adam", default=False):
         return None
-    try:
-        from transformer_engine.pytorch.optimizers.fused_adam import FusedAdam
-    except ImportError:
-        return None
+    from transformer_engine.pytorch.optimizers.fused_adam import FusedAdam
 
     all_param_list = list(all_params)
     master_weights = get_bool_opt(
@@ -503,10 +499,23 @@ def iter_torch_optimizers(optimizer: Any) -> Iterable[torch.optim.Optimizer]:
 
 
 def dtensor_from_local(
-    local_tensor: torch.Tensor, device_mesh: Any, placements: Any
+    local_tensor: torch.Tensor,
+    device_mesh: Any,
+    placements: Any,
+    *,
+    shape: Any = None,
+    stride: Any = None,
 ) -> torch.Tensor:
     from torch.distributed.tensor import DTensor
 
+    # ``DTensor.from_local`` infers the global shape as local_shard * mesh, which
+    # is WRONG for unevenly-sharded params (FSDP2 pads the last shard). Pass the
+    # original global shape/stride so the round-trip is exact for any dim not
+    # divisible by the mesh size.
+    if shape is not None:
+        return DTensor.from_local(
+            local_tensor, device_mesh, placements, shape=shape, stride=stride
+        )
     return DTensor.from_local(local_tensor, device_mesh, placements)
 
 
diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/grad_clip.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/grad_clip.py
index f14bdd7fd83..d93149008ac 100644
--- a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/grad_clip.py
+++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/grad_clip.py
@@ -6,7 +6,6 @@
 import math
 from collections import defaultdict
 from collections.abc import Callable, Iterable
-from numbers import Number
 from typing import Any
 
 import torch
@@ -239,20 +238,13 @@ def _local_grad(grad: torch.Tensor, meta: Any | None) -> torch.Tensor:
 
 
 def _scale_grad_(grad: torch.Tensor, scale: float | torch.Tensor) -> None:
-    to_local = getattr(grad, "to_local", None)
-    if not callable(to_local):
-        grad.mul_(_scale_for_tensor(scale, grad))
-        return
-    local_grad = to_local()
-    local_grad.mul_(_scale_for_tensor(scale, local_grad))
-    if DTensor is not None and isinstance(grad, DTensor):
-        grad.copy_(DTensor.from_local(local_grad, grad.device_mesh, grad.placements))
-
-
-def _scale_for_tensor(scale: float | torch.Tensor, tensor: torch.Tensor) -> float | torch.Tensor:
-    if isinstance(scale, Number):
-        return float(scale)
-    return scale.to(device=tensor.device, dtype=tensor.dtype)
+    # clip_coef is a scalar; scale every shard by it in place. A plain scalar mul_
+    # is correct for ANY DTensor placement (Shard/Replicate/Partial) and avoids a
+    # to_local()/from_local() round-trip, which mis-reconstructs the global shape of
+    # an unevenly-sharded param -- e.g. a (3,) mHC scale FSDP-sharded over 8 ranks:
+    # from_local assumes even sharding and infers dim0=8, so copy_ raises
+    # "tensor a (3) must match tensor b (8) at dim 0".
+    grad.mul_(float(scale) if isinstance(scale, torch.Tensor) else scale)
 
 
 def _dtensor_meta(param: nn.Parameter, grad: torch.Tensor) -> Any | None:
diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/state.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/state.py
index 3dd0f06df4c..b8efb65d0fc 100644
--- a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/state.py
+++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/state.py
@@ -21,6 +21,10 @@ class OffloadedStateEntry:
     is_dtensor: bool = False
     device_mesh: Any | None = None
     placements: Any | None = None
+    # Global shape/stride captured so an unevenly-sharded DTensor reconstructs
+    # exactly on reload (from_local would otherwise infer local_shard * mesh).
+    global_shape: Any | None = None
+    global_stride: Any | None = None
 
 
 def move_optimizer_state_to_cpu(
@@ -48,6 +52,8 @@ def move_optimizer_state_to_cpu(
                         is_dtensor=True,
                         device_mesh=value.device_mesh,
                         placements=value.placements,
+                        global_shape=tuple(value.shape),
+                        global_stride=tuple(value.stride()),
                     )
                     param_state[key] = local_value.detach().to("cpu")
                     continue
@@ -80,7 +86,11 @@ def move_offloaded_optimizer_state_to_device(
                     device_value = value.to(entry.device, non_blocking=True)
                     if entry.is_dtensor:
                         param_state[state_key] = dtensor_from_local(
-                            device_value, entry.device_mesh, entry.placements
+                            device_value,
+                            entry.device_mesh,
+                            entry.placements,
+                            shape=entry.global_shape,
+                            stride=entry.global_stride,
                         )
                     else:
                         param_state[state_key] = device_value
diff --git a/experimental/lite/megatron/lite/primitive/parallel/__init__.py b/experimental/lite/megatron/lite/primitive/parallel/__init__.py
index c6fc5701e8a..8ba90b5a1de 100644
--- a/experimental/lite/megatron/lite/primitive/parallel/__init__.py
+++ b/experimental/lite/megatron/lite/primitive/parallel/__init__.py
@@ -4,14 +4,19 @@
 from __future__ import annotations
 
 from megatron.lite.primitive.parallel.cp import (
+    contiguous_to_zigzag_chunks,
     split_packed_for_cp,
     zigzag_position_ids_for_cp,
     zigzag_reconstruct_from_cp_parts,
     zigzag_slice_for_cp,
     zigzag_split_for_cp,
+    zigzag_to_contiguous_chunks,
 )
 from megatron.lite.primitive.parallel.pipeline import forward_backward_pipelining
-from megatron.lite.primitive.parallel.pp import PipelineChunkLayout, build_pipeline_chunk_layout
+from megatron.lite.primitive.parallel.pp import (
+    PipelineChunkLayout,
+    build_pipeline_chunk_layout,
+)
 from megatron.lite.primitive.parallel.sp import (
     gather_for_non_sp_head,
     gather_from_sequence_parallel,
@@ -21,7 +26,11 @@
 from megatron.lite.primitive.parallel.thd import (
     PackedSeqParams,
     PackedTHDBatch,
+    has_packed_thd_params,
     pack_nested_thd,
+    parallel_state_from_model,
+    prepare_packed_thd_for_context_parallel,
+    prepare_packed_thd_kwargs_for_context_parallel,
     reconstruct_packed_from_cp_parts,
     roll_packed_thd_left,
     split_packed_to_cp_local,
@@ -48,6 +57,7 @@ def __getattr__(name: str):
 
 __all__ = [
     "ColumnParallelLinear",
+    "contiguous_to_zigzag_chunks",
     "PackedSeqParams",
     "PackedTHDBatch",
     "PipelineChunkLayout",
@@ -60,9 +70,13 @@ def __getattr__(name: str):
     "forward_backward_pipelining",
     "gather_for_non_sp_head",
     "gather_from_sequence_parallel",
+    "has_packed_thd_params",
     "init_parallel",
     "pad_vocab_for_tp",
     "pack_nested_thd",
+    "parallel_state_from_model",
+    "prepare_packed_thd_for_context_parallel",
+    "prepare_packed_thd_kwargs_for_context_parallel",
     "reconstruct_packed_from_cp_parts",
     "roll_packed_thd_left",
     "scatter_to_sequence_parallel",
@@ -73,4 +87,5 @@ def __getattr__(name: str):
     "zigzag_reconstruct_from_cp_parts",
     "zigzag_slice_for_cp",
     "zigzag_split_for_cp",
+    "zigzag_to_contiguous_chunks",
 ]
diff --git a/experimental/lite/megatron/lite/primitive/parallel/cp.py b/experimental/lite/megatron/lite/primitive/parallel/cp.py
index 2b1f4c21490..26820e188f4 100644
--- a/experimental/lite/megatron/lite/primitive/parallel/cp.py
+++ b/experimental/lite/megatron/lite/primitive/parallel/cp.py
@@ -3,11 +3,17 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
+import torch.distributed as dist
 
 
 def zigzag_split_for_cp(
-    tensor: torch.Tensor, cp_rank: int, cp_size: int, seq_dim: int = 1
+    tensor: torch.Tensor,
+    cp_rank: int,
+    cp_size: int,
+    seq_dim: int = 1,
 ) -> torch.Tensor:
     """Split tensor along sequence dim using zigzag (striped) pattern for CP.
 
@@ -27,7 +33,11 @@ def zigzag_split_for_cp(
     shape = list(tensor.shape)
     shape[seq_dim : seq_dim + 1] = [2 * cp_size, seq_len // (2 * cp_size)]
     tensor = tensor.view(*shape)
-    idx = torch.tensor([cp_rank, 2 * cp_size - cp_rank - 1], dtype=torch.long, device=tensor.device)
+    idx = torch.tensor(
+        [cp_rank, 2 * cp_size - cp_rank - 1],
+        dtype=torch.long,
+        device=tensor.device,
+    )
     tensor = tensor.index_select(seq_dim, idx)
     shape[seq_dim : seq_dim + 2] = [seq_len // cp_size]
     return tensor.reshape(*shape)
@@ -79,8 +89,203 @@ def zigzag_slice_for_cp(
     return torch.cat((first, second), dim=seq_dim).contiguous()
 
 
+def contiguous_slice_for_cp(
+    tensor: torch.Tensor, cp_rank: int, cp_size: int, seq_dim: int = 1
+) -> torch.Tensor:
+    """Return one rank's contiguous CP shard from a full sequence tensor."""
+    if cp_size <= 1:
+        return tensor
+    seq_len = tensor.shape[seq_dim]
+    if seq_len % cp_size != 0:
+        raise ValueError(f"seq_len={seq_len} must be divisible by cp_size={cp_size}")
+    local_len = seq_len // cp_size
+    return tensor.narrow(seq_dim, cp_rank * local_len, local_len).contiguous()
+
+
+def contiguous_position_ids_for_cp(
+    seq_len: int,
+    cp_rank: int,
+    cp_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Return global position IDs for this CP rank under contiguous splitting."""
+    if cp_size <= 1:
+        return torch.arange(seq_len, device=device).unsqueeze(0)
+    if seq_len % cp_size != 0:
+        raise ValueError(f"seq_len={seq_len} must be divisible by cp_size={cp_size}")
+    local_len = seq_len // cp_size
+    start = cp_rank * local_len
+    return torch.arange(start, start + local_len, device=device).unsqueeze(0)
+
+
+def local_position_ids_for_cp(position_ids, *, batch, local_seq_len, cp_rank, cp_size):
+    """Validate and contiguous-slice full-length position_ids to this CP rank."""
+    if position_ids.dim() == 1:
+        position_ids = position_ids.unsqueeze(0)
+    if position_ids.dim() != 2:
+        raise ValueError("position_ids must have shape (S,) or (B, S).")
+    if position_ids.size(0) == 1 and batch > 1:
+        position_ids = position_ids.expand(batch, -1)
+    if position_ids.size(0) != batch:
+        raise ValueError(
+            f"position_ids batch={position_ids.size(0)} does not match input batch={batch}."
+        )
+    if cp_size <= 1 or position_ids.size(1) == local_seq_len:
+        return position_ids
+
+    full_seq_len = local_seq_len * cp_size
+    if position_ids.size(1) != full_seq_len:
+        raise ValueError(
+            "CP expects position_ids to be either CP-local or full-length; "
+            f"got {position_ids.size(1)} for local_seq_len={local_seq_len}, cp={cp_size}."
+        )
+    return contiguous_slice_for_cp(position_ids, cp_rank, cp_size, seq_dim=1)
+
+
+def local_sequence_tensor_for_cp(
+    tensor,
+    *,
+    local_seq_len,
+    cp_rank,
+    cp_size,
+    seq_dim=1,
+    name: str = "tensor",
+    unsqueeze_1d: bool = True,
+):
+    """Validate and contiguous-slice a full-length sequence tensor to this CP rank."""
+    if tensor is None or cp_size <= 1:
+        return tensor
+    if unsqueeze_1d and tensor.dim() == 1:
+        tensor = tensor.unsqueeze(0)
+    full_seq_len = local_seq_len * cp_size
+    seq_len = tensor.size(seq_dim)
+    if seq_len == local_seq_len:
+        return tensor
+    if seq_len != full_seq_len:
+        raise ValueError(
+            f"CP expects {name} to be either CP-local or full-length; "
+            f"got {seq_len} for local_seq_len={local_seq_len}, cp={cp_size}."
+        )
+    return contiguous_slice_for_cp(tensor, cp_rank, cp_size, seq_dim=seq_dim)
+
+
+def zigzag_to_contiguous_chunks(
+    tensor: torch.Tensor,
+    cp_group: dist.ProcessGroup | None,
+    seq_dim: int = 1,
+) -> torch.Tensor:
+    """Swap a CP-local tensor from Megatron zigzag layout to contiguous chunks.
+
+    Zigzag CP layout assigns rank ``r`` global chunks ``[r, 2*cp-r-1]``.
+    Linear-attention all-gather CP kernels expect rank ``r`` to hold chunks
+    ``[2*r, 2*r+1]``. The conversion is a chunk-level all-to-all and preserves
+    the local tensor shape.
+    """
+    return _zigzag_contiguous_chunk_swap(tensor, cp_group, seq_dim, to_contiguous=True)
+
+
+def contiguous_to_zigzag_chunks(
+    tensor: torch.Tensor,
+    cp_group: dist.ProcessGroup | None,
+    seq_dim: int = 1,
+) -> torch.Tensor:
+    """Inverse of :func:`zigzag_to_contiguous_chunks`."""
+    return _zigzag_contiguous_chunk_swap(tensor, cp_group, seq_dim, to_contiguous=False)
+
+
+def _zigzag_contiguous_chunk_swap(
+    tensor: torch.Tensor,
+    cp_group: Optional[dist.ProcessGroup],
+    seq_dim: int,
+    *,
+    to_contiguous: bool,
+) -> torch.Tensor:
+    cp_size = dist.get_world_size(cp_group) if cp_group is not None else 1
+    if cp_size <= 1:
+        return tensor
+    cp_rank = dist.get_rank(cp_group)
+
+    if seq_dim != 0:
+        tensor = tensor.movedim(seq_dim, 0)
+    tensor = tensor.contiguous()
+
+    local_len = tensor.size(0)
+    if local_len % 2 != 0:
+        raise ValueError(
+            f"zigzag/contiguous CP chunk swap requires even local sequence length, got {local_len}."
+        )
+    chunk_len = local_len // 2
+
+    def rank_to_chunks(rank: int, in_zigzag: bool) -> tuple[int, int]:
+        if in_zigzag:
+            return rank, 2 * cp_size - rank - 1
+        return 2 * rank, 2 * rank + 1
+
+    def chunk_to_dest(chunk_idx: int, target_zigzag: bool) -> tuple[int, int]:
+        if target_zigzag:
+            if chunk_idx < cp_size:
+                return chunk_idx, 0
+            return 2 * cp_size - chunk_idx - 1, 1
+        return chunk_idx // 2, chunk_idx % 2
+
+    source_in_zigzag = to_contiguous
+    target_in_zigzag = not to_contiguous
+    local_chunks = [tensor[:chunk_len], tensor[chunk_len:]]
+    local_chunk_indices = rank_to_chunks(cp_rank, source_in_zigzag)
+    local_dests = [chunk_to_dest(chunk_idx, target_in_zigzag) for chunk_idx in local_chunk_indices]
+    local_slot_order = sorted(range(2), key=lambda slot: local_dests[slot])
+    send_buf = torch.cat([local_chunks[slot] for slot in local_slot_order], dim=0).contiguous()
+
+    input_split_chunks = [0] * cp_size
+    for dst_rank, _dst_slot in local_dests:
+        input_split_chunks[dst_rank] += 1
+
+    output_split_chunks = [0] * cp_size
+    recv_dst_slots_per_source: list[list[int]] = [[] for _ in range(cp_size)]
+    for src_rank in range(cp_size):
+        src_chunks = rank_to_chunks(src_rank, source_in_zigzag)
+        src_dests = [chunk_to_dest(chunk_idx, target_in_zigzag) for chunk_idx in src_chunks]
+        src_slot_order = sorted(range(2), key=lambda slot: src_dests[slot])
+        for slot in src_slot_order:
+            dst_rank, dst_slot = src_dests[slot]
+            if dst_rank == cp_rank:
+                output_split_chunks[src_rank] += 1
+                recv_dst_slots_per_source[src_rank].append(dst_slot)
+
+    input_split_sizes = [count * chunk_len for count in input_split_chunks]
+    output_split_sizes = [count * chunk_len for count in output_split_chunks]
+    recv_shape = (sum(output_split_sizes), *send_buf.shape[1:])
+    recv_buf = torch.empty(recv_shape, dtype=send_buf.dtype, device=send_buf.device)
+    from torch.distributed.nn.functional import all_to_all_single
+
+    recv_buf = all_to_all_single(
+        recv_buf,
+        send_buf,
+        output_split_sizes=output_split_sizes,
+        input_split_sizes=input_split_sizes,
+        group=cp_group,
+    )
+
+    target_slots: list[torch.Tensor | None] = [None, None]
+    offset = 0
+    for src_rank in range(cp_size):
+        for dst_slot in recv_dst_slots_per_source[src_rank]:
+            target_slots[dst_slot] = recv_buf[offset : offset + chunk_len]
+            offset += chunk_len
+    if any(slot is None for slot in target_slots):
+        raise RuntimeError("Incomplete CP chunk reassembly.")
+
+    out = torch.cat([slot for slot in target_slots if slot is not None], dim=0)
+    if seq_dim != 0:
+        out = out.movedim(0, seq_dim)
+    return out.contiguous()
+
+
 def zigzag_position_ids_for_cp(
-    seq_len: int, cp_rank: int, cp_size: int, device: torch.device
+    seq_len: int,
+    cp_rank: int,
+    cp_size: int,
+    device: torch.device,
 ) -> torch.Tensor:
     """Return global position IDs for this CP rank under zigzag splitting.
 
@@ -146,7 +351,11 @@ def split_packed_for_cp(
 
 
 __all__ = [
+    "contiguous_position_ids_for_cp",
+    "contiguous_slice_for_cp",
+    "contiguous_to_zigzag_chunks",
     "split_packed_for_cp",
+    "zigzag_to_contiguous_chunks",
     "zigzag_reconstruct_from_cp_parts",
     "zigzag_position_ids_for_cp",
     "zigzag_slice_for_cp",
diff --git a/experimental/lite/megatron/lite/primitive/parallel/mhc.py b/experimental/lite/megatron/lite/primitive/parallel/mhc.py
new file mode 100644
index 00000000000..9a48022a300
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/parallel/mhc.py
@@ -0,0 +1,52 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+from typing import Any
+
+import torch
+
+
+def expand_mhc_hidden_for_pipeline(hidden: torch.Tensor, *, hc_mult: int) -> torch.Tensor:
+    if hidden.dim() == 3:
+        return hidden.unsqueeze(2).expand(-1, -1, hc_mult, -1).contiguous()
+    return hidden
+
+
+def fold_mhc_hidden_for_pipeline(hidden: torch.Tensor) -> torch.Tensor:
+    """Fold the 4-D hc streams [B, S, hc_mult, H] into [B, S, hc_mult * H] for PP P2P.
+
+    The pipeline P2P recv buffer is 3-D, so the hc_mult parallel residual streams must be
+    flattened into the hidden dimension before crossing a stage boundary. No-op for 3-D input.
+    """
+    if hidden.dim() == 4:
+        b, s, m, h = hidden.shape
+        return hidden.reshape(b, s, m * h).contiguous()
+    return hidden
+
+
+def unfold_mhc_hidden_from_pipeline(hidden: torch.Tensor, *, hc_mult: int) -> torch.Tensor:
+    """Inverse of :func:`fold_mhc_hidden_for_pipeline`: [B, S, hc_mult * H] -> [B, S, hc_mult, H].
+
+    Used on non-first PP stages to restore the hc_mult streams received over P2P. No-op when the
+    tensor is already 4-D.
+    """
+    if hidden.dim() == 4:
+        return hidden
+    b, s, mh = hidden.shape
+    return hidden.reshape(b, s, hc_mult, mh // hc_mult).contiguous()
+
+
+def contract_mhc_hidden_for_pipeline(
+    hidden: torch.Tensor,
+    *,
+    norm: Any,
+    head: Any,
+    return_source: bool = False,
+):
+    if head is None or norm is None:
+        if return_source:
+            return hidden, None
+        return hidden
+    source = hidden
+    contracted = norm(head(hidden))
+    if return_source:
+        return contracted, source
+    return contracted
diff --git a/experimental/lite/megatron/lite/primitive/parallel/pipeline.py b/experimental/lite/megatron/lite/primitive/parallel/pipeline.py
index d52271862a8..5cee84b543a 100644
--- a/experimental/lite/megatron/lite/primitive/parallel/pipeline.py
+++ b/experimental/lite/megatron/lite/primitive/parallel/pipeline.py
@@ -11,6 +11,7 @@
 import torch.distributed as dist  # pyright: ignore[reportMissingImports]
 
 from megatron.lite.primitive.utils import ensure_divisible
+from megatron.lite.runtime.contracts.loss import split_loss_context, use_loss_context
 
 if TYPE_CHECKING:
     from megatron.lite.primitive.parallel.state import ParallelState
@@ -133,19 +134,16 @@ def _batch_get(batch, key: str):
     return getattr(batch, key, None)
 
 
-def _batch_input_ids(batch):
-    input_ids = _batch_get(batch, "input_ids")
-    if input_ids is not None and input_ids.dim() == 1:
-        input_ids = input_ids.unsqueeze(0)
-    return input_ids
-
-
 def _apply_external_loss(
-    out: dict, batch, loss_fn
+    out: dict, batch, loss_fn, loss_context=None
 ) -> tuple[torch.Tensor, dict] | tuple[None, None]:
     if loss_fn is None:
         return None, None
-    loss, metrics = loss_fn(out, batch)
+    # Mirror run_microbatch_loop: pass loss_context as 3rd arg when present.
+    if loss_context is None:
+        loss, metrics = loss_fn(out, batch)
+    else:
+        loss, metrics = loss_fn(out, batch, loss_context)
     out["loss"] = loss
     out["_loss_fn_metrics"] = metrics
     return loss, metrics
@@ -155,15 +153,15 @@ def _compact_pipeline_output(out: dict | None) -> dict:
     if not out:
         return {}
     compact: dict = {}
-    if "_verl_model_output" in out:
-        compact["model_output"] = out["_verl_model_output"]
+    if "model_output" in out:
+        compact["model_output"] = out["model_output"]
     if "loss" in out and out["loss"] is not None:
         loss = out["loss"]
         compact["loss"] = loss.detach().item() if isinstance(loss, torch.Tensor) else float(loss)
     if "_loss_fn_metrics" in out:
         compact["metrics"] = out["_loss_fn_metrics"]
-    elif "_verl_metrics" in out:
-        compact["metrics"] = out["_verl_metrics"]
+    elif "metrics" in out:
+        compact["metrics"] = out["metrics"]
     return compact
 
 
@@ -236,7 +234,9 @@ def _1f1b_schedule(
     num_warmup = min(ps.pp_size - ps.pp_rank - 1, num_microbatches)
     num_steady = num_microbatches - num_warmup
 
-    batches = [next(data_iter) for _ in range(num_microbatches)]
+    # Split each microbatch into (PackedBatch, LossContext) like run_microbatch_loop; the connector
+    # yields (batch, loss_context) tuples, so forward_step must receive the unwrapped batch.
+    batches = [split_loss_context(next(data_iter)) for _ in range(num_microbatches)]
     mb_idx = 0
 
     input_tensors: list[torch.Tensor | None] = []
@@ -256,31 +256,18 @@ def _1f1b_schedule(
         else None
     )
 
-    def _run_forward(input_tensor, batch):
+    def _run_forward(input_tensor, batch, loss_context=None):
         _set_aux_loss_scale(pre_forward_hook, num_microbatches)
-        position_ids = _batch_get(batch, "position_ids")
-        packed_seq_params = _batch_get(batch, "packed_seq_params")
-        if ps.pp_is_first:
-            return forward_step_fn(model, batch)
-        if ps.pp_is_last:
-            out = model(
-                input_ids=_batch_input_ids(batch),
-                hidden_states=input_tensor,
-                position_ids=position_ids,
-                packed_seq_params=packed_seq_params,
-                labels=_batch_get(batch, "labels"),
-                loss_mask=_batch_get(batch, "loss_mask"),
-                temperature=_batch_get(batch, "temperature") or 1.0,
-                use_fused_kernels=bool(_batch_get(batch, "use_fused_kernels") or False),
-                calculate_entropy=bool(_batch_get(batch, "calculate_entropy") or False),
-            )
-            _apply_external_loss(out, batch, loss_fn)
-            return out
-        return model(
-            hidden_states=input_tensor,
-            position_ids=position_ids,
-            packed_seq_params=packed_seq_params,
-        )
+        if not ps.pp_is_first:
+            # `model` is the dist_opt DDP-wrapped chunk; set_input_tensor lives on the base lite model.
+            from megatron.lite.primitive.ckpt.hf_weights import unwrap_model
+
+            unwrap_model(model).set_input_tensor(input_tensor)
+        with use_loss_context(loss_context):
+            out = forward_step_fn(model, batch)
+            if ps.pp_is_last:
+                _apply_external_loss(out, batch, loss_fn, loss_context)
+        return out
 
     def _run_backward(inp_t, hid_t, loss_t, grad_t):
         if ps.pp_is_last:
@@ -309,10 +296,10 @@ def _p2p(send_fwd=None, send_bwd=None, recv_fwd=False, recv_bwd=False):
         if not ps.pp_is_first and k == 0:
             fwd_input, _ = _p2p(recv_fwd=True)
 
-        batch = batches[mb_idx]
+        batch, loss_ctx = batches[mb_idx]
         mb_idx += 1
         current_input = fwd_input
-        out = _run_forward(fwd_input, batch)
+        out = _run_forward(fwd_input, batch, loss_ctx)
         hidden = out.get("hidden_states")
         loss_s = out["loss"] / num_microbatches if "loss" in out and ps.pp_is_last else None
 
@@ -336,9 +323,9 @@ def _p2p(send_fwd=None, send_bwd=None, recv_fwd=False, recv_bwd=False):
         if not ps.pp_is_first and k == 0 and num_warmup == 0:
             fwd_input, _ = _p2p(recv_fwd=True)
 
-        batch = batches[mb_idx]
+        batch, loss_ctx = batches[mb_idx]
         mb_idx += 1
-        out = _run_forward(fwd_input, batch)
+        out = _run_forward(fwd_input, batch, loss_ctx)
         hidden = out.get("hidden_states")
         loss_s = out["loss"] / num_microbatches if "loss" in out and ps.pp_is_last else None
 
@@ -502,18 +489,13 @@ def _pipeline_stage_barrier(ps: ParallelState) -> None:
         dist.barrier(group=ps.pp_cpu_group)
 
 
-def _set_virtual_pipeline_rank(chunk_id: int | None, num_chunks: int) -> None:
+def _set_virtual_pipeline_rank(ps: ParallelState, chunk_id: int | None, num_chunks: int) -> None:
     if chunk_id is None or num_chunks <= 1:
+        ps.virtual_pipeline_size = None
+        ps.virtual_pipeline_rank = None
         return
-    try:
-        from megatron.core import parallel_state as mpu  # pyright: ignore[reportMissingImports]
-    except Exception:
-        return
-    if not mpu.is_initialized():
-        return
-    vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()
-    if vpp_size is not None and vpp_size > 1:
-        mpu.set_virtual_pipeline_model_parallel_rank(chunk_id)
+    ps.virtual_pipeline_size = num_chunks
+    ps.virtual_pipeline_rank = chunk_id
 
 
 def _run_pipeline_chunk_forward(
@@ -529,27 +511,12 @@ def _run_pipeline_chunk_forward(
     loss_fn=None,
 ) -> dict:
     _set_aux_loss_scale(pre_forward_hook, num_microbatches)
-    position_ids = _batch_get(batch, "position_ids")
-    packed_seq_params = _batch_get(batch, "packed_seq_params")
-    if is_first_stage:
-        return forward_step_fn(model, batch)
+    if not is_first_stage:
+        model.set_input_tensor(input_tensor)
+    out = forward_step_fn(model, batch)
     if is_last_stage:
-        out = model(
-            input_ids=_batch_input_ids(batch),
-            hidden_states=input_tensor,
-            position_ids=position_ids,
-            packed_seq_params=packed_seq_params,
-            labels=_batch_get(batch, "labels"),
-            loss_mask=_batch_get(batch, "loss_mask"),
-            temperature=_batch_get(batch, "temperature") or 1.0,
-            use_fused_kernels=bool(_batch_get(batch, "use_fused_kernels") or False),
-            calculate_entropy=bool(_batch_get(batch, "calculate_entropy") or False),
-        )
         _apply_external_loss(out, batch, loss_fn)
-        return out
-    return model(
-        hidden_states=input_tensor, position_ids=position_ids, packed_seq_params=packed_seq_params
-    )
+    return out
 
 
 def _forward_only_pipeline_schedule(
@@ -579,7 +546,7 @@ def _forward_only_pipeline_schedule(
             hidden: torch.Tensor | None = None
             if is_local_stage:
                 chunk_id = stage_id // ps.pp_size
-                _set_virtual_pipeline_rank(chunk_id, num_chunks)
+                _set_virtual_pipeline_rank(ps, chunk_id, num_chunks)
                 model = model_chunks[chunk_id]
                 is_first_stage = stage_id == 0
                 is_last_stage = stage_id == total_stages - 1
@@ -619,6 +586,7 @@ def _forward_only_pipeline_schedule(
 
         outputs.append(_compact_pipeline_output(last_output) if last_output is not None else {})
 
+    _set_virtual_pipeline_rank(ps, None, num_chunks)
     return outputs
 
 
@@ -675,7 +643,7 @@ def _interleaved_1f1b_schedule(
             hidden: torch.Tensor | None = None
             if is_local_stage:
                 chunk_id = stage_id // ps.pp_size
-                _set_virtual_pipeline_rank(chunk_id, num_chunks)
+                _set_virtual_pipeline_rank(ps, chunk_id, num_chunks)
                 model = model_chunks[chunk_id]
                 is_first_stage = stage_id == 0
                 is_last_stage = stage_id == total_stages - 1
@@ -747,7 +715,7 @@ def _interleaved_1f1b_schedule(
             inp_grad: torch.Tensor | None = None
             if is_local_stage:
                 chunk_id = stage_id // ps.pp_size
-                _set_virtual_pipeline_rank(chunk_id, num_chunks)
+                _set_virtual_pipeline_rank(ps, chunk_id, num_chunks)
                 is_first_stage = stage_id == 0
                 is_last_stage = stage_id == total_stages - 1
                 inp, out_t, loss, _out = saved[stage_id]
@@ -788,6 +756,7 @@ def _interleaved_1f1b_schedule(
         if _dbg:
             print(f"[VPP r{rank}] mb={mb_id} complete", flush=True)
 
+    _set_virtual_pipeline_rank(ps, None, num_chunks)
     return outputs
 
 
diff --git a/experimental/lite/megatron/lite/primitive/parallel/pp.py b/experimental/lite/megatron/lite/primitive/parallel/pp.py
index e5da34b8a37..61077f07d56 100644
--- a/experimental/lite/megatron/lite/primitive/parallel/pp.py
+++ b/experimental/lite/megatron/lite/primitive/parallel/pp.py
@@ -1,13 +1,21 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-"""Context/pipeline parallel sequence splitting utilities."""
+"""Pipeline parallel layer layout — a thin wrapper over Megatron-core's
+``PipelineParallelLayerLayout`` (authorized mcore-reuse), which owns the per-stage
+split for non-divisible decoder counts plus MTP. Two pp-only modes:
+
+* **auto** (default, only ``pp`` set): balance ``[E, decoder*N, mtp*K, loss]`` across stages.
+* **custom** (``ParallelConfig.pp_layout``): an explicit mcore layout string/list,
+  e.g. ``"E|t*5|t*6|t,m,L"``.
+
+Not supported (raise, never mis-place): VPP, and standalone MTP (``m`` off the
+final/loss stage — mlite's MTP shares the head there; cross-stage MTP is a follow-up).
+"""
 
 from __future__ import annotations
 
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
-from megatron.lite.primitive.utils import ensure_divisible
-
 if TYPE_CHECKING:
     from megatron.lite.primitive.parallel.state import ParallelState
 
@@ -17,6 +25,24 @@ class PipelineChunkLayout:
     layer_indices: list[int] = field(default_factory=list)
     has_embed: bool = False
     has_head: bool = False
+    has_mtp: bool = False
+
+
+def _auto_layout(num_hidden_layers: int, pp_size: int, num_mtp_layers: int):
+    """Balance ``[E, decoder*N, mtp*K, loss]`` into even contiguous chunks; the
+    embedding/MTP/loss slots make their stages carry fewer decoders (e.g. 6/pp4 ->
+    [1,2,2,1]) — Megatron's embedding/loss split accounting."""
+    from megatron.core.transformer.pipeline_parallel_layer_layout import (
+        PipelineParallelLayerLayout,
+    )
+
+    units = ["embedding"] + ["decoder"] * num_hidden_layers + ["mtp"] * max(num_mtp_layers, 0) + ["loss"]
+    base, remainder = divmod(len(units), pp_size)
+    rows, pos = [], 0
+    for size in (base + (1 if s < remainder else 0) for s in range(pp_size)):
+        rows.append(units[pos : pos + size])
+        pos += size
+    return PipelineParallelLayerLayout(rows, pipeline_model_parallel_size=pp_size)
 
 
 def build_pipeline_chunk_layout(
@@ -24,30 +50,50 @@ def build_pipeline_chunk_layout(
     ps: ParallelState,
     vpp: int | None = None,
     vpp_chunk_id: int | None = None,
+    *,
+    num_mtp_layers: int = 0,
 ) -> PipelineChunkLayout:
-    """Compute layer_indices, has_embed, has_head for this PP rank / VPP chunk."""
-    if vpp_chunk_id is not None:
-        assert vpp is not None
-        layers_per_chunk = ensure_divisible(num_hidden_layers, ps.pp_size * vpp)
-        start = ps.pp_rank * layers_per_chunk + vpp_chunk_id * (ps.pp_size * layers_per_chunk)
-        layer_indices = list(range(start, start + layers_per_chunk))
-        has_embed = ps.pp_is_first and vpp_chunk_id == 0
-        has_head = ps.pp_is_last and vpp_chunk_id == vpp - 1
-    elif vpp is not None:
-        layers_per_chunk = ensure_divisible(num_hidden_layers, ps.pp_size * vpp)
-        layer_indices = []
-        for chunk in range(vpp):
-            start = ps.pp_rank * layers_per_chunk + chunk * (ps.pp_size * layers_per_chunk)
-            layer_indices.extend(range(start, start + layers_per_chunk))
-        has_embed = ps.pp_is_first
-        has_head = ps.pp_is_last
+    """``layer_indices`` / ``has_embed`` / ``has_head`` / ``has_mtp`` for this PP rank,
+    from ``ps.pp_layout`` (custom) or an auto-balanced layout. ``has_mtp`` follows the
+    layout's ``m`` placement, so MTP is built where the layout says, not a fixed rank."""
+    if (vpp is not None and vpp > 1) or vpp_chunk_id is not None:
+        raise NotImplementedError("VPP / interleaved pipeline layout is not supported (use vpp=1).")
+
+    if ps.pp_size <= 1:  # no pipeline: this stage owns everything
+        return PipelineChunkLayout(
+            layer_indices=list(range(num_hidden_layers)),
+            has_embed=True,
+            has_head=True,
+            has_mtp=num_mtp_layers > 0,
+        )
+
+    from megatron.core.transformer.enums import LayerType
+    from megatron.core.transformer.pipeline_parallel_layer_layout import (
+        PipelineParallelLayerLayout,
+    )
+
+    pp_layout = getattr(ps, "pp_layout", None)
+    if pp_layout is not None:
+        layout = PipelineParallelLayerLayout(pp_layout, pipeline_model_parallel_size=ps.pp_size)
+        if layout.virtual_pipeline_model_parallel_size > 1:
+            raise NotImplementedError("VPP pp_layout is not supported (one stage per pp rank).")
     else:
-        layers_per_stage = ensure_divisible(num_hidden_layers, ps.pp_size)
-        start = ps.pp_rank * layers_per_stage
-        layer_indices = list(range(start, start + layers_per_stage))
-        has_embed = ps.pp_is_first
-        has_head = ps.pp_is_last
-    return PipelineChunkLayout(layer_indices=layer_indices, has_embed=has_embed, has_head=has_head)
+        layout = _auto_layout(num_hidden_layers, ps.pp_size, num_mtp_layers)
+
+    # validate_layer_layout checks legality and returns mtp_standalone=True when `m` is
+    # off the final stage — which mlite's head-coupled MTP cannot run.
+    if layout.validate_layer_layout(num_hidden_layers, num_mtp_layers or None):
+        raise NotImplementedError(
+            "Standalone MTP (pp_layout with `m` off the final/loss stage) is not "
+            "implemented yet — mlite's MTP shares the output head there. Use the auto "
+            "layout (set only `pp`), or place `m` on the same stage as `L`."
+        )
+    return PipelineChunkLayout(
+        layer_indices=layout.get_layer_id_list(LayerType.decoder, vp_stage=0, pp_rank=ps.pp_rank),
+        has_embed=ps.pp_is_first,
+        has_head=ps.pp_is_last,
+        has_mtp=layout.get_num_layers_to_build(LayerType.mtp, vp_stage=0, pp_rank=ps.pp_rank) > 0,
+    )
 
 
 __all__ = ["PipelineChunkLayout", "build_pipeline_chunk_layout"]
diff --git a/experimental/lite/megatron/lite/primitive/parallel/state.py b/experimental/lite/megatron/lite/primitive/parallel/state.py
index c282364a714..93481126953 100644
--- a/experimental/lite/megatron/lite/primitive/parallel/state.py
+++ b/experimental/lite/megatron/lite/primitive/parallel/state.py
@@ -47,6 +47,11 @@ class ParallelState:
     pp_is_last: bool = True
     pp_next_rank: int = -1
     pp_prev_rank: int = -1
+    virtual_pipeline_size: int | None = None
+    virtual_pipeline_rank: int | None = None
+
+    # Optional explicit mcore pipeline layout (custom mode); None -> auto-infer.
+    pp_layout: str | list | None = None
 
 
 def init_parallel(config) -> ParallelState:
@@ -66,6 +71,7 @@ def init_parallel(config) -> ParallelState:
     expert_dp = ensure_divisible(world, etp * ep * pp)
 
     ps = ParallelState()
+    ps.pp_layout = getattr(config, "pp_layout", None)
     ps.tp_size, ps.ep_size, ps.etp_size = tp, ep, etp
     ps.cp_size, ps.pp_size, ps.dp_size = cp, pp, dense_dp
     ps.expert_dp_size = expert_dp
diff --git a/experimental/lite/megatron/lite/primitive/parallel/thd.py b/experimental/lite/megatron/lite/primitive/parallel/thd.py
index ac928bd8b5b..65a61bc0e01 100644
--- a/experimental/lite/megatron/lite/primitive/parallel/thd.py
+++ b/experimental/lite/megatron/lite/primitive/parallel/thd.py
@@ -3,48 +3,14 @@
 
 from __future__ import annotations
 
+from copy import copy
 from dataclasses import dataclass
 from typing import Any
 
 import torch
 import torch.distributed as dist
 
-
-@dataclass
-class PackedSeqParams:
-    """Parameters for THD-format packed-sequence attention.
-
-    Mirrors the fields consumed by TE's ``DotProductAttention.forward()``
-    when ``qkv_format="thd"`` is used (total-tokens x heads x dim).
-
-    Typical construction::
-
-        params = PackedSeqParams.from_cu_seqlens(batch.cu_seqlens, batch.max_seqlen)
-    """
-
-    qkv_format: str = "thd"
-    cu_seqlens_q: torch.Tensor | None = None
-    cu_seqlens_kv: torch.Tensor | None = None
-    max_seqlen_q: int | None = None
-    max_seqlen_kv: int | None = None
-    cu_seqlens_q_padded: torch.Tensor | None = None
-    cu_seqlens_kv_padded: torch.Tensor | None = None
-    local_cp_size: int | None = None
-    cp_group: Any | None = None
-    cp_rank: int | None = None
-
-    @staticmethod
-    def from_cu_seqlens(cu_seqlens: torch.Tensor, max_seqlen: int) -> PackedSeqParams:
-        """Build from shared Q/KV cu_seqlens (self-attention)."""
-        return PackedSeqParams(
-            qkv_format="thd",
-            cu_seqlens_q=cu_seqlens,
-            cu_seqlens_kv=cu_seqlens,
-            max_seqlen_q=max_seqlen,
-            max_seqlen_kv=max_seqlen,
-            cu_seqlens_q_padded=cu_seqlens,
-            cu_seqlens_kv_padded=cu_seqlens,
-        )
+from megatron.lite.primitive.utils.packed_seq import PackedSeqParams
 
 
 @dataclass(frozen=True)
@@ -77,34 +43,17 @@ def _make_packed_seq_params(
         extra_args["local_cp_size"] = cp_size
         if cp_group is not None:
             extra_args["cp_group"] = cp_group
-    try:
-        from megatron.core.packed_seq_params import PackedSeqParams as MCorePackedSeqParams
-
-        params = MCorePackedSeqParams(
-            qkv_format="thd",
-            cu_seqlens_q=cu_seqlens_padded,
-            cu_seqlens_kv=cu_seqlens_padded,
-            max_seqlen_q=max_seqlen,
-            max_seqlen_kv=max_seqlen,
-            cu_seqlens_q_padded=cu_seqlens_padded,
-            cu_seqlens_kv_padded=cu_seqlens_padded,
-            **extra_args,
-        )
-        # MCore's PackedSeqParams does not carry rank, but local rolling needs it.
-        params.cp_rank = cp_rank
-        return params
-    except Exception:
-        return PackedSeqParams(
-            qkv_format="thd",
-            cu_seqlens_q=cu_seqlens_padded,
-            cu_seqlens_kv=cu_seqlens_padded,
-            max_seqlen_q=max_seqlen,
-            max_seqlen_kv=max_seqlen,
-            cu_seqlens_q_padded=cu_seqlens_padded,
-            cu_seqlens_kv_padded=cu_seqlens_padded,
-            cp_rank=cp_rank,
-            **extra_args,
-        )
+    return PackedSeqParams(
+        qkv_format="thd",
+        cu_seqlens_q=cu_seqlens_padded,
+        cu_seqlens_kv=cu_seqlens_padded,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_kv=max_seqlen,
+        cu_seqlens_q_padded=cu_seqlens_padded,
+        cu_seqlens_kv_padded=cu_seqlens_padded,
+        cp_rank=cp_rank,
+        **extra_args,
+    )
 
 
 def _slice_along_dim(tensor: torch.Tensor, dim: int, start: int, end: int) -> torch.Tensor:
@@ -198,6 +147,188 @@ def split_packed_to_cp_local(
     )
 
 
+def _packed_cu_seqlens(packed_seq_params: Any) -> torch.Tensor | None:
+    if packed_seq_params is None:
+        return None
+    cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q_padded", None)
+    if cu_seqlens is None:
+        cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q", None)
+    return cu_seqlens
+
+
+def has_packed_thd_params(packed_seq_params: Any) -> bool:
+    return _packed_cu_seqlens(packed_seq_params) is not None
+
+
+def _sequence_dim(tensor: torch.Tensor) -> int:
+    return tensor.dim() - 1 if tensor.dim() > 1 else 0
+
+
+def _with_cp_metadata(packed_seq_params: Any, *, cp_size: int, cp_rank: int, cp_group: Any):
+    updated = copy(packed_seq_params)
+    updated.local_cp_size = cp_size
+    updated.cp_rank = cp_rank
+    updated.cp_group = cp_group
+    return updated
+
+
+def parallel_state_from_model(model: Any) -> Any:
+    """Return the MLite parallel state from a raw model or wrapper."""
+
+    current = model
+    seen: set[int] = set()
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        ps = getattr(current, "ps", None)
+        if ps is not None:
+            return ps
+        ps = getattr(current, "_parallel_state", None)
+        if ps is not None:
+            return ps
+        current = getattr(current, "module", None)
+    return None
+
+
+def prepare_packed_thd_for_context_parallel(
+    packed_seq_params: Any,
+    tensors: tuple[torch.Tensor | None, ...],
+    *,
+    cp_size: int,
+    cp_rank: int,
+    cp_group: Any = None,
+    dims: tuple[int | None, ...] | None = None,
+) -> tuple[Any, tuple[torch.Tensor | None, ...]]:
+    """Split plain packed THD tensors to one CP rank without knowing batch keys."""
+
+    tensor_tuple = tuple(tensors)
+    cu_seqlens = _packed_cu_seqlens(packed_seq_params)
+    if cu_seqlens is None or cp_size <= 1:
+        return packed_seq_params, tensor_tuple
+    if int(getattr(packed_seq_params, "local_cp_size", None) or 1) > 1:
+        return packed_seq_params, tensor_tuple
+    if dims is not None and len(dims) != len(tensor_tuple):
+        raise ValueError(
+            f"dims length {len(dims)} does not match tensors length {len(tensor_tuple)}."
+        )
+
+    local_tensors: list[torch.Tensor | None] = []
+    for idx, tensor in enumerate(tensor_tuple):
+        if tensor is None:
+            local_tensors.append(None)
+            continue
+        dim = _sequence_dim(tensor) if dims is None or dims[idx] is None else int(dims[idx])
+        local_tensors.append(
+            _split_full_to_cp_local(
+                tensor,
+                cu_seqlens_padded=cu_seqlens,
+                cp_size=cp_size,
+                cp_rank=cp_rank,
+                dim=dim,
+            )
+        )
+    return (
+        _with_cp_metadata(packed_seq_params, cp_size=cp_size, cp_rank=cp_rank, cp_group=cp_group),
+        tuple(local_tensors),
+    )
+
+
+def prepare_packed_thd_kwargs_for_context_parallel(
+    model: Any,
+    kwargs: dict[str, Any],
+    *,
+    tensor_keys: tuple[str, ...] = ("input_ids", "labels", "loss_mask", "position_ids"),
+) -> None:
+    packed_seq_params = kwargs.get("packed_seq_params")
+    if not has_packed_thd_params(packed_seq_params):
+        return
+
+    ps = parallel_state_from_model(model)
+    packed_seq_params, tensors = prepare_packed_thd_for_context_parallel(
+        packed_seq_params,
+        tuple(kwargs.get(key) for key in tensor_keys),
+        cp_size=int(getattr(ps, "cp_size", 1) or 1),
+        cp_rank=int(getattr(ps, "cp_rank", 0) or 0),
+        cp_group=getattr(ps, "cp_group", None),
+    )
+    if packed_seq_params is not None or "packed_seq_params" in kwargs:
+        kwargs["packed_seq_params"] = packed_seq_params
+    for key, tensor in zip(tensor_keys, tensors, strict=True):
+        if tensor is not None or key in kwargs:
+            kwargs[key] = tensor
+
+
+@dataclass(frozen=True)
+class ThdPackMeta:
+    """Per-sequence THD padding layout, recomputable from true seq lengths.
+
+    Shared by a model's pack/unpack pair so the connector never needs to hold
+    ``PackedSeqParams`` or padded token tensors to reverse a model output.
+    """
+
+    lengths: torch.Tensor
+    padded_lengths: torch.Tensor
+    cu_seqlens_padded: torch.Tensor
+    cp_size: int
+    cp_group: Any | None
+
+
+def thd_pack_meta(
+    seq_lens: torch.Tensor,
+    *,
+    tp_size: int = 1,
+    cp_size: int = 1,
+    cp_group: Any | None = None,
+    contiguous: bool = False,
+) -> ThdPackMeta:
+    """Compute the padded THD layout for ``seq_lens`` without copying tokens.
+
+    ``contiguous=False`` aligns to Megatron/TE zigzag CP (``2*cp``); ``True``
+    aligns to contiguous CP (``cp``). Mirrors :func:`pack_nested_thd` padding.
+    """
+    lengths = seq_lens.to(dtype=torch.int32)
+    cp_align = (cp_size if contiguous else 2 * cp_size) if cp_size > 1 else 1
+    align_size = max(int(tp_size), 1) * cp_align
+    pad_size = (align_size - lengths % align_size) % align_size
+    padded_lengths = lengths + pad_size
+    cu_seqlens_padded = torch.zeros(
+        lengths.numel() + 1, dtype=torch.int32, device=lengths.device
+    )
+    cu_seqlens_padded[1:] = torch.cumsum(padded_lengths, dim=0)
+    return ThdPackMeta(lengths, padded_lengths, cu_seqlens_padded, cp_size, cp_group)
+
+
+def unpack_thd_to_nested(
+    output: torch.Tensor, meta: ThdPackMeta, *, contiguous: bool = False
+) -> torch.Tensor:
+    """Reverse a model output back to jagged true-length form using ``meta``.
+
+    Gathers CP-local shards (zigzag or contiguous reconstruct) then slices each
+    sequence's true length out of the padded layout.
+    """
+    if output.dim() >= 2 and output.shape[0] == 1:
+        flat = output[0]
+    elif output.dim() >= 2 and output.shape[1] == 1:
+        flat = output[:, 0]
+    else:
+        flat = output
+
+    if meta.cp_size > 1:
+        parts = _all_gather_cp_tensor(flat, cp_size=meta.cp_size, cp_group=meta.cp_group)
+        if contiguous:
+            flat = torch.cat(parts, dim=0)
+        else:
+            flat = _reconstruct_full_from_cp_parts(
+                parts, cu_seqlens_padded=meta.cu_seqlens_padded, cp_size=meta.cp_size, dim=0
+            )
+
+    pieces = []
+    for idx, length_t in enumerate(meta.lengths):
+        length = int(length_t.item())
+        start = int(meta.cu_seqlens_padded[idx].item())
+        pieces.append(flat[start : start + length])
+    return torch.nested.as_nested_tensor(pieces, layout=torch.jagged)
+
+
 def _all_gather_cp_tensor(
     tensor: torch.Tensor, *, cp_size: int, cp_group: Any
 ) -> list[torch.Tensor]:
@@ -285,6 +416,7 @@ def pack_nested_thd(
     cp_size: int = 1,
     cp_rank: int = 0,
     cp_group: Any | None = None,
+    split_cp: bool = True,
     labels: torch.Tensor | None = None,
     roll_labels: bool = False,
     loss_mask: torch.Tensor | None = None,
@@ -295,7 +427,10 @@ def pack_nested_thd(
     Mirrors VERL/Megatron's THD engine convention: each sequence is
     padded to the tensor-parallel alignment, concatenated, then represented as
     a single ``[1, local_padded_tokens]`` token row plus ``PackedSeqParams``.
-    For CP>1 the local row uses Megatron/TE zigzag chunking.
+    For CP>1 the local row uses Megatron/TE zigzag chunking unless
+    ``split_cp=False``.  The latter keeps full packed tokens and plain
+    ``PackedSeqParams`` while still padding each sample to CP-compatible
+    alignment, matching the external VERL runtime contract.
     """
 
     if cp_size < 1:
@@ -323,7 +458,7 @@ def pack_nested_thd(
     cu_seqlens_padded = torch.zeros(lengths.numel() + 1, dtype=torch.int32, device=device)
     cu_seqlens_padded[1:] = torch.cumsum(padded_lengths, dim=0)
     total_padded = int(cu_seqlens_padded[-1].item())
-    total_local = total_padded // cp_size
+    total_local = total_padded // cp_size if split_cp else total_padded
     max_seqlen = int(padded_lengths.max().item()) if padded_lengths.numel() else 0
 
     packed_input = torch.zeros(total_local, dtype=input_ids.dtype, device=device)
@@ -343,7 +478,7 @@ def pack_nested_thd(
         length = int(length_t.item())
         padded_length = int(padded_lengths[idx].item())
         full_start = int(cu_seqlens_padded[idx].item())
-        local_start = full_start // cp_size
+        local_start = full_start // cp_size if split_cp else full_start
 
         seq_input = torch.zeros(padded_length, dtype=input_ids.dtype, device=device)
         seq_input[:length] = input_ids[idx]
@@ -366,17 +501,9 @@ def pack_nested_thd(
         seq_positions = torch.zeros(padded_length, dtype=torch.long, device=device)
         seq_positions[:length] = torch.arange(length, dtype=torch.long, device=device)
 
-        local_input = _split_full_to_cp_local(
-            seq_input,
-            cu_seqlens_padded=torch.tensor([0, padded_length], dtype=torch.int32, device=device),
-            cp_size=cp_size,
-            cp_rank=cp_rank,
-            dim=0,
-        )
-        packed_input[local_start : local_start + local_input.numel()] = local_input
-        if seq_labels is not None:
-            local_labels = _split_full_to_cp_local(
-                seq_labels,
+        local_input = (
+            _split_full_to_cp_local(
+                seq_input,
                 cu_seqlens_padded=torch.tensor(
                     [0, padded_length], dtype=torch.int32, device=device
                 ),
@@ -384,17 +511,39 @@ def pack_nested_thd(
                 cp_rank=cp_rank,
                 dim=0,
             )
+            if split_cp
+            else seq_input
+        )
+        packed_input[local_start : local_start + local_input.numel()] = local_input
+        if seq_labels is not None:
+            local_labels = (
+                _split_full_to_cp_local(
+                    seq_labels,
+                    cu_seqlens_padded=torch.tensor(
+                        [0, padded_length], dtype=torch.int32, device=device
+                    ),
+                    cp_size=cp_size,
+                    cp_rank=cp_rank,
+                    dim=0,
+                )
+                if split_cp
+                else seq_labels
+            )
             assert packed_labels is not None
             packed_labels[local_start : local_start + local_labels.numel()] = local_labels
         if seq_loss_mask is not None:
-            local_loss_mask = _split_full_to_cp_local(
-                seq_loss_mask,
-                cu_seqlens_padded=torch.tensor(
-                    [0, padded_length], dtype=torch.int32, device=device
-                ),
-                cp_size=cp_size,
-                cp_rank=cp_rank,
-                dim=0,
+            local_loss_mask = (
+                _split_full_to_cp_local(
+                    seq_loss_mask,
+                    cu_seqlens_padded=torch.tensor(
+                        [0, padded_length], dtype=torch.int32, device=device
+                    ),
+                    cp_size=cp_size,
+                    cp_rank=cp_rank,
+                    dim=0,
+                )
+                if split_cp
+                else seq_loss_mask
             )
             assert packed_loss_mask is not None
             packed_loss_mask[local_start : local_start + local_loss_mask.numel()] = local_loss_mask
@@ -403,14 +552,14 @@ def pack_nested_thd(
     return PackedTHDBatch(
         input_ids=packed_input.unsqueeze(0),
         labels=packed_labels.unsqueeze(0) if packed_labels is not None else None,
-        loss_mask=packed_loss_mask.unsqueeze(0) if packed_loss_mask is not None else None,
+        loss_mask=(packed_loss_mask.unsqueeze(0) if packed_loss_mask is not None else None),
         position_ids=position_ids.unsqueeze(0),
         packed_seq_params=_make_packed_seq_params(
             cu_seqlens_padded=cu_seqlens_padded,
             max_seqlen=max_seqlen,
-            cp_size=cp_size,
-            cp_rank=cp_rank,
-            cp_group=cp_group,
+            cp_size=cp_size if split_cp else 1,
+            cp_rank=cp_rank if split_cp else 0,
+            cp_group=cp_group if split_cp else None,
         ),
         cu_seqlens_padded=cu_seqlens_padded,
         lengths=lengths,
@@ -449,9 +598,16 @@ def unpack_packed_thd_to_nested(output: torch.Tensor, batch: PackedTHDBatch) ->
 __all__ = [
     "PackedSeqParams",
     "PackedTHDBatch",
+    "ThdPackMeta",
+    "has_packed_thd_params",
     "pack_nested_thd",
+    "parallel_state_from_model",
+    "prepare_packed_thd_for_context_parallel",
+    "prepare_packed_thd_kwargs_for_context_parallel",
     "reconstruct_packed_from_cp_parts",
     "roll_packed_thd_left",
     "split_packed_to_cp_local",
+    "thd_pack_meta",
     "unpack_packed_thd_to_nested",
+    "unpack_thd_to_nested",
 ]
diff --git a/experimental/lite/megatron/lite/primitive/train_step.py b/experimental/lite/megatron/lite/primitive/train_step.py
index 630dc5f7701..88eeb896945 100644
--- a/experimental/lite/megatron/lite/primitive/train_step.py
+++ b/experimental/lite/megatron/lite/primitive/train_step.py
@@ -7,9 +7,9 @@
 
 import torch
 import torch.distributed as dist
-
 from megatron.lite.primitive.parallel import ParallelState
 from megatron.lite.primitive.protocols import ExpertClassifierFn, default_expert_classifier
+from megatron.lite.runtime.contracts.loss import split_loss_context, use_loss_context
 
 
 def run_microbatch_loop(
@@ -21,6 +21,7 @@ def run_microbatch_loop(
     dist_opt: bool = False,
     pre_forward_hook: Callable[[torch.Tensor], None] | None = None,
     loss_fn: Callable | None = None,
+    forward_only: bool = False,
 ):
     """Run forward-backward over microbatches with loss accumulation.
 
@@ -40,23 +41,32 @@ def run_microbatch_loop(
             ``loss_fn(model_output: dict, batch) -> (loss: Tensor, metrics: dict)``.
             When provided, ``forward_fn`` output is passed to ``loss_fn`` instead of
             reading ``out["loss"]`` directly. This enables RLHF policy/value losses.
+        forward_only: When True, run forward without backward (e.g. validation /
+            ``infer_batch``). The caller (verl) runs the forward under ``no_grad`` so the
+            loss has no ``grad_fn``; calling ``.backward()`` then raises. Mirrors the
+            pipeline path, which already threads ``forward_only`` to skip backward.
     """
     last_out = None
     all_metrics: list[dict] = []
     for mb in range(num_microbatches):
-        batch = next(data_iter)
+        batch, loss_context = split_loss_context(next(data_iter))
         if pre_forward_hook is not None:
             scale = torch.tensor(1.0 / num_microbatches, device="cuda")
             pre_forward_hook(scale)
-        out = forward_fn(model, batch)
+        with use_loss_context(loss_context):
+            out = forward_fn(model, batch)
         if dist_opt and optimizer is not None and mb == num_microbatches - 1:
             optimizer.grad_sync_enabled = True
         if loss_fn is not None:
-            loss, metrics = loss_fn(out, batch)
-            (loss / num_microbatches).backward()
+            if loss_context is None:
+                loss, metrics = loss_fn(out, batch)
+            else:
+                loss, metrics = loss_fn(out, batch, loss_context)
+            if not forward_only:
+                (loss / num_microbatches).backward()
             out["loss"] = loss.detach()
             all_metrics.append(metrics)
-        else:
+        elif not forward_only:
             (out["loss"] / num_microbatches).backward()
         last_out = out
     if last_out is not None and all_metrics:
diff --git a/experimental/lite/megatron/lite/primitive/utils.py b/experimental/lite/megatron/lite/primitive/utils/__init__.py
similarity index 100%
rename from experimental/lite/megatron/lite/primitive/utils.py
rename to experimental/lite/megatron/lite/primitive/utils/__init__.py
diff --git a/experimental/lite/megatron/lite/primitive/utils/moe.py b/experimental/lite/megatron/lite/primitive/utils/moe.py
new file mode 100644
index 00000000000..23a2cfc591a
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/utils/moe.py
@@ -0,0 +1,422 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""MoE routing, permutation, and router GEMM helpers for MLite primitives."""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+
+from transformer_engine.pytorch.cpp_extensions import general_gemm
+from transformer_engine.pytorch.permutation import moe_permute as fused_permute
+from transformer_engine.pytorch.permutation import (
+    moe_permute_and_pad_with_probs as fused_permute_and_pad_with_probs,
+)
+from transformer_engine.pytorch.permutation import (
+    moe_permute_with_probs as fused_permute_with_probs,
+)
+from transformer_engine.pytorch.permutation import moe_unpermute as fused_unpermute
+from transformer_engine.pytorch.router import (
+    fused_compute_score_for_moe_aux_loss,
+    fused_moe_aux_loss,
+    fused_topk_with_score_function,
+)
+
+
+def _te_general_gemm(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    out_dtype: torch.dtype | None = None,
+    *,
+    layout: str = "TN",
+    out: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+    grad: bool = False,
+):
+    kwargs = dict(
+        out_dtype=out_dtype,
+        quantization_params=None,
+        gelu=None,
+        gelu_in=None,
+        accumulate=False,
+        layout=layout,
+        out=out,
+        bias=bias,
+        use_split_accumulator=False,
+        grad=grad,
+        ub=None,
+        ub_type=None,
+        extra_output=None,
+        bulk_overlap=False,
+    )
+    return general_gemm(a, b, **kwargs)
+
+
+def switch_load_balancing_loss_func(
+    probs: torch.Tensor,
+    tokens_per_expert: torch.Tensor,
+    total_num_tokens: int,
+    topk: int,
+    num_experts: int,
+    moe_aux_loss_coeff: float,
+    *,
+    fused: bool = False,
+    padding_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if padding_mask is not None:
+        probs = probs * padding_mask.unsqueeze(-1)
+
+    if fused:
+        return fused_moe_aux_loss(
+            probs=probs,
+            tokens_per_expert=tokens_per_expert,
+            total_num_tokens=total_num_tokens,
+            topk=topk,
+            num_experts=num_experts,
+            coeff=moe_aux_loss_coeff,
+        )
+
+    aggregated_probs_per_expert = probs.sum(dim=0)
+    return torch.sum(aggregated_probs_per_expert * tokens_per_expert) * (
+        num_experts * moe_aux_loss_coeff / (topk * total_num_tokens * total_num_tokens)
+    )
+
+
+def permute(
+    tokens: torch.Tensor,
+    routing_map: torch.Tensor,
+    probs: Optional[torch.Tensor] = None,
+    num_out_tokens: Optional[int] = None,
+    fused: bool = False,
+    drop_and_pad: bool = False,
+    tokens_per_expert: Optional[torch.Tensor] = None,
+    align_size: int = 0,
+) -> Tuple[
+    torch.Tensor,
+    Optional[torch.Tensor],
+    torch.Tensor,
+    Optional[torch.Tensor],
+    Optional[torch.Tensor],
+]:
+    if fused and probs is None:
+        permuted_input, sorted_indices = fused_permute(
+            tokens, routing_map, num_out_tokens=num_out_tokens
+        )
+        return permuted_input, None, sorted_indices, None, tokens_per_expert
+
+    if fused and probs is not None:
+        if tokens_per_expert is not None and align_size > 0:
+            return fused_permute_and_pad_with_probs(
+                tokens, probs, routing_map, tokens_per_expert, align_size
+            )
+        output, permuted_probs, row_id_map = fused_permute_with_probs(
+            tokens, probs, routing_map, num_out_tokens=num_out_tokens
+        )
+        return output, permuted_probs, row_id_map, None, tokens_per_expert
+
+    num_tokens, _hidden = tokens.shape
+    num_experts = routing_map.shape[1]
+    permuted_probs = None
+    if drop_and_pad and num_out_tokens is not None:
+        capacity = num_out_tokens // num_experts
+        assert not routing_map.requires_grad
+        routing_map = routing_map.to(dtype=torch.int8).T.contiguous()
+        sorted_indices = routing_map.argsort(dim=-1, descending=True, stable=True)[
+            :, :capacity
+        ].contiguous()
+        sorted_indices = sorted_indices.view(-1)
+
+        if probs is not None:
+            probs_t_1d = probs.T.contiguous().view(-1)
+            indices_dim0 = torch.arange(num_experts, device=routing_map.device).unsqueeze(-1)
+            indices_dim1 = sorted_indices.view(num_experts, capacity)
+            indices_1d = (indices_dim0 * num_tokens + indices_dim1).view(-1)
+            permuted_probs = probs_t_1d.index_select(0, indices_1d)
+    else:
+        if num_out_tokens is None:
+            raise AssertionError("num_out_tokens is required for argsort-based permute")
+
+        routing_map = routing_map.bool().T.contiguous()
+        flat_sorted = routing_map.reshape(-1).argsort(descending=True, stable=True)
+        flat_sorted = flat_sorted[:num_out_tokens]
+        sorted_indices = flat_sorted % num_tokens
+
+        if probs is not None:
+            permuted_probs = probs.T.contiguous().reshape(-1)[flat_sorted]
+
+    return (
+        tokens.index_select(0, sorted_indices),
+        permuted_probs,
+        sorted_indices,
+        None,
+        tokens_per_expert,
+    )
+
+
+def unpermute(
+    permuted_tokens: torch.Tensor,
+    sorted_indices: torch.Tensor,
+    restore_shape: torch.Size,
+    probs: Optional[torch.Tensor] = None,
+    routing_map: Optional[torch.Tensor] = None,
+    fused: bool = False,
+    drop_and_pad: bool = False,
+    pad_offsets: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if fused:
+        kwargs = {}
+        if pad_offsets is not None:
+            kwargs["pad_offsets"] = pad_offsets
+        return fused_unpermute(
+            permuted_tokens,
+            sorted_indices,
+            merging_probs=probs,
+            restore_shape=restore_shape,
+            **kwargs,
+        )
+
+    _, hidden = restore_shape
+    input_dtype = permuted_tokens.dtype
+
+    if probs is not None:
+        assert routing_map is not None, "Mask must be provided to permute the probs."
+        if drop_and_pad:
+            num_experts = routing_map.size(1)
+            num_permuted_tokens = sorted_indices.size(0)
+            capacity = num_permuted_tokens // num_experts
+            num_unpermuted_tokens = probs.size(0)
+            probs_t_1d = probs.T.contiguous().view(-1)
+            indices_dim0 = torch.arange(num_experts, device=routing_map.device).unsqueeze(-1)
+            indices_dim1 = sorted_indices.view(num_experts, capacity)
+            indices_1d = (indices_dim0 * num_unpermuted_tokens + indices_dim1).view(-1)
+            permuted_probs = probs_t_1d.index_select(0, indices_1d)
+        else:
+            permuted_probs = probs.T.contiguous().masked_select(routing_map.T.contiguous())
+        permuted_tokens = permuted_tokens * permuted_probs.unsqueeze(-1)
+
+    output_tokens = torch.zeros(
+        restore_shape, dtype=permuted_tokens.dtype, device=permuted_tokens.device
+    )
+    if torch.are_deterministic_algorithms_enabled():
+        output_tokens.index_add_(0, sorted_indices, permuted_tokens)
+    else:
+        output_tokens.scatter_add_(
+            0, sorted_indices.unsqueeze(1).expand(-1, hidden), permuted_tokens
+        )
+    return output_tokens.to(dtype=input_dtype)
+
+
+def group_limited_topk(
+    scores: torch.Tensor,
+    topk: int,
+    num_tokens: int,
+    num_experts: int,
+    num_groups: int,
+    group_topk: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    group_scores = (
+        scores.view(num_tokens, num_groups, -1).topk(topk // group_topk, dim=-1)[0].sum(dim=-1)
+    )
+    group_idx = torch.topk(group_scores, k=group_topk, dim=-1, sorted=False)[1]
+    group_mask = torch.zeros_like(group_scores)
+    group_mask.scatter_(1, group_idx, 1)
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_tokens, num_groups, num_experts // num_groups)
+        .reshape(num_tokens, -1)
+    )
+    masked_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))
+    return torch.topk(masked_scores, k=topk, dim=-1)
+
+
+def topk_routing_with_score_function(
+    logits: torch.Tensor,
+    topk: int,
+    use_pre_softmax: bool = False,
+    num_groups: Optional[int] = None,
+    group_topk: Optional[int] = None,
+    scaling_factor: Optional[float] = None,
+    score_function: str = "softmax",
+    expert_bias: Optional[torch.Tensor] = None,
+    fused: bool = False,
+    dense_output: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}."
+    num_tokens, num_experts = logits.shape
+    if fused:
+        return fused_topk_with_score_function(
+            logits=logits,
+            topk=topk,
+            use_pre_softmax=use_pre_softmax,
+            num_groups=num_groups,
+            group_topk=group_topk,
+            scaling_factor=scaling_factor,
+            score_function=score_function,
+            expert_bias=expert_bias,
+        )
+
+    def compute_topk(
+        scores: torch.Tensor,
+        k: int,
+        groups: Optional[int] = None,
+        groups_topk: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if groups_topk:
+            assert groups is not None
+            return group_limited_topk(
+                scores=scores,
+                topk=k,
+                num_tokens=num_tokens,
+                num_experts=num_experts,
+                num_groups=groups,
+                group_topk=groups_topk,
+            )
+        return torch.topk(scores, k=k, dim=1, sorted=torch.is_grad_enabled())
+
+    if score_function == "softmax":
+        if use_pre_softmax:
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+            probs, top_indices = compute_topk(scores, topk, num_groups, group_topk)
+        else:
+            scores, top_indices = compute_topk(logits, topk, num_groups, group_topk)
+            probs = torch.softmax(scores, dim=-1, dtype=torch.float32)
+    elif score_function in ("sigmoid", "sqrtsoftplus"):
+        if score_function == "sigmoid":
+            scores = torch.sigmoid(logits.float())
+        else:
+            scores = torch.nn.functional.softplus(logits.float()).sqrt()
+        if expert_bias is not None:
+            scores_for_routing = scores + expert_bias.float()
+            _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)
+            scores = torch.gather(scores, dim=1, index=top_indices)
+        else:
+            scores, top_indices = compute_topk(scores, topk, num_groups, group_topk)
+        probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores
+    else:
+        raise ValueError(f"Invalid score_function: {score_function}")
+
+    if scaling_factor:
+        probs = probs * scaling_factor
+    probs = probs.type_as(logits)
+
+    if dense_output:
+        return probs, top_indices
+
+    if torch.are_deterministic_algorithms_enabled():
+        routing_probs = torch.zeros_like(logits)
+        rows = torch.arange(num_tokens, device=logits.device).unsqueeze(1)
+        routing_probs.index_put_((rows, top_indices), probs, accumulate=False)
+        routing_map = torch.zeros_like(logits, dtype=logits.dtype)
+        routing_map.index_put_(
+            (rows, top_indices), torch.ones_like(probs, dtype=routing_map.dtype), accumulate=False
+        )
+        routing_map = routing_map.bool()
+    else:
+        routing_probs = torch.zeros_like(logits).scatter(1, top_indices, probs)
+        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    return routing_probs, routing_map
+
+
+def compute_routing_scores_for_aux_loss(
+    logits: torch.Tensor,
+    topk: int,
+    score_function: str,
+    *,
+    fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if fused:
+        routing_map, scores = fused_compute_score_for_moe_aux_loss(
+            logits=logits, topk=topk, score_function=score_function
+        )
+    else:
+        if score_function == "softmax":
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+        elif score_function == "sigmoid":
+            scores = torch.sigmoid(logits.float())
+            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+        elif score_function == "sqrtsoftplus":
+            scores = torch.nn.functional.softplus(logits.float()).sqrt()
+            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+        else:
+            raise ValueError(f"Invalid score_function: {score_function}")
+        _, top_indices = torch.topk(scores, k=topk, dim=1)
+        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+
+    if padding_mask is not None:
+        valid_mask = (~padding_mask).unsqueeze(-1)
+        routing_map = routing_map * valid_mask
+        scores = scores * valid_mask
+    return routing_map, scores
+
+
+class RouterGatingLinearFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        inp: torch.Tensor,
+        weight: torch.Tensor,
+        bias: torch.Tensor | None,
+        router_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        ctx.save_for_backward(inp, weight, bias)
+        ctx.router_dtype = router_dtype
+        ctx.input_dtype = inp.dtype
+        ctx.weight_dtype = weight.dtype
+        inp_shape = inp.shape
+        inp = inp.view(-1, inp_shape[-1])
+
+        gemm_out = None
+        if router_dtype != torch.float64:
+            gemm_out = _te_general_gemm(weight, inp, router_dtype, layout="TN", bias=bias)
+        if gemm_out is not None:
+            output = gemm_out[0]
+        elif bias is None:
+            output = torch.mm(inp.to(router_dtype), weight.to(router_dtype).t())
+        else:
+            output = torch.addmm(
+                bias.to(router_dtype), inp.to(router_dtype), weight.to(router_dtype).t()
+            )
+        return output.view(*inp_shape[:-1], -1)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        inp, weight, bias = ctx.saved_tensors
+        inp_shape = inp.shape
+        grad_shape = grad_output.shape
+        inp = inp.view(-1, inp_shape[-1])
+        grad_output = grad_output.view(-1, grad_shape[-1])
+
+        grad_input_out = grad_weight_out = None
+        if ctx.router_dtype != torch.float64:
+            grad_input_out = _te_general_gemm(
+                weight.to(ctx.router_dtype), grad_output, ctx.router_dtype, layout="NN", grad=True
+            )
+            grad_weight_out = _te_general_gemm(
+                inp.to(ctx.router_dtype), grad_output, ctx.router_dtype, layout="NT", grad=True
+            )
+        if grad_input_out is not None and grad_weight_out is not None:
+            grad_input = grad_input_out[0].to(ctx.input_dtype)
+            grad_weight = grad_weight_out[0].to(ctx.weight_dtype)
+        else:
+            grad_input = torch.mm(grad_output, weight.to(ctx.router_dtype)).to(ctx.input_dtype)
+            grad_weight = torch.mm(grad_output.t(), inp.to(ctx.router_dtype)).to(ctx.weight_dtype)
+        grad_bias = grad_output.sum(dim=0).to(ctx.weight_dtype) if bias is not None else None
+        return grad_input.view(*inp_shape), grad_weight, grad_bias, None
+
+
+def router_gating_linear(
+    inp: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, router_dtype: torch.dtype
+) -> torch.Tensor:
+    return RouterGatingLinearFunction.apply(inp, weight, bias, router_dtype)
+
+
+__all__ = [
+    "compute_routing_scores_for_aux_loss",
+    "group_limited_topk",
+    "permute",
+    "router_gating_linear",
+    "switch_load_balancing_loss_func",
+    "topk_routing_with_score_function",
+    "unpermute",
+]
diff --git a/experimental/lite/megatron/lite/primitive/utils/packed_seq.py b/experimental/lite/megatron/lite/primitive/utils/packed_seq.py
new file mode 100644
index 00000000000..d28b28fb850
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/utils/packed_seq.py
@@ -0,0 +1,61 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Packed sequence parameter containers for THD-format primitives."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class PackedSeqParams:
+    """Parameters consumed by TE DotProductAttention and MLite THD helpers."""
+
+    qkv_format: str | None = "thd"
+    cu_seqlens_q: Tensor | None = None
+    cu_seqlens_kv: Tensor | None = None
+    cu_seqlens_q_padded: Tensor | None = None
+    cu_seqlens_kv_padded: Tensor | None = None
+    max_seqlen_q: int | None = None
+    max_seqlen_kv: int | None = None
+    local_cp_size: int | None = None
+    cp_group: Any | None = None
+    total_tokens: int | None = None
+    seq_idx: Tensor | None = None
+    cp_rank: int | None = None
+
+    def __post_init__(self) -> None:
+        cu_seqlens = (
+            self.cu_seqlens_q_padded if self.cu_seqlens_q_padded is not None else self.cu_seqlens_q
+        )
+        if isinstance(cu_seqlens, Tensor) and self.total_tokens is not None:
+            total_tokens_tensor = torch.tensor(
+                [self.total_tokens], dtype=cu_seqlens.dtype, device=cu_seqlens.device
+            )
+            cu_seqlens_with_max = torch.cat([cu_seqlens, total_tokens_tensor])
+            seq_lengths = (cu_seqlens_with_max[1:] - cu_seqlens_with_max[:-1]).clamp(min=0)
+            self.seq_idx = (
+                torch.repeat_interleave(
+                    torch.arange(seq_lengths.numel(), device=cu_seqlens.device), seq_lengths
+                )
+                .to(torch.int32)
+                .unsqueeze(0)
+            )
+
+    @staticmethod
+    def from_cu_seqlens(cu_seqlens: Tensor, max_seqlen: int) -> PackedSeqParams:
+        return PackedSeqParams(
+            qkv_format="thd",
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_kv=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_kv=max_seqlen,
+            cu_seqlens_q_padded=cu_seqlens,
+            cu_seqlens_kv_padded=cu_seqlens,
+        )
+
+
+__all__ = ["PackedSeqParams"]
diff --git a/experimental/lite/megatron/lite/primitive/utils/rope.py b/experimental/lite/megatron/lite/primitive/utils/rope.py
new file mode 100644
index 00000000000..115ede80f67
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/utils/rope.py
@@ -0,0 +1,144 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""RoPE helpers for MLite attention primitives."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+def get_pos_emb_on_this_cp_rank(
+    pos_emb: Tensor, seq_dim: int, cp_group: torch.distributed.ProcessGroup
+) -> Tensor:
+    if cp_group is None:
+        raise ValueError("cp_group must be provided to get positional embedding per CP rank")
+    cp_size = cp_group.size()
+    cp_rank = cp_group.rank()
+    cp_idx = torch.tensor(
+        [cp_rank, (2 * cp_size - cp_rank - 1)], device=pos_emb.device, dtype=torch.long
+    )
+    pos_emb = pos_emb.view(
+        *pos_emb.shape[:seq_dim], 2 * cp_size, -1, *pos_emb.shape[(seq_dim + 1) :]
+    )
+    pos_emb = pos_emb.index_select(seq_dim, cp_idx)
+    return pos_emb.view(*pos_emb.shape[:seq_dim], -1, *pos_emb.shape[(seq_dim + 2) :])
+
+
+def _rotate_half(x: Tensor, rotary_interleaved: bool) -> Tensor:
+    if not rotary_interleaved:
+        x1, x2 = torch.chunk(x, 2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    x1 = x[:, :, :, ::2]
+    x2 = x[:, :, :, 1::2]
+    x_new = torch.stack((-x2, x1), dim=-1)
+    return x_new.view(x_new.shape[0], x_new.shape[1], x_new.shape[2], -1)
+
+
+def _apply_rotary_pos_emb_bshd(
+    t: Tensor,
+    freqs: Tensor,
+    rotary_interleaved: bool = False,
+    mla_rotary_interleaved: bool = False,
+    mscale: float = 1.0,
+    multi_latent_attention: Optional[bool] = None,
+) -> Tensor:
+    if multi_latent_attention is not None:
+        warnings.warn(
+            "multi_latent_attention is deprecated. Use mla_rotary_interleaved instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        mla_rotary_interleaved = multi_latent_attention
+
+    rot_dim = freqs.shape[-1]
+    t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+
+    if mla_rotary_interleaved:
+        x1 = t[..., 0::2]
+        x2 = t[..., 1::2]
+        t = torch.cat((x1, x2), dim=-1)
+
+    cos_ = (torch.cos(freqs) * mscale).to(t.dtype)
+    sin_ = (torch.sin(freqs) * mscale).to(t.dtype)
+    t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_)
+    return torch.cat((t, t_pass), dim=-1)
+
+
+def _get_thd_freqs_on_this_cp_rank(
+    cp_rank: int, cp_size: int, x: Tensor, freqs: Tensor, offset: int = 0
+) -> Tensor:
+    if cp_size > 1:
+        cp_seg = x.size(0) // 2
+        full_seqlen = cp_size * x.size(0)
+        return torch.cat(
+            [
+                freqs[offset + cp_rank * cp_seg : offset + (cp_rank + 1) * cp_seg],
+                freqs[
+                    offset
+                    + full_seqlen
+                    - (cp_rank + 1) * cp_seg : offset
+                    + full_seqlen
+                    - cp_rank * cp_seg
+                ],
+            ]
+        )
+    return freqs[offset : offset + x.size(0)]
+
+
+def _apply_rotary_pos_emb_thd(
+    t: Tensor,
+    cu_seqlens: Tensor,
+    freqs: Tensor,
+    rotary_interleaved: bool = False,
+    mla_rotary_interleaved: bool = False,
+    mscale: float = 1.0,
+    cp_group: torch.distributed.ProcessGroup = None,
+    multi_latent_attention: Optional[bool] = None,
+) -> Tensor:
+    if multi_latent_attention is not None:
+        warnings.warn(
+            "multi_latent_attention is deprecated. Use mla_rotary_interleaved instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        mla_rotary_interleaved = multi_latent_attention
+
+    if cp_group is None:
+        raise ValueError("cp_group must be provided for THD format RoPE")
+    cp_size = cp_group.size()
+    cp_rank = cp_group.rank()
+    seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist()
+
+    sequence_splits = torch.split(t, seqlens)
+    if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens[-1]:
+        freq_slices = []
+        for i, x in enumerate(sequence_splits):
+            seq_start_offset = cu_seqlens[i].item()
+            freq_slices.append(
+                _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset)
+            )
+        freqs_packed = torch.cat(freq_slices, dim=0)
+    else:
+        freqs_packed = torch.cat(
+            [_get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) for x in sequence_splits],
+            dim=0,
+        )
+
+    return _apply_rotary_pos_emb_bshd(
+        t.unsqueeze(1),
+        freqs_packed,
+        rotary_interleaved=rotary_interleaved,
+        mla_rotary_interleaved=mla_rotary_interleaved,
+        mscale=mscale,
+    ).squeeze(1)
+
+
+__all__ = [
+    "_apply_rotary_pos_emb_bshd",
+    "_apply_rotary_pos_emb_thd",
+    "_rotate_half",
+    "get_pos_emb_on_this_cp_rank",
+]
diff --git a/experimental/lite/megatron/lite/primitive/utils/rotary.py b/experimental/lite/megatron/lite/primitive/utils/rotary.py
new file mode 100644
index 00000000000..df5ecbee508
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/utils/rotary.py
@@ -0,0 +1,299 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Rotary embedding modules for MLite primitives."""
+
+from __future__ import annotations
+
+import math
+from functools import lru_cache
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+
+from megatron.lite.primitive.utils.rope import get_pos_emb_on_this_cp_rank
+
+
+def _default_rope_device(use_cpu_initialization: bool) -> str | torch.device:
+    if use_cpu_initialization or not torch.cuda.is_available():
+        return "cpu"
+    return torch.device("cuda", torch.cuda.current_device())
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary embedding with optional context-parallel slicing."""
+
+    def __init__(
+        self,
+        kv_channels: int,
+        rotary_percent: float = 1.0,
+        rotary_interleaved: bool = False,
+        seq_len_interpolation_factor: float | None = None,
+        rotary_base: float = 10000,
+        rope_scaling: bool = False,
+        rope_scaling_factor: float = 8.0,
+        use_cpu_initialization: bool = False,
+        cp_group: Optional[torch.distributed.ProcessGroup] = None,
+    ) -> None:
+        super().__init__()
+        dim = kv_channels
+        if rotary_percent < 1.0:
+            dim = int(dim * rotary_percent)
+        self.rotary_interleaved = rotary_interleaved
+        self.seq_len_interpolation_factor = seq_len_interpolation_factor
+        device = _default_rope_device(use_cpu_initialization)
+        self.inv_freq = 1.0 / (
+            rotary_base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+        )
+        if rope_scaling:
+            self.inv_freq = self._apply_scaling(self.inv_freq, factor=rope_scaling_factor)
+        self.cp_group = cp_group
+
+    def _apply_scaling(
+        self,
+        freqs: Tensor,
+        factor: float = 8,
+        low_freq_factor: float = 1,
+        high_freq_factor: float = 4,
+        original_max_position_embeddings: int = 8192,
+    ) -> Tensor:
+        low_freq_wavelen = original_max_position_embeddings / low_freq_factor
+        high_freq_wavelen = original_max_position_embeddings / high_freq_factor
+
+        wavelen = 2 * math.pi / freqs
+        inv_freq_llama = torch.where(wavelen > low_freq_wavelen, freqs / factor, freqs)
+        smooth_factor = (original_max_position_embeddings / wavelen - low_freq_factor) / (
+            high_freq_factor - low_freq_factor
+        )
+        smoothed_inv_freq = (
+            1 - smooth_factor
+        ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
+        is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
+        return torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
+
+    def get_freqs_non_repeated(self, max_seq_len: int, offset: int = 0) -> Tensor:
+        seq = (
+            torch.arange(max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+            + offset
+        )
+        if self.seq_len_interpolation_factor is not None:
+            seq *= 1 / self.seq_len_interpolation_factor
+        return torch.outer(seq, self.inv_freq)
+
+    def get_cos_sin(self, max_seq_len: int, offset: int = 0) -> tuple[Tensor, Tensor]:
+        freqs = self.get_freqs_non_repeated(max_seq_len, offset)
+        return torch.cos(freqs), torch.sin(freqs)
+
+    def get_emb(self, max_seq_len: int, offset: int = 0) -> Tensor:
+        if self.inv_freq.device.type == "cpu" and torch.cuda.is_available():
+            self.inv_freq = self.inv_freq.to(device=torch.cuda.current_device())
+
+        freqs = self.get_freqs_non_repeated(max_seq_len, offset)
+        if not self.rotary_interleaved:
+            emb = torch.cat((freqs, freqs), dim=-1)
+        else:
+            emb = torch.stack((freqs.view(-1, 1), freqs.view(-1, 1)), dim=-1).view(
+                freqs.shape[0], -1
+            )
+        return emb[:, None, None, :]
+
+    @lru_cache(maxsize=32)
+    def forward(
+        self,
+        max_seq_len: int,
+        offset: int = 0,
+        packed_seq: bool = False,
+        cp_group: Optional[torch.distributed.ProcessGroup] = None,
+    ) -> Tensor:
+        emb = self.get_emb(max_seq_len, offset)
+        if cp_group is None:
+            cp_group = self.cp_group
+        if cp_group is not None and cp_group.size() > 1 and not packed_seq:
+            emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group)
+        return emb
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        state_dict.pop(f"{prefix}inv_freq", None)
+        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+
+class YarnRotaryEmbedding(RotaryEmbedding):
+    """YARN rotary embedding variant used by MLA-style models."""
+
+    def __init__(
+        self,
+        kv_channels: int,
+        rotary_percent: float = 1.0,
+        rotary_interleaved: bool = False,
+        seq_len_interpolation_factor: Optional[float] = None,
+        rotary_base: float = 10000.0,
+        use_cpu_initialization: bool = False,
+        scaling_factor: float = 1.0,
+        original_max_position_embeddings: int = 4096,
+        beta_fast: float = 32.0,
+        beta_slow: float = 1.0,
+        mscale: float = 1.0,
+        mscale_all_dim: float = 0.0,
+        correction_range_round_to_int: bool = True,
+        cp_group: Optional[torch.distributed.ProcessGroup] = None,
+    ):
+        self.dim = kv_channels
+        self.rotary_base = rotary_base
+        self.scaling_factor = scaling_factor
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.beta_fast = beta_fast
+        self.beta_slow = beta_slow
+        self.mscale = mscale
+        self.mscale_all_dim = mscale_all_dim
+        self.correction_range_round_to_int = correction_range_round_to_int
+
+        device = _default_rope_device(use_cpu_initialization)
+        self.inv_freq_extra = 1.0 / (
+            self.rotary_base
+            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
+        )
+        self.inv_freq_inter = 1.0 / (
+            self.scaling_factor
+            * self.rotary_base
+            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
+        )
+        super().__init__(
+            kv_channels=kv_channels,
+            rotary_percent=rotary_percent,
+            rotary_interleaved=rotary_interleaved,
+            seq_len_interpolation_factor=seq_len_interpolation_factor,
+            rotary_base=rotary_base,
+            use_cpu_initialization=use_cpu_initialization,
+            cp_group=cp_group,
+        )
+        self._set_cos_sin_cache(
+            self.original_max_position_embeddings, offset=0, dtype=torch.get_default_dtype()
+        )
+        self.forward.cache_clear()
+
+    def get_emb(self, max_seq_len: int, offset: int = 0) -> tuple[Tensor, float]:
+        if self.rotary_interleaved:
+            raise AssertionError("YARN RoPE does not support interleaved rotary embeddings")
+        if self.inv_freq_extra.device.type == "cpu" and torch.cuda.is_available():
+            self.inv_freq_extra = self.inv_freq_extra.to(device=torch.cuda.current_device())
+        if self.inv_freq_inter.device.type == "cpu" and torch.cuda.is_available():
+            self.inv_freq_inter = self.inv_freq_inter.to(device=torch.cuda.current_device())
+
+        low, high = _yarn_find_correction_range(
+            self.beta_fast,
+            self.beta_slow,
+            self.dim,
+            self.rotary_base,
+            self.original_max_position_embeddings,
+            self.correction_range_round_to_int,
+        )
+        inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(
+            low, high, self.dim // 2, device=self.inv_freq_extra.device
+        ).to(dtype=torch.float32)
+        inv_freq = self.inv_freq_inter * (1 - inv_freq_mask) + self.inv_freq_extra * inv_freq_mask
+        seq = (
+            torch.arange(
+                max_seq_len, device=self.inv_freq_extra.device, dtype=self.inv_freq_extra.dtype
+            )
+            + offset
+        )
+        freqs = torch.outer(seq, inv_freq)
+        concentration = _yarn_get_concentration_factor(
+            self.scaling_factor, self.mscale, self.mscale_all_dim
+        )
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb[:, None, None, :], concentration
+
+    @lru_cache(maxsize=32)
+    def forward(
+        self,
+        max_seq_len: int,
+        offset: int = 0,
+        packed_seq: bool = False,
+        cp_group: Optional[torch.distributed.ProcessGroup] = None,
+    ) -> tuple[Tensor, float]:
+        emb, concentration = self.get_emb(max_seq_len, offset)
+        if cp_group is None:
+            cp_group = self.cp_group
+        if cp_group is not None and cp_group.size() > 1 and not packed_seq:
+            emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group)
+        return emb, concentration
+
+    def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False, cp_group=None):
+        self.max_seq_len_cached = seq_len
+        self.offset_cached = offset
+        self.dtype_cached = dtype
+        self.packed_seq_cached = packed_seq
+        emb, concentration = self.forward(seq_len, offset, packed_seq=packed_seq, cp_group=cp_group)
+        self.register_buffer(
+            "cos_cached", (emb.cos() * concentration).to(dtype).contiguous(), persistent=False
+        )
+        self.register_buffer(
+            "sin_cached", (emb.sin() * concentration).to(dtype).contiguous(), persistent=False
+        )
+
+    def get_cached_cos_sin(
+        self, seq_len, offset=0, dtype=torch.get_default_dtype(), packed_seq=False, cp_group=None
+    ):
+        if (
+            seq_len > self.max_seq_len_cached
+            or offset != self.offset_cached
+            or dtype != self.dtype_cached
+            or packed_seq != self.packed_seq_cached
+        ):
+            self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq, cp_group)
+        return self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]
+
+
+def _yarn_find_correction_dim(
+    num_rotations: float, dim: int, rotary_base: float = 10000, max_position_embeddings: int = 2048
+) -> float:
+    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
+        2 * math.log(rotary_base)
+    )
+
+
+def _yarn_find_correction_range(
+    low_rot: float,
+    high_rot: float,
+    dim: int,
+    rotary_base: float = 10000,
+    max_position_embeddings: int = 2048,
+    round_to_int: bool = True,
+) -> tuple[int, int]:
+    low = _yarn_find_correction_dim(low_rot, dim, rotary_base, max_position_embeddings)
+    high = _yarn_find_correction_dim(high_rot, dim, rotary_base, max_position_embeddings)
+    if round_to_int:
+        low = math.floor(low)
+        high = math.ceil(high)
+    return max(low, 0), min(high, dim - 1)
+
+
+def _yarn_linear_ramp_mask(
+    minimum: float, maximum: float, dim: int, device: torch.device
+) -> Tensor:
+    if minimum == maximum:
+        maximum += 0.001
+    linear_func = (torch.arange(dim, dtype=torch.float32, device=device) - minimum) / (
+        maximum - minimum
+    )
+    return torch.clamp(linear_func, 0, 1)
+
+
+def _yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+@lru_cache(maxsize=8)
+def _yarn_get_concentration_factor(
+    scaling_factor: float, mscale: Optional[float], mscale_all_dim: Optional[float]
+) -> float:
+    if mscale is None or mscale_all_dim is None:
+        return _yarn_get_mscale(scaling_factor)
+    return float(
+        _yarn_get_mscale(scaling_factor, mscale) / _yarn_get_mscale(scaling_factor, mscale_all_dim)
+    )
+
+
+__all__ = ["RotaryEmbedding", "YarnRotaryEmbedding", "_yarn_get_mscale"]
diff --git a/experimental/lite/megatron/lite/runtime/contracts/__init__.py b/experimental/lite/megatron/lite/runtime/contracts/__init__.py
index e8485baca29..7087a5d70a9 100644
--- a/experimental/lite/megatron/lite/runtime/contracts/__init__.py
+++ b/experimental/lite/megatron/lite/runtime/contracts/__init__.py
@@ -27,6 +27,7 @@
         TrainBatch,
     )
     from megatron.lite.runtime.contracts.handle import ModelHandle
+    from megatron.lite.runtime.contracts.loss import LossContext
 
 __all__ = [
     "MegatronLiteConfig",
@@ -34,6 +35,7 @@
     "BridgeConfig",
     "DebugConfig",
     "ForwardResult",
+    "LossContext",
     "ModelHandle",
     "ModelOutputs",
     "OptimizerConfig",
@@ -51,6 +53,7 @@ def __getattr__(name: str):
         "DebugConfig": "megatron.lite.runtime.backends.mlite.config",
         "MegatronLiteConfig": "megatron.lite.runtime.backends.mlite.config",
         "ForwardResult": "megatron.lite.runtime.contracts.data",
+        "LossContext": "megatron.lite.runtime.contracts.loss",
         "ModelHandle": "megatron.lite.runtime.contracts.handle",
         "ModelOutputs": "megatron.lite.runtime.contracts.data",
         "OptimizerConfig": "megatron.lite.runtime.contracts.config",
diff --git a/experimental/lite/megatron/lite/runtime/contracts/config.py b/experimental/lite/megatron/lite/runtime/contracts/config.py
index 7ef5d29108e..74d3af85f97 100644
--- a/experimental/lite/megatron/lite/runtime/contracts/config.py
+++ b/experimental/lite/megatron/lite/runtime/contracts/config.py
@@ -28,6 +28,10 @@ class ParallelConfig:
     pp: int = 1
     vpp: int = 1
     cp: int = 1
+    # Optional explicit mcore pipeline layout for advanced users (custom mode),
+    # e.g. "E|t*5|t*6|t,m,L" (MTP `m` must sit on the final/loss stage; standalone MTP
+    # is not supported yet). None -> auto-balanced from (num_layers, pp, mtp).
+    pp_layout: str | list | None = None
 
 
 @dataclass
diff --git a/experimental/lite/megatron/lite/runtime/contracts/data.py b/experimental/lite/megatron/lite/runtime/contracts/data.py
index ff028a4fa20..d431a728bfd 100644
--- a/experimental/lite/megatron/lite/runtime/contracts/data.py
+++ b/experimental/lite/megatron/lite/runtime/contracts/data.py
@@ -28,10 +28,6 @@ def sizes(self) -> torch.Tensor:
         """Per-sequence token counts.  Shape ``[num_seqs]``."""
         raise NotImplementedError
 
-    def __getitem__(self, key: str) -> Any:
-        """Dict-like field access (``batch["input_ids"]``)."""
-        raise NotImplementedError
-
 
 @dataclass(slots=True)
 class PackedBatch(Batch):
@@ -55,11 +51,6 @@ def __len__(self) -> int:
     def sizes(self) -> torch.Tensor:
         return self.seq_lens
 
-    def __getitem__(self, key: str) -> Any:
-        if key in self.__slots__:
-            return getattr(self, key)
-        return self.extras[key]
-
     @property
     def cu_seqlens(self) -> torch.Tensor:
         """Cumulative sequence lengths for THD attention.  Shape ``[num_seqs+1]``."""
@@ -120,4 +111,10 @@ class ForwardResult:
     metrics: dict[str, Any] = field(default_factory=dict)
 
 
-__all__ = ["Batch", "ForwardResult", "ModelOutputs", "PackedBatch", "TrainBatch"]
+__all__ = [
+    "Batch",
+    "ForwardResult",
+    "ModelOutputs",
+    "PackedBatch",
+    "TrainBatch",
+]
diff --git a/experimental/lite/megatron/lite/runtime/contracts/loss.py b/experimental/lite/megatron/lite/runtime/contracts/loss.py
new file mode 100644
index 00000000000..47bcdb6114d
--- /dev/null
+++ b/experimental/lite/megatron/lite/runtime/contracts/loss.py
@@ -0,0 +1,52 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Loss-side runtime context, kept separate from batch data contracts."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True, slots=True)
+class LossContext:
+    """Per-microbatch loss/output policy for model-owned loss helpers."""
+
+    temperature: float = 1.0
+    calculate_entropy: bool = False
+    return_log_probs: bool = True
+    loss_scale: float = 1.0
+    source_batch: Any | None = None
+
+
+_CURRENT_LOSS_CONTEXT: ContextVar[LossContext | None] = ContextVar(
+    "megatron_lite_loss_context", default=None
+)
+
+
+def get_loss_context() -> LossContext | None:
+    return _CURRENT_LOSS_CONTEXT.get()
+
+
+@contextmanager
+def use_loss_context(loss_context: LossContext | None) -> Iterator[None]:
+    token = _CURRENT_LOSS_CONTEXT.set(loss_context)
+    try:
+        yield
+    finally:
+        _CURRENT_LOSS_CONTEXT.reset(token)
+
+
+def split_loss_context(item):
+    if (
+        isinstance(item, tuple)
+        and len(item) == 2
+        and (item[1] is None or isinstance(item[1], LossContext))
+    ):
+        return item
+    return item, None
+
+
+__all__ = ["LossContext", "get_loss_context", "split_loss_context", "use_loss_context"]
diff --git a/experimental/lite/tests/smoke/model/test_qwen_lite_forward_smoke.py b/experimental/lite/tests/smoke/model/test_qwen_lite_forward_smoke.py
index b00ee6b9646..d5a03c32d20 100644
--- a/experimental/lite/tests/smoke/model/test_qwen_lite_forward_smoke.py
+++ b/experimental/lite/tests/smoke/model/test_qwen_lite_forward_smoke.py
@@ -88,7 +88,7 @@ def _tiny_qwen35_config():
         max_position_embeddings=16,
         partial_rotary_factor=1.0,
         mrope_section=[1, 1, 0],
-        layer_types=["full_attention"],
+        layer_types=["linear_attention"],
     )
 
 
@@ -115,13 +115,19 @@ def _assert_loss_and_backward(output: dict, model: torch.nn.Module):
         param for param in model.parameters() if param.requires_grad and param.grad is not None
     ]
     assert grad_params
-    assert all(torch.isfinite(param.grad.detach().float()).all() for param in grad_params)
+    assert all(
+        torch.isfinite(param.grad.detach().float()).all() for param in grad_params
+    )
 
 
 def test_qwen3_moe_lite_tiny_forward_backward_smoke():
     _Qwen3MoEConfig, Qwen3MoEModel = _qwen3_symbols()
     config = _tiny_qwen3_config()
-    model = Qwen3MoEModel(config, _parallel_state(), use_deepep=False).cuda().to(torch.bfloat16)
+    model = (
+        Qwen3MoEModel(config, _parallel_state(), use_deepep=False)
+        .cuda()
+        .to(torch.bfloat16)
+    )
     input_ids, labels = _token_batch(config.vocab_size)
 
     output = model(input_ids=input_ids, labels=labels, return_log_probs=True)
@@ -146,7 +152,9 @@ def test_qwen35_lite_tiny_forward_backward_smoke():
         recompute_modules=[],
         deterministic=True,
     )
-    model = Qwen35Model(config, train_config, _parallel_state()).cuda().to(torch.bfloat16)
+    model = (
+        Qwen35Model(config, train_config, _parallel_state()).cuda().to(torch.bfloat16)
+    )
     input_ids, labels = _token_batch(config.vocab_size)
 
     output = model(input_ids=input_ids, labels=labels)
diff --git a/experimental/lite/tests/smoke/primitive/test_parallel_topologies_smoke.py b/experimental/lite/tests/smoke/primitive/test_parallel_topologies_smoke.py
index 3ee1f4e54da..1a1c0f3d224 100644
--- a/experimental/lite/tests/smoke/primitive/test_parallel_topologies_smoke.py
+++ b/experimental/lite/tests/smoke/primitive/test_parallel_topologies_smoke.py
@@ -8,7 +8,13 @@
 import torch
 import torch.distributed as dist
 
-from megatron.lite.primitive.parallel import init_parallel
+from megatron.lite.primitive.parallel import (
+    PackedSeqParams,
+    contiguous_to_zigzag_chunks,
+    init_parallel,
+    zigzag_split_for_cp,
+    zigzag_to_contiguous_chunks,
+)
 
 pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed]
 
@@ -94,3 +100,49 @@ def test_parallel_state_builds_expected_primitive_groups():
         if cfg.pp > 1:
             assert ps.pp_next_rank in ps.pp_global_ranks
             assert ps.pp_prev_rank in ps.pp_global_ranks
+
+
+def test_cp_zigzag_contiguous_chunk_swap_roundtrip():
+    if dist.get_world_size() < 2:
+        pytest.skip("CP chunk swap smoke requires at least 2 ranks.")
+
+    ps = init_parallel(SimpleNamespace(tp=1, ep=1, etp=1, cp=2, pp=1))
+    full = torch.arange(8, device="cuda", dtype=torch.float32).reshape(1, 8, 1) + 100 * ps.dp_rank
+    local_zigzag = zigzag_split_for_cp(full, ps.cp_rank, cp_size=2, seq_dim=1)
+
+    local_contiguous = zigzag_to_contiguous_chunks(local_zigzag, ps.cp_group, seq_dim=1)
+    expected = full[:, ps.cp_rank * 4 : (ps.cp_rank + 1) * 4, :]
+    assert torch.equal(local_contiguous, expected)
+
+    restored = contiguous_to_zigzag_chunks(local_contiguous, ps.cp_group, seq_dim=1)
+    assert torch.equal(restored, local_zigzag)
+
+
+def test_gdn_rejects_thd_context_parallel_until_validated():
+    if dist.get_world_size() < 2:
+        pytest.skip("GDN THD+CP guard smoke requires at least 2 ranks.")
+
+    pytest.importorskip("transformer_engine.pytorch")
+    from megatron.lite.primitive.modules.gated_delta_net import GatedDeltaNet
+
+    ps = init_parallel(SimpleNamespace(tp=1, ep=1, etp=1, cp=2, pp=1))
+    gdn = (
+        GatedDeltaNet(
+            hidden_size=16,
+            linear_num_key_heads=2,
+            linear_key_head_dim=4,
+            linear_num_value_heads=2,
+            linear_value_head_dim=4,
+            linear_conv_kernel_dim=2,
+            rms_norm_eps=1e-6,
+            ps=ps,
+        )
+        .cuda()
+        .to(torch.bfloat16)
+    )
+    cu_seqlens = torch.tensor([0, 8], dtype=torch.int32, device="cuda")
+    packed_seq_params = PackedSeqParams.from_cu_seqlens(cu_seqlens, max_seqlen=8)
+    local_hidden = torch.randn(4, 1, 16, device="cuda", dtype=torch.bfloat16)
+
+    with pytest.raises(NotImplementedError, match="all-gather CP"):
+        gdn(local_hidden, packed_seq_params=packed_seq_params)
diff --git a/experimental/lite/tests/unit/primitive/test_attention_moe_unit.py b/experimental/lite/tests/unit/primitive/test_attention_moe_unit.py
index c1e61aacba4..3945e32c283 100644
--- a/experimental/lite/tests/unit/primitive/test_attention_moe_unit.py
+++ b/experimental/lite/tests/unit/primitive/test_attention_moe_unit.py
@@ -9,6 +9,11 @@
 pytestmark = pytest.mark.mlite
 
 
+@pytest.fixture(autouse=True)
+def _te_import_stub(transformer_engine_import_stub):
+    transformer_engine_import_stub()
+
+
 def _split_grouped_qkvg():
     from megatron.lite.primitive.modules import split_grouped_qkvg
 
@@ -22,14 +27,10 @@ def _moe_aux_scaler():
 
 
 def _router_and_parallel_state(monkeypatch):
-    from megatron.core.transformer.moe import moe_utils
-
-    if not hasattr(moe_utils, "te_general_gemm"):
-        monkeypatch.setattr(moe_utils, "te_general_gemm", None, raising=False)
-
     from megatron.lite.primitive.modules.router import TopKRouter
     from megatron.lite.primitive.parallel import ParallelState
 
+    del monkeypatch
     return TopKRouter, ParallelState
 
 
diff --git a/experimental/lite/tests/unit/primitive/test_parallel_dimensions_independent_unit.py b/experimental/lite/tests/unit/primitive/test_parallel_dimensions_independent_unit.py
index 482d5449856..fe3ecfb7b94 100644
--- a/experimental/lite/tests/unit/primitive/test_parallel_dimensions_independent_unit.py
+++ b/experimental/lite/tests/unit/primitive/test_parallel_dimensions_independent_unit.py
@@ -85,11 +85,15 @@ def test_cp_packed_split_handles_each_sample_independently():
     assert rank1[3] == 4
 
 
-def test_pp_layout_rejects_non_divisible_layer_counts():
-    ps = ParallelState(pp_size=2, pp_rank=0, pp_is_first=True, pp_is_last=False)
-
-    with pytest.raises(ValueError, match="not divisible"):
-        build_pipeline_chunk_layout(7, ps)
+def test_pp_layout_auto_balances_non_divisible_layer_counts():
+    # Non-divisible counts no longer raise "not divisible": mcore's layout balances
+    # them. 7/pp2 with embedding/loss accounting -> [4, 3]. (Needs megatron.core.)
+    pytest.importorskip("megatron.core.transformer.pipeline_parallel_layer_layout")
+    rank0 = ParallelState(pp_size=2, pp_rank=0, pp_is_first=True, pp_is_last=False)
+    rank1 = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True)
+
+    assert build_pipeline_chunk_layout(7, rank0).layer_indices == [0, 1, 2, 3]
+    assert build_pipeline_chunk_layout(7, rank1).layer_indices == [4, 5, 6]
 
 
 def test_dp_dimension_controls_dense_microbatch_contract():
diff --git a/experimental/lite/tests/unit/primitive/test_parallel_unit.py b/experimental/lite/tests/unit/primitive/test_parallel_unit.py
index 6eb42553542..add7ff1e7be 100644
--- a/experimental/lite/tests/unit/primitive/test_parallel_unit.py
+++ b/experimental/lite/tests/unit/primitive/test_parallel_unit.py
@@ -3,10 +3,12 @@
 
 import pytest
 import torch
-
 from megatron.lite.primitive.parallel import (
     ParallelState,
     build_pipeline_chunk_layout,
+    pack_nested_thd,
+    parallel_state_from_model,
+    prepare_packed_thd_for_context_parallel,
     reconstruct_packed_from_cp_parts,
     roll_packed_thd_left,
     split_packed_to_cp_local,
@@ -41,10 +43,49 @@ def test_cp_position_ids_follow_zigzag_order():
     )
 
 
-def test_pp_layout_marks_stage_boundaries_and_vpp_chunks():
-    rank0 = ParallelState(pp_size=2, pp_rank=0, pp_is_first=True, pp_is_last=False)
-    rank1 = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True)
+# The pipeline layout wires to Megatron-core's layer-layout machinery, so these
+# tests need megatron.core importable (present in the mlite GPU/smoke containers).
+_mcore_layout = pytest.importorskip(
+    "megatron.core.transformer.pipeline_parallel_layer_layout"
+)
+
+
+def _ranks(pp_size: int, pp_layout=None) -> list[ParallelState]:
+    return [
+        ParallelState(
+            pp_size=pp_size,
+            pp_rank=r,
+            pp_is_first=(r == 0),
+            pp_is_last=(r == pp_size - 1),
+            pp_layout=pp_layout,
+        )
+        for r in range(pp_size)
+    ]
+
+
+def _layout_indices(num_layers: int, pp_size: int, *, pp_layout=None, **kw) -> list[list[int]]:
+    return [
+        build_pipeline_chunk_layout(num_layers, ps, **kw).layer_indices
+        for ps in _ranks(pp_size, pp_layout)
+    ]
+
 
+def _mtp_flags(num_layers: int, pp_size: int, *, pp_layout=None, **kw) -> list[bool]:
+    return [
+        build_pipeline_chunk_layout(num_layers, ps, **kw).has_mtp
+        for ps in _ranks(pp_size, pp_layout)
+    ]
+
+
+def _assert_full_contiguous_cover(indices: list[list[int]], num_layers: int) -> None:
+    """Every decoder 0..N-1 placed exactly once, in order, each stage a contiguous run."""
+    assert [i for stage in indices for i in stage] == list(range(num_layers))
+    for stage in indices:
+        assert not stage or stage == list(range(stage[0], stage[0] + len(stage)))
+
+
+def test_pp_layout_marks_stage_boundaries():
+    rank0, rank1 = _ranks(2)
     assert build_pipeline_chunk_layout(8, rank0).layer_indices == [0, 1, 2, 3]
     assert build_pipeline_chunk_layout(8, rank0).has_embed is True
     assert build_pipeline_chunk_layout(8, rank0).has_head is False
@@ -52,12 +93,158 @@ def test_pp_layout_marks_stage_boundaries_and_vpp_chunks():
     assert build_pipeline_chunk_layout(8, rank1).has_embed is False
     assert build_pipeline_chunk_layout(8, rank1).has_head is True
 
-    vpp_rank0_chunk1 = build_pipeline_chunk_layout(8, rank0, vpp=2, vpp_chunk_id=1)
-    vpp_rank1_chunk1 = build_pipeline_chunk_layout(8, rank1, vpp=2, vpp_chunk_id=1)
-    assert vpp_rank0_chunk1.layer_indices == [4, 5]
-    assert vpp_rank0_chunk1.has_head is False
-    assert vpp_rank1_chunk1.layer_indices == [6, 7]
-    assert vpp_rank1_chunk1.has_head is True
+
+def test_pp_layout_vpp_is_not_supported_yet():
+    # pp-only: requesting VPP raises rather than silently mis-splitting.
+    rank0 = ParallelState(pp_size=2, pp_rank=0, pp_is_first=True, pp_is_last=False)
+    with pytest.raises(NotImplementedError):
+        build_pipeline_chunk_layout(8, rank0, vpp=2)
+    with pytest.raises(NotImplementedError):
+        build_pipeline_chunk_layout(8, rank0, vpp=2, vpp_chunk_id=1)
+
+
+def test_pp_layout_auto_balances_non_divisible_counts():
+    # account_for_embedding/loss: the embedding/loss stages each give up a decoder,
+    # so 6/pp4 balances to [1,2,2,1] (not [2,2,1,1]) and never raises "not divisible".
+    assert _layout_indices(6, 4) == [[0], [1, 2], [3, 4], [5]]
+    assert _layout_indices(6, 4, vpp=1) == _layout_indices(6, 4)  # vpp=1 == pp-only
+    # An MTP head occupies the last stage's slot -> it gives up its decoder: [2,2,2,0].
+    assert _layout_indices(6, 4, num_mtp_layers=1) == [[0, 1], [2, 3], [4, 5], []]
+
+
+def test_pp_layout_accounts_for_head_tail_even_when_divisible():
+    # Embedding/loss occupy head/tail slots, so 8/pp4 is [2,3,2,1] not [2,2,2,2].
+    assert _layout_indices(8, 4) == [[0, 1], [2, 3, 4], [5, 6], [7]]
+    assert _layout_indices(8, 4, num_mtp_layers=1) == [[0, 1], [2, 3, 4], [5, 6, 7], []]
+    assert _layout_indices(8, 2) == [[0, 1, 2, 3], [4, 5, 6, 7]]
+
+
+# --- Correctness matrix: real delivery counts x MTP{off,on} x PP{2,4,8} ---
+_REAL_LAYERS = {"deepseek_v4": 43, "kimi_k2": 61, "glm5": 78}
+_CORRECTNESS_CASES = [(m, pp, mtp) for m in _REAL_LAYERS for pp in (2, 4, 8) for mtp in (0, 1)]
+
+
+def _mcore_reference_decoder_ids(num_layers: int, pp: int, mtp: int) -> list[list[int]]:
+    """Per-stage decoder ids from mcore's PipelineParallelLayerLayout, built directly
+    from the canonical unit sequence — the independent reference the glue must match."""
+    from megatron.core.transformer.enums import LayerType
+    from megatron.core.transformer.pipeline_parallel_layer_layout import (
+        PipelineParallelLayerLayout,
+    )
+
+    units = ["embedding"] + ["decoder"] * num_layers + ["mtp"] * mtp + ["loss"]
+    base, rem = divmod(len(units), pp)
+    rows, pos = [], 0
+    for size in (base + (1 if s < rem else 0) for s in range(pp)):
+        rows.append(units[pos : pos + size])
+        pos += size
+    ref = PipelineParallelLayerLayout(rows, pipeline_model_parallel_size=pp)
+    return [ref.get_layer_id_list(LayerType.decoder, vp_stage=0, pp_rank=r) for r in range(pp)]
+
+
+@pytest.mark.parametrize("model, pp, mtp", _CORRECTNESS_CASES)
+def test_auto_layout_is_bit_equal_to_mcore(model, pp, mtp):
+    # Faithful reuse: the glue's per-stage decoder ids are exactly mcore's for the same
+    # canonical layout, so correctness is inherited from mcore (not re-derived).
+    n = _REAL_LAYERS[model]
+    assert _layout_indices(n, pp, num_mtp_layers=mtp) == _mcore_reference_decoder_ids(n, pp, mtp)
+
+
+@pytest.mark.parametrize("model, pp, mtp", _CORRECTNESS_CASES)
+def test_auto_layout_is_legal_and_balanced(model, pp, mtp):
+    n = _REAL_LAYERS[model]
+    chunks = [build_pipeline_chunk_layout(n, ps, num_mtp_layers=mtp) for ps in _ranks(pp)]
+    # complete + ordered, no missing/dup decoder -> sum == num_layers
+    assert [i for c in chunks for i in c.layer_indices] == list(range(n))
+    # embedding only on stage 0; loss/head only on the last stage
+    assert [c.has_embed for c in chunks] == [r == 0 for r in range(pp)]
+    assert [c.has_head for c in chunks] == [r == pp - 1 for r in range(pp)]
+    # MTP placed exactly `mtp` times, on the final (head) stage
+    assert sum(c.has_mtp for c in chunks) == mtp and chunks[-1].has_mtp == bool(mtp)
+    # balanced: per-stage unit cells (decoders + embedding/loss/mtp slots) differ by <=1
+    cells = [len(c.layer_indices) + c.has_embed + c.has_head + (mtp if c.has_mtp else 0) for c in chunks]
+    assert max(cells) - min(cells) <= 1
+
+
+def test_pp_layout_custom_string_mode_is_used_verbatim():
+    # Advanced users pin an explicit mcore layout string (custom mode) instead of the
+    # auto split. "Ettt|t|t|tL" front-loads 3 decoders onto stage 0; mcore parses
+    # E=embedding t=decoder L=loss. This is NOT the auto [1,2,2,1] for 6/pp4.
+    custom = _layout_indices(6, 4, pp_layout="Ettt|t|t|tL")
+    assert [len(stage) for stage in custom] == [3, 1, 1, 1]
+    _assert_full_contiguous_cover(custom, 6)
+    assert custom != _layout_indices(6, 4)  # custom overrides auto
+
+    # A custom string carrying MTP after the decoders validates against num_mtp_layers.
+    mtp_custom = _layout_indices(6, 4, pp_layout="Ett|tt|t|tmL", num_mtp_layers=1)
+    assert [len(stage) for stage in mtp_custom] == [2, 2, 1, 1]
+    _assert_full_contiguous_cover(mtp_custom, 6)
+
+
+def test_pp_layout_custom_string_rejects_vpp_multi_chunk():
+    # A custom layout with more stages than pp implies VPP, which is not supported.
+    with pytest.raises(NotImplementedError):
+        _layout_indices(6, 2, pp_layout="Et|t|t|ttL")  # 4 stages over pp2 -> vpp=2
+
+
+def test_pp_layout_marks_mtp_stage_from_the_layout_not_a_fixed_rank():
+    # auto: the MTP slot sits just before loss, so it lands on the final (head) stage
+    # -- has_mtp marks exactly that stage, and shifts the decoder split [2,2,2,0].
+    assert _mtp_flags(6, 4, num_mtp_layers=1) == [False, False, False, True]
+    assert _layout_indices(6, 4, num_mtp_layers=1) == [[0, 1], [2, 3], [4, 5], []]
+    # no MTP requested -> no stage owns MTP.
+    assert _mtp_flags(6, 4) == [False, False, False, False]
+    # single stage owns the MTP too.
+    assert build_pipeline_chunk_layout(
+        5, ParallelState(pp_size=1, pp_rank=0, pp_is_first=True, pp_is_last=True), num_mtp_layers=1
+    ).has_mtp is True
+
+
+def test_pp_layout_custom_string_mtp_lands_on_the_designated_stage():
+    # The `m` token co-located with loss on the final stage: has_mtp marks that stage
+    # (== the head stage), driven by the layout, not a hard-coded rank.
+    flags = _mtp_flags(6, 4, pp_layout="Ett|tt|t|tmL", num_mtp_layers=1)
+    assert flags == [False, False, False, True]
+    assert _layout_indices(6, 4, pp_layout="Ett|tt|t|tmL", num_mtp_layers=1) == [
+        [0, 1], [2, 3], [4], [5]
+    ]
+
+
+@pytest.mark.xfail(
+    reason="Standalone MTP (a custom pp_layout placing `m` off the final/head stage) is a "
+    "planned follow-up: mlite's MTP shares the output head on the loss stage, so it "
+    "currently raises NotImplementedError instead of building MTP on the `m` stage. "
+    "When cross-stage standalone MTP lands, this should pass and the marker be removed.",
+    strict=True,
+    raises=NotImplementedError,
+)
+def test_pp_layout_custom_string_standalone_mtp_builds_on_designated_stage():
+    # DESIRED (follow-up): `m` on a non-final stage builds MTP on that stage.
+    # "E|ttttttm|L" over pp3 -> [E], [t*6, m], [L]; MTP belongs on the middle stage.
+    assert _mtp_flags(6, 3, pp_layout="E|ttttttm|L", num_mtp_layers=1) == [False, True, False]
+
+
+def test_pp_layout_single_stage_owns_all_layers():
+    ps = ParallelState(pp_size=1, pp_rank=0, pp_is_first=True, pp_is_last=True)
+    layout = build_pipeline_chunk_layout(5, ps)
+    assert layout.layer_indices == [0, 1, 2, 3, 4]
+    assert layout.has_embed and layout.has_head
+
+
+def test_virtual_pipeline_rank_is_tracked_on_lite_parallel_state():
+    from megatron.lite.primitive.parallel.pipeline import _set_virtual_pipeline_rank
+
+    ps = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True)
+
+    _set_virtual_pipeline_rank(ps, chunk_id=1, num_chunks=2)
+
+    assert ps.virtual_pipeline_size == 2
+    assert ps.virtual_pipeline_rank == 1
+
+    _set_virtual_pipeline_rank(ps, chunk_id=None, num_chunks=2)
+
+    assert ps.virtual_pipeline_size is None
+    assert ps.virtual_pipeline_rank is None
 
 
 def test_thd_roll_keeps_sequence_boundaries():
@@ -84,3 +271,98 @@ def test_thd_cp_split_and_reconstruct_roundtrip():
         reconstruct_packed_from_cp_parts(parts, cu_seqlens_padded=cu_seqlens, cp_size=2, dim=0),
         tensor,
     )
+
+
+def test_plain_thd_batch_is_split_by_protocol_context_parallel_helper():
+    ids = torch.nested.as_nested_tensor(
+        [torch.arange(1, 6), torch.arange(11, 18)],
+        layout=torch.jagged,
+    )
+    labels = torch.nested.as_nested_tensor(
+        [torch.arange(101, 106), torch.arange(111, 118)],
+        layout=torch.jagged,
+    )
+    loss_mask = torch.nested.as_nested_tensor(
+        [torch.ones(5), torch.ones(7)],
+        layout=torch.jagged,
+    )
+    packed = pack_nested_thd(
+        ids,
+        cp_size=2,
+        split_cp=False,
+        labels=labels,
+        loss_mask=loss_mask,
+    )
+
+    assert packed.input_ids.shape == (1, 16)
+    assert packed.cp_size == 2
+    assert packed.packed_seq_params.local_cp_size is None
+
+    local_params, local_tensors = prepare_packed_thd_for_context_parallel(
+        packed.packed_seq_params,
+        (packed.input_ids, packed.labels, packed.loss_mask, packed.position_ids),
+        cp_size=2,
+        cp_rank=0,
+    )
+
+    expected_ids = split_packed_to_cp_local(
+        packed.input_ids,
+        cu_seqlens_padded=packed.cu_seqlens_padded,
+        cp_size=2,
+        cp_rank=0,
+        dim=1,
+    )
+    expected_pos = split_packed_to_cp_local(
+        packed.position_ids,
+        cu_seqlens_padded=packed.cu_seqlens_padded,
+        cp_size=2,
+        cp_rank=0,
+        dim=1,
+    )
+    local_ids, local_labels, local_loss_mask, local_pos = local_tensors
+    assert torch.equal(local_ids, expected_ids)
+    assert torch.equal(local_pos, expected_pos)
+    assert local_labels is not None
+    assert local_loss_mask is not None
+    assert local_params.local_cp_size == 2
+    assert local_params.cp_rank == 0
+
+
+def test_parallel_state_from_model_unwraps_ddp_style_module():
+    class Model:
+        ps = ParallelState(cp_size=2, cp_rank=1)
+
+    class Wrapper:
+        module = Model()
+
+    assert parallel_state_from_model(Wrapper()).cp_rank == 1
+
+
+def test_protocol_context_parallel_helper_keeps_pre_split_thd_batch_idempotent():
+    ids = torch.nested.as_nested_tensor([torch.arange(8)], layout=torch.jagged)
+    packed = pack_nested_thd(ids, cp_size=2, cp_rank=1)
+
+    local_params, local_tensors = prepare_packed_thd_for_context_parallel(
+        packed.packed_seq_params,
+        (packed.input_ids, packed.position_ids),
+        cp_size=2,
+        cp_rank=1,
+    )
+
+    assert local_params is packed.packed_seq_params
+    assert torch.equal(local_tensors[0], packed.input_ids)
+    assert local_params.local_cp_size == 2
+
+
+def test_protocol_context_parallel_helper_is_noop_without_packed_thd_params():
+    tensor = torch.arange(8)
+
+    local_params, local_tensors = prepare_packed_thd_for_context_parallel(
+        None,
+        (tensor,),
+        cp_size=2,
+        cp_rank=0,
+    )
+
+    assert local_params is None
+    assert local_tensors[0] is tensor