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 1/4] [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

From ae4e853e467999116b6c6301901caff05cd4b64c Mon Sep 17 00:00:00 2001
From: Yan Bai <bayan@nvidia.com>
Date: Wed, 24 Jun 2026 07:51:28 -0700
Subject: [PATCH 2/4] [dev] Megatron Lite (2/4): runtime layering guard
 (variable-leakage governance)

Signed-off-by: Yan Bai <bayan@nvidia.com>
---
 experimental/lite/examples/bench/bench.py     |  17 +-
 .../lite/examples/bench/correctness.py        |  54 +--
 experimental/lite/examples/bench/session.py   |  42 ++-
 .../verl/scripts/run_qwen3moe_gsm8k_grpo.sh   |  11 +-
 .../examples/verl/scripts/run_qwen3moe_sft.sh |  22 +-
 .../lite/examples/verl/verl_mlite/compat.py   |  40 +++
 .../verl/verl_mlite/config/engine/mlite.yaml  |   1 +
 .../examples/verl/verl_mlite/engine/config.py |   2 +
 .../verl/verl_mlite/engine/mlite_engine.py    | 339 +++++++++++-------
 .../lite/examples/verl/verl_mlite/launch.py   |  26 ++
 .../lite/model/qwen3_5/lite/protocol.py       |  74 ++--
 .../lite/model/qwen3_moe/lite/protocol.py     |  49 +--
 .../primitive/optimizers/megatron_wrap.py     | 118 +++---
 .../lite/runtime/backends/bridge/config.py    |   5 +
 .../lite/runtime/backends/bridge/runtime.py   | 162 ++++++++-
 .../lite/runtime/backends/mbridge/runtime.py  |   2 +-
 .../lite/runtime/backends/mlite/runtime.py    |  54 ++-
 .../megatron/lite/runtime/megatron_utils.py   |  22 +-
 experimental/lite/tests/conftest.py           |  37 +-
 .../lite/tests/run_layering_contracts.sh      |   8 +
 .../test_distopt_checkpoint_smoke.py          |  72 ++--
 ...test_qwen3_moe_distopt_checkpoint_smoke.py |  46 ++-
 .../unit/primitive/test_checkpoint_runtime.py |   6 +-
 .../unit/primitive/test_checkpoint_unit.py    |   6 +-
 .../primitive/test_dist_opt_validation.py     |  15 +-
 .../primitive/test_training_checkpoint.py     |  19 +-
 .../unit/runtime/test_layering_contracts.py   | 235 ++++++++++++
 .../unit/runtime/test_packed_batch_bridge.py  | 142 ++++++++
 .../unit/runtime/test_runtime_backend_unit.py | 147 ++++++++
 .../unit/verl/test_mlite_engine_config.py     |  37 +-
 30 files changed, 1385 insertions(+), 425 deletions(-)
 create mode 100644 experimental/lite/examples/verl/verl_mlite/launch.py
 create mode 100755 experimental/lite/tests/run_layering_contracts.sh
 create mode 100644 experimental/lite/tests/unit/runtime/test_layering_contracts.py
 create mode 100644 experimental/lite/tests/unit/runtime/test_packed_batch_bridge.py

diff --git a/experimental/lite/examples/bench/bench.py b/experimental/lite/examples/bench/bench.py
index be6192c9c20..709057388e0 100644
--- a/experimental/lite/examples/bench/bench.py
+++ b/experimental/lite/examples/bench/bench.py
@@ -156,18 +156,18 @@ def keep_experts_hook(model_cfg):
 
         def truncate_layers_hook(model_cfg):
             old_layers = getattr(model_cfg, "num_hidden_layers", None)
-            layer_types = getattr(model_cfg, "layer_types", None)
-            if old_layers is None or layer_types is None:
-                raise ValueError("truncate_layers requires num_hidden_layers and layer_types.")
+            if old_layers is None:
+                raise ValueError("truncate_layers requires num_hidden_layers.")
             if keep_layers <= 0 or keep_layers > old_layers:
                 raise ValueError(
                     f"truncate_layers must be in [1, {old_layers}], got {keep_layers}."
                 )
-            return replace(
-                model_cfg,
-                num_hidden_layers=keep_layers,
-                layer_types=list(layer_types[:keep_layers]),
-            )
+            updates = {"num_hidden_layers": keep_layers}
+            # layer_types is qwen-style metadata; deepseek_v4 / kimi_k2 configs don't carry it.
+            layer_types = getattr(model_cfg, "layer_types", None)
+            if layer_types is not None:
+                updates["layer_types"] = list(layer_types[:keep_layers])
+            return replace(model_cfg, **updates)
 
         hooks.append(truncate_layers_hook)
 
@@ -289,6 +289,7 @@ def build_runtime_config(cfg: BenchCliConfig) -> RuntimeConfig:
             model_name=cfg.model_name,
             parallel=parallel,
             optimizer=optimizer,
+            use_thd=cfg.use_thd,
             load_hf_weights=not cfg.skip_load_hf_weights,
             build_optimizer=not cfg.skip_optimizer_build,
             override_ddp_config=_json_mapping(cfg.override_ddp_json, name="override_ddp_json"),
diff --git a/experimental/lite/examples/bench/correctness.py b/experimental/lite/examples/bench/correctness.py
index fd534fdc05f..31199d9bd50 100644
--- a/experimental/lite/examples/bench/correctness.py
+++ b/experimental/lite/examples/bench/correctness.py
@@ -22,11 +22,12 @@
 if str(_REPO_ROOT) not in sys.path:
     sys.path.insert(1, str(_REPO_ROOT))
 
+from megatron.lite.primitive.deterministic import set_deterministic
+from megatron.lite.runtime import create_runtime
+
 from examples.bench.bench import BenchCliConfig, build_runtime_config, build_session_config
 from examples.bench.results import compare_correctness_artifacts, load_result_artifact
 from examples.bench.session import _make_data_iter
-from megatron.lite.primitive.deterministic import set_deterministic
-from megatron.lite.runtime import create_runtime
 
 
 def _distributed_rank() -> int:
@@ -299,48 +300,15 @@ def _weight_fingerprint(rt, handle) -> dict[str, Any]:
     return result
 
 
-def _batch_without_labels(batch: Any) -> dict[str, Any]:
-    if not isinstance(batch, dict):
-        return {
-            "input_ids": batch["input_ids"],
-            "position_ids": getattr(batch, "position_ids", None),
-            "packed_seq_params": getattr(batch, "packed_seq_params", None),
-        }
-    return {k: v for k, v in batch.items() if k != "labels"}
-
-
 def _forward_logits(rt, handle, batch: Any) -> torch.Tensor | None:
-    sample = _batch_without_labels(batch)
-    if "forward_step" in handle._extras:
-        try:
-            out = handle._model(**sample)
-        except (KeyError, TypeError):
-            out = handle._extras["forward_step"](handle._model, sample)
-        if isinstance(out, dict):
-            logits = out.get("logits")
-            if logits is not None:
-                return logits
-            return out.get("vocab_parallel_logits")
-        return out if isinstance(out, torch.Tensor) else None
-
-    model_list = handle._extras.get("model_list")
-    if model_list:
-        out = model_list[0](
-            input_ids=sample.get("input_ids"),
-            position_ids=sample.get("position_ids"),
-            attention_mask=sample.get("attention_mask"),
-            packed_seq_params=sample.get("packed_seq_params"),
-        )
-        if isinstance(out, tuple):
-            out = out[0]
-        if isinstance(out, dict):
-            logits = out.get("logits")
-            if logits is not None:
-                return logits
-            return out.get("vocab_parallel_logits")
-        return out if isinstance(out, torch.Tensor) else None
-
-    return None
+    result = rt.forward_backward(
+        handle,
+        iter([batch]),
+        loss_fn=None,
+        num_microbatches=1,
+        forward_only=True,
+    )
+    return result.model_output.vocab_parallel_logits
 
 
 def run_backend(
diff --git a/experimental/lite/examples/bench/session.py b/experimental/lite/examples/bench/session.py
index 47e7a25789a..75e76f16ca5 100644
--- a/experimental/lite/examples/bench/session.py
+++ b/experimental/lite/examples/bench/session.py
@@ -10,7 +10,6 @@
 from typing import Any
 
 import torch
-
 from megatron.lite.runtime.backends import Runtime
 from megatron.lite.runtime.contracts.handle import ModelHandle
 
@@ -66,26 +65,33 @@ def _resolve_vocab_size(handle: ModelHandle) -> int:
     return 151936
 
 
-def _make_data_iter(handle: ModelHandle, cfg: PretrainSessionConfig):
-    data_seed = cfg.seed if cfg.same_data_across_dp else cfg.seed + handle.dp_rank
-    vocab_size = _resolve_vocab_size(handle)
-
-    if cfg.use_thd:
-        from megatron.lite.primitive.data import infinite_batches_thd
-
-        ps = handle._parallel_state
-        return infinite_batches_thd(
-            vocab_size,
-            cfg.seq_len,
-            cp_size=getattr(ps, "cp_size", 1),
-            cp_rank=getattr(ps, "cp_rank", 0),
-            device=cfg.device,
-            seed=data_seed,
+def _infinite_packed_batches(
+    vocab_size: int, seq_len: int, *, device: str, seed: int
+):
+    """Yield raw, model-agnostic :class:`PackedBatch` objects for the bench.
+
+    The bench is the single source of truth for one unpadded packed batch (1-D
+    ``input_ids``/``labels`` plus true per-sequence ``seq_lens``). Padding, CP
+    layout and THD metadata (``packed_seq_params``) are derived by whichever
+    runtime/model consumes the batch at the forward boundary, never baked into
+    bench data — that is what keeps the mlite-vs-bridge comparison fair.
+    """
+    from megatron.lite.runtime.contracts.data import PackedBatch
+
+    g = torch.Generator(device=device).manual_seed(seed)
+    seq_lens = torch.tensor([seq_len], dtype=torch.int64, device=device)
+    while True:
+        yield PackedBatch(
+            input_ids=torch.randint(0, vocab_size, (seq_len,), device=device, generator=g),
+            labels=torch.randint(0, vocab_size, (seq_len,), device=device, generator=g),
+            seq_lens=seq_lens.clone(),
         )
 
-    from megatron.lite.primitive.data import infinite_batches
 
-    return infinite_batches(vocab_size, cfg.seq_len, device=cfg.device, seed=data_seed)
+def _make_data_iter(handle: ModelHandle, cfg: PretrainSessionConfig):
+    data_seed = cfg.seed if cfg.same_data_across_dp else cfg.seed + handle.dp_rank
+    vocab_size = _resolve_vocab_size(handle)
+    return _infinite_packed_batches(vocab_size, cfg.seq_len, device=cfg.device, seed=data_seed)
 
 
 def _calc_tflops_per_gpu(
diff --git a/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh b/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh
index c20f1464db6..01baa935446 100755
--- a/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh
+++ b/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh
@@ -85,9 +85,9 @@ MLITE_MODEL_NAME="${MLITE_MODEL_NAME:-auto}"
 MLITE_IMPL="${MLITE_IMPL:-lite}"
 ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}"
 # Optimizer backend:
-# - distopt (default): Megatron-Core DDP + distributed optimizer.
+# - dist_opt (default): Megatron-Core DDP + distributed optimizer.
 # - fsdp2: Megatron Lite FSDP2 wrapper + optimizer.
-MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-distopt}"
+MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-dist_opt}"
 
 ACTOR_LR="${ACTOR_LR:-1e-6}"
 POLICY_LOSS_MODE="${POLICY_LOSS_MODE:-vanilla}"
@@ -124,14 +124,14 @@ if [[ "${INFER_BACKEND}" != "vllm" && "${INFER_BACKEND}" != "sglang" && "${INFER
 fi
 
 case "${MLITE_OPTIMIZER_BACKEND}" in
-  distopt)
-    MLITE_IMPL_OPTIMIZER="mc"
+  dist_opt)
+    MLITE_IMPL_OPTIMIZER="dist_opt"
     ;;
   fsdp2)
     MLITE_IMPL_OPTIMIZER="fsdp2"
     ;;
   *)
-    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected distopt or fsdp2." >&2
+    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected dist_opt or fsdp2." >&2
     exit 1
     ;;
 esac
@@ -184,7 +184,6 @@ DATA=(
 MODEL=(
   "actor_rollout_ref.model.path=${MODEL_PATH}"
   "actor_rollout_ref.model.trust_remote_code=True"
-  "actor_rollout_ref.model.use_remove_padding=True"
   "actor_rollout_ref.model.use_fused_kernels=False"
 )
 
diff --git a/experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh b/experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh
index eb4c903b539..8ef6d15179e 100755
--- a/experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh
+++ b/experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh
@@ -24,6 +24,10 @@ add_pythonpath "${VERL_ROOT:-}"
 add_pythonpath "${MEGATRON_ROOT:-}"
 
 export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}"
+if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
+  unset ROCR_VISIBLE_DEVICES
+  unset HIP_VISIBLE_DEVICES
+fi
 
 : "${MODEL_PATH:?set MODEL_PATH to a Hugging Face checkpoint directory or model id}"
 : "${TRAIN_FILES:?set TRAIN_FILES to a messages parquet path or comma-separated parquet paths}"
@@ -46,13 +50,14 @@ SAVE_FREQ="${SAVE_FREQ:-${TOTAL_STEPS}}"
 TEST_FREQ="${TEST_FREQ:--1}"
 RESUME_MODE="${RESUME_MODE:-disable}"
 RESUME_FROM_PATH="${RESUME_FROM_PATH:-null}"
+CHECKPOINT_SAVE_CONTENTS="${CHECKPOINT_SAVE_CONTENTS:-[model,optimizer,extra]}"
+LOAD_HF_WEIGHTS="${LOAD_HF_WEIGHTS:-True}"
 TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-64}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
 MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}"
 MAX_LENGTH="${MAX_LENGTH:-${MAX_TOKENS_PER_GPU}}"
 PAD_MODE="${PAD_MODE:-no_padding}"
 USE_DYNAMIC_BSZ="${USE_DYNAMIC_BSZ:-True}"
-USE_REMOVE_PADDING="${USE_REMOVE_PADDING:-True}"
 IGNORE_INPUT_IDS_MISMATCH="${IGNORE_INPUT_IDS_MISMATCH:-True}"
 TRUST_REMOTE_CODE="${TRUST_REMOTE_CODE:-True}"
 MESSAGES_KEY="${MESSAGES_KEY:-messages}"
@@ -70,9 +75,9 @@ MLITE_MODEL_NAME="${MLITE_MODEL_NAME:-auto}"
 MLITE_IMPL="${MLITE_IMPL:-lite}"
 ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}"
 # Optimizer backend:
-# - distopt (default): Megatron-Core DDP + distributed optimizer.
+# - dist_opt (default): Megatron-Core DDP + distributed optimizer.
 # - fsdp2: Megatron Lite FSDP2 wrapper + optimizer.
-MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-distopt}"
+MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-dist_opt}"
 
 LR="${LR:-1e-5}"
 MIN_LR="${MIN_LR:-${LR}}"
@@ -102,14 +107,14 @@ if [[ "${PAD_MODE}" != "no_padding" ]]; then
 fi
 
 case "${MLITE_OPTIMIZER_BACKEND}" in
-  distopt)
-    MLITE_IMPL_OPTIMIZER="mc"
+  dist_opt)
+    MLITE_IMPL_OPTIMIZER="dist_opt"
     ;;
   fsdp2)
     MLITE_IMPL_OPTIMIZER="fsdp2"
     ;;
   *)
-    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected distopt or fsdp2." >&2
+    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected dist_opt or fsdp2." >&2
     exit 1
     ;;
 esac
@@ -149,7 +154,6 @@ COMMON_ARGS=(
   "model=hf_model"
   "model.path=${MODEL_PATH}"
   "model.trust_remote_code=${TRUST_REMOTE_CODE}"
-  "model.use_remove_padding=${USE_REMOVE_PADDING}"
   "optim=megatron"
   "optim.lr=${LR}"
   "optim.min_lr=${MIN_LR}"
@@ -172,7 +176,7 @@ COMMON_ARGS=(
   "trainer.resume_from_path=${RESUME_FROM_PATH}"
   "trainer.nnodes=${NNODES}"
   "trainer.n_gpus_per_node=${NPROC_PER_NODE}"
-  "checkpoint.save_contents=[model,optimizer,extra]"
+  "checkpoint.save_contents=${CHECKPOINT_SAVE_CONTENTS}"
 )
 
 if [[ -n "${VAL_FILES}" ]]; then
@@ -195,6 +199,7 @@ BACKEND_ARGS=(
   "engine.optimizer_offload=${OPTIMIZER_OFFLOAD}"
   "engine.grad_offload=${GRAD_OFFLOAD}"
   "engine.attention_backend_override=${ATTENTION_BACKEND}"
+  "engine.load_hf_weights=${LOAD_HF_WEIGHTS}"
   "engine.impl_cfg.use_thd=True"
   "+engine.impl_cfg.optimizer=${MLITE_IMPL_OPTIMIZER}"
 )
@@ -215,6 +220,7 @@ COMMAND=(
   --master_port="${MASTER_PORT}"
   --nproc_per_node="${NPROC_PER_NODE}"
   -m
+  verl_mlite.launch
   verl.trainer.sft_trainer
   "${COMMON_ARGS[@]}"
   "${BACKEND_ARGS[@]}"
diff --git a/experimental/lite/examples/verl/verl_mlite/compat.py b/experimental/lite/examples/verl/verl_mlite/compat.py
index bb5ac89ba97..f49d6f50949 100644
--- a/experimental/lite/examples/verl/verl_mlite/compat.py
+++ b/experimental/lite/examples/verl/verl_mlite/compat.py
@@ -3,8 +3,11 @@
 
 from __future__ import annotations
 
+import importlib.util
+import sys
 from collections.abc import Iterable
 from functools import wraps
+from pathlib import Path
 from typing import Any
 
 
@@ -54,3 +57,40 @@ def patched(*args: Any, **kwargs: Any) -> Any:
 
 def apply_runtime_patches() -> None:
     _patch_transformers_rope_ignore_keys()
+
+
+def _load_verl_file(relative_path: str, module_name: str):
+    spec = importlib.util.find_spec("verl")
+    if spec is None or spec.submodule_search_locations is None:
+        raise ModuleNotFoundError("No module named 'verl'")
+
+    path = Path(next(iter(spec.submodule_search_locations))) / relative_path
+    file_spec = importlib.util.spec_from_file_location(module_name, path)
+    if file_spec is None or file_spec.loader is None:
+        raise ImportError(f"Unable to load VERL module from {path}")
+
+    module = importlib.util.module_from_spec(file_spec)
+    sys.modules[module_name] = module
+    file_spec.loader.exec_module(module)
+    return module
+
+
+def load_verl_engine_api():
+    # Prefer the canonical package import so the MLite engine registers into the
+    # SAME EngineRegistry that verl's trainers resolve against. Loading base.py as
+    # a standalone module (below) creates a *duplicate* registry, which silently
+    # drops the mlite backend ("Unknown backend: mlite"). The file-load path is
+    # only a fallback for environments where verl isn't importable as a package.
+    try:
+        from verl.workers.engine.base import BaseEngine, BaseEngineCtx, EngineRegistry
+        from verl.workers.engine.utils import postprocess_batch_func, prepare_micro_batches
+    except (ModuleNotFoundError, ImportError):
+        base = _load_verl_file("workers/engine/base.py", "_verl_mlite_verl_engine_base")
+        utils = _load_verl_file("workers/engine/utils.py", "_verl_mlite_verl_engine_utils")
+        BaseEngine = base.BaseEngine
+        BaseEngineCtx = base.BaseEngineCtx
+        EngineRegistry = base.EngineRegistry
+        postprocess_batch_func = utils.postprocess_batch_func
+        prepare_micro_batches = utils.prepare_micro_batches
+
+    return BaseEngine, BaseEngineCtx, EngineRegistry, postprocess_batch_func, prepare_micro_batches
diff --git a/experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml b/experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml
index 7e3e4814239..dbdd132feb2 100644
--- a/experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml
+++ b/experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml
@@ -8,6 +8,7 @@ grad_offload: false
 forward_only: false
 dtype: bfloat16
 export_dtype: bfloat16
+load_hf_weights: true
 
 model_name: auto
 impl: lite
diff --git a/experimental/lite/examples/verl/verl_mlite/engine/config.py b/experimental/lite/examples/verl/verl_mlite/engine/config.py
index fa2bb461ff4..4ee184b752b 100644
--- a/experimental/lite/examples/verl/verl_mlite/engine/config.py
+++ b/experimental/lite/examples/verl/verl_mlite/engine/config.py
@@ -28,7 +28,9 @@ class MegatronLiteEngineConfig(EngineConfig):
 
     attention_backend_override: str | None = "flash"
     router_aux_loss_coef: float | None = None
+    cross_entropy_fusion: bool | None = None
     export_dtype: str | None = "bfloat16"
+    load_hf_weights: bool = True
     impl_cfg: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
diff --git a/experimental/lite/examples/verl/verl_mlite/engine/mlite_engine.py b/experimental/lite/examples/verl/verl_mlite/engine/mlite_engine.py
index 32e81f23fb3..561bf1526b8 100644
--- a/experimental/lite/examples/verl/verl_mlite/engine/mlite_engine.py
+++ b/experimental/lite/examples/verl/verl_mlite/engine/mlite_engine.py
@@ -3,31 +3,43 @@
 
 from __future__ import annotations
 
+import math
 import os
+from enum import Enum
 from typing import Any
 
 import torch
 import torch.distributed as dist
-from tensordict import TensorDict
-from verl.trainer.config import CheckpointConfig
-from verl.utils import tensordict_utils as tu
-from verl.utils.dataset.dataset_utils import DatasetPadMode
-from verl.utils.device import get_device_id, get_device_name
-from verl.workers.config import HFModelConfig, OptimizerConfig
-from verl.workers.engine.base import BaseEngine, BaseEngineCtx, EngineRegistry
-from verl.workers.engine.utils import postprocess_batch_func, prepare_micro_batches
-
 from megatron.lite.model import resolve_model_type_from_hf
 from megatron.lite.primitive.ckpt import load_training_checkpoint, save_training_checkpoint
-from megatron.lite.primitive.parallel import pack_nested_thd, unpack_packed_thd_to_nested
 from megatron.lite.primitive.protocols import default_expert_classifier, default_placement_fn
 from megatron.lite.runtime import create_runtime
 from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig
+from megatron.lite.runtime.contracts import LossContext, PackedBatch
 from megatron.lite.runtime.contracts.config import OptimizerConfig as MegatronLiteOptimizerConfig
 from megatron.lite.runtime.contracts.config import ParallelConfig, RuntimeConfig
+from tensordict import TensorDict
+
+from verl.trainer.config import CheckpointConfig
+from verl.utils import tensordict_utils as tu
+from verl.utils.device import get_device_id, get_device_name
+from verl.workers.config import HFModelConfig, OptimizerConfig
+from verl_mlite.compat import load_verl_engine_api
 
 from .config import MegatronLiteEngineConfig
 
+BaseEngine, BaseEngineCtx, EngineRegistry, postprocess_batch_func, prepare_micro_batches = (
+    load_verl_engine_api()
+)
+
+try:
+    from verl.utils.dataset.dataset_utils import DatasetPadMode
+except ImportError:
+
+    class DatasetPadMode(Enum):
+        NO_PADDING = "no_padding"
+
+
 _LR_SCHEDULER_STATE = "lr_scheduler.pt"
 
 
@@ -47,14 +59,129 @@ def _isolate_compile_cache_per_rank() -> None:
         os.environ[var] = rank_dir
 
 
+def _is_no_padding_pad_mode(pad_mode: Any) -> bool:
+    return (
+        pad_mode == DatasetPadMode.NO_PADDING
+        or getattr(pad_mode, "name", None) == "NO_PADDING"
+        or getattr(pad_mode, "value", None) == "no_padding"
+        or str(pad_mode) in {"no_padding", "DatasetPadMode.NO_PADDING"}
+    )
+
+
+class _MegatronLiteLRScheduler:
+    def __init__(
+        self,
+        optimizer,
+        *,
+        init_lr: float,
+        max_lr: float,
+        min_lr: float,
+        lr_warmup_steps: int,
+        lr_decay_steps: int,
+        lr_decay_style: str,
+        start_wd: float,
+        end_wd: float,
+        wd_incr_steps: int,
+        wd_incr_style: str,
+        wsd_decay_steps: int | None,
+        lr_wsd_decay_style: str,
+    ):
+        self.optimizer = optimizer
+        self.init_lr = init_lr
+        self.max_lr = max_lr
+        self.min_lr = min_lr
+        self.lr_warmup_steps = max(lr_warmup_steps, 0)
+        self.lr_decay_steps = max(lr_decay_steps, self.lr_warmup_steps + 1)
+        self.lr_decay_style = lr_decay_style.lower()
+        self.start_wd = start_wd
+        self.end_wd = end_wd
+        self.wd_incr_steps = max(wd_incr_steps, 1)
+        self.wd_incr_style = wd_incr_style.lower()
+        self.wsd_decay_steps = wsd_decay_steps
+        self.lr_wsd_decay_style = lr_wsd_decay_style.lower()
+        self.num_steps = 0
+        self._apply()
+
+    def state_dict(self) -> dict[str, Any]:
+        return {"num_steps": self.num_steps}
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        self.num_steps = int(state.get("num_steps", state.get("step", 0)))
+        self._apply()
+
+    def step(self, increment: int = 1) -> None:
+        self.num_steps += increment
+        self._apply()
+
+    def get_last_lr(self) -> list[float]:
+        return [group["lr"] for group in self.optimizer.param_groups]
+
+    def _apply(self) -> None:
+        lr = self._get_lr()
+        wd = self._get_wd()
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] = lr
+            if param_group.get("weight_decay", None) is not None:
+                param_group["weight_decay"] = wd
+
+    def _get_lr(self) -> float:
+        if self.lr_warmup_steps > 0 and self.num_steps <= self.lr_warmup_steps:
+            ratio = self.num_steps / self.lr_warmup_steps
+            return self.init_lr + (self.max_lr - self.init_lr) * ratio
+
+        if self.lr_decay_style == "constant":
+            return self.max_lr
+
+        if self.lr_decay_style == "inverse-square-root":
+            warmup = max(self.lr_warmup_steps, 1)
+            step = max(self.num_steps, 1)
+            return max(self.min_lr, self.max_lr * math.sqrt(warmup) / math.sqrt(step))
+
+        if self.lr_decay_style == "wsd":
+            return self._get_wsd_lr()
+
+        decay_span = max(self.lr_decay_steps - self.lr_warmup_steps, 1)
+        ratio = min(max((self.num_steps - self.lr_warmup_steps) / decay_span, 0.0), 1.0)
+        return self._decay(self.max_lr, self.min_lr, ratio, self.lr_decay_style)
+
+    def _get_wsd_lr(self) -> float:
+        decay_steps = self.wsd_decay_steps or 0
+        decay_start = max(self.lr_decay_steps - decay_steps, self.lr_warmup_steps)
+        if decay_steps <= 0 or self.num_steps <= decay_start:
+            return self.max_lr
+        ratio = min((self.num_steps - decay_start) / max(decay_steps, 1), 1.0)
+        return self._decay(self.max_lr, self.min_lr, ratio, self.lr_wsd_decay_style)
+
+    def _get_wd(self) -> float:
+        if self.wd_incr_style == "constant":
+            return self.end_wd
+        ratio = min(max(self.num_steps / self.wd_incr_steps, 0.0), 1.0)
+        return self._decay(self.start_wd, self.end_wd, ratio, self.wd_incr_style)
+
+    @staticmethod
+    def _decay(start: float, end: float, ratio: float, style: str) -> float:
+        if style == "linear":
+            return start + (end - start) * ratio
+        if style == "cosine":
+            coeff = 0.5 * (math.cos(math.pi * ratio) + 1.0)
+            return end + (start - end) * coeff
+        if style == "exponential":
+            if start == 0.0:
+                return 0.0
+            if end == 0.0:
+                return start * (1.0 - ratio)
+            return start * ((end / start) ** ratio)
+        if style == "constant":
+            return start
+        raise ValueError(f"Unsupported scheduler decay style: {style!r}")
+
+
 def _build_lr_scheduler(optimizer, opt: MegatronLiteOptimizerConfig):
     """Build a Megatron-style LR scheduler for Megatron Lite's optimizer."""
     total_steps = opt.total_training_steps
     if total_steps <= 0:
         return None
 
-    from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
-
     warmup_steps = opt.lr_warmup_steps if opt.lr_warmup_steps is not None else -1
     if warmup_steps <= 0 and opt.lr_warmup_steps_ratio > 0:
         warmup_steps = int(opt.lr_warmup_steps_ratio * total_steps)
@@ -66,7 +193,7 @@ def _build_lr_scheduler(optimizer, opt: MegatronLiteOptimizerConfig):
         if param_group.get("min_lr") is None:
             param_group["min_lr"] = min_lr
 
-    return OptimizerParamScheduler(
+    return _MegatronLiteLRScheduler(
         optimizer,
         init_lr=opt.lr_warmup_init,
         max_lr=opt.lr,
@@ -78,8 +205,6 @@ def _build_lr_scheduler(optimizer, opt: MegatronLiteOptimizerConfig):
         end_wd=opt.weight_decay,
         wd_incr_steps=total_steps,
         wd_incr_style=opt.weight_decay_incr_style,
-        use_checkpoint_opt_param_scheduler=opt.use_checkpoint_opt_param_scheduler,
-        override_opt_param_scheduler=not opt.use_checkpoint_opt_param_scheduler,
         wsd_decay_steps=opt.lr_wsd_decay_steps,
         lr_wsd_decay_style=opt.lr_wsd_decay_style,
     )
@@ -203,7 +328,7 @@ def forward_backward_batch(
         pad_mode = tu.get_non_tensor_data(
             data=data, key="pad_mode", default=DatasetPadMode.NO_PADDING
         )
-        if pad_mode != DatasetPadMode.NO_PADDING:
+        if not _is_no_padding_pad_mode(pad_mode):
             raise NotImplementedError(
                 "MegatronLiteEngine only supports pad_mode=no_padding for now."
             )
@@ -224,72 +349,15 @@ def forward_backward_batch(
             data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True
         )
 
-        if self._use_runtime_forward_backward():
-            return self._forward_backward_batch_with_runtime(
-                data=data,
-                micro_batches=micro_batches,
-                indices=indices,
-                loss_function=loss_function,
-                forward_only=forward_only,
-            )
-
-        outputs = []
-        num_micro_batches = len(micro_batches)
-        for micro_idx, micro_batch in enumerate(micro_batches):
-            tu.assign_non_tensor(micro_batch, micro_batch_idx=micro_idx)
-            micro_batch = micro_batch.to(get_device_id())
-            model_inputs = self._make_model_inputs(micro_batch)
-
-            pre_forward_hook = self.handle._extras.get("pre_forward_hook")
-            if pre_forward_hook is not None:
-                pre_forward_hook(torch.tensor(1.0 / num_micro_batches, device=get_device_id()))
-
-            with torch.no_grad() if forward_only else torch.enable_grad():
-                raw_output = self.module(
-                    input_ids=model_inputs["input_ids"],
-                    position_ids=model_inputs["position_ids"],
-                    packed_seq_params=model_inputs["packed_seq_params"],
-                    labels=model_inputs["labels"],
-                    loss_mask=model_inputs.get("loss_mask"),
-                    temperature=model_inputs["temperature"],
-                    use_fused_kernels=model_inputs["use_fused_kernels"],
-                    calculate_entropy=model_inputs["calculate_entropy"],
-                )
-
-                model_output = self._build_verl_model_output(
-                    raw_output=raw_output, micro_batch=micro_batch, inputs=model_inputs
-                )
-
-                if loss_function is not None:
-                    loss, metrics = loss_function(
-                        model_output=model_output,
-                        data=micro_batch,
-                        dp_group=self.get_data_parallel_group(),
-                    )
-                else:
-                    loss = torch.zeros((), device=get_device_id(), dtype=torch.float32)
-                    metrics = {}
-                if raw_output.get("mtp_loss") is not None:
-                    metrics = dict(metrics)
-                    mtp_loss = self._reduce_mtp_metric(raw_output["mtp_loss"])
-                    metrics["mtp_losses/mtp_1_loss"] = (
-                        float(mtp_loss.item()) if mtp_loss.numel() == 1 else mtp_loss.cpu().tolist()
-                    )
-
-            if not forward_only and loss_function is not None:
-                loss.backward()
-
-            outputs.append(
-                {"model_output": model_output, "loss": loss.detach().item(), "metrics": metrics}
-            )
-
-        if not forward_only:
-            finalize_grads = self.handle._extras.get("finalize_grads")
-            if finalize_grads is not None:
-                finalize_grads()
-
-        result = postprocess_batch_func(output_lst=outputs, indices=indices, data=data)
-        return result
+        # Megatron drives every forward through the runtime's forward_backward
+        # callback; the engine never calls the module directly.
+        return self._forward_backward_batch_with_runtime(
+            data=data,
+            micro_batches=micro_batches,
+            indices=indices,
+            loss_function=loss_function,
+            forward_only=forward_only,
+        )
 
     def get_per_tensor_param(self, **kwargs):
         self._require_initialized()
@@ -459,6 +527,7 @@ def _build_mlite_config(self) -> MegatronLiteConfig:
             optimizer=self._build_mlite_optimizer_config(),
             attention_backend_override=self.engine_config.attention_backend_override,
             router_aux_loss_coef=self.engine_config.router_aux_loss_coef,
+            load_hf_weights=self.engine_config.load_hf_weights,
             impl_cfg=self._build_impl_cfg(),
         )
 
@@ -474,6 +543,10 @@ def _build_impl_cfg(self) -> dict[str, Any]:
                 "MegatronLiteEngine supports only THD/no-padding SFT; set engine.impl_cfg.use_thd=True."
             )
         impl_cfg["use_thd"] = True
+        cross_entropy_fusion = getattr(self.engine_config, "cross_entropy_fusion", None)
+        if cross_entropy_fusion is None:
+            cross_entropy_fusion = getattr(self.engine_config, "use_fused_kernels", False)
+        impl_cfg.setdefault("cross_entropy_fusion", bool(cross_entropy_fusion))
         mtp_cfg = getattr(self.model_config, "mtp", None)
         if mtp_cfg is not None:
             mtp_enable = bool(getattr(mtp_cfg, "enable", False))
@@ -559,10 +632,6 @@ def _extract_primary_module(self):
             return model[0]
         return model
 
-    def _use_runtime_forward_backward(self) -> bool:
-        ps = self.handle._parallel_state
-        return ps.pp_size > 1
-
     def _forward_backward_batch_with_runtime(
         self,
         *,
@@ -585,21 +654,11 @@ def _forward_backward_batch_with_runtime(
         for micro_idx, micro_batch in enumerate(micro_batches):
             tu.assign_non_tensor(micro_batch, micro_batch_idx=micro_idx)
             micro_batch = micro_batch.to(get_device_id())
-            model_inputs = self._make_model_inputs(micro_batch)
             runtime_batches.append(
-                {
-                    "input_ids": model_inputs["input_ids"],
-                    "position_ids": model_inputs["position_ids"],
-                    "packed_seq_params": model_inputs["packed_seq_params"],
-                    "labels": model_inputs["labels"],
-                    "loss_mask": model_inputs.get("loss_mask"),
-                    "loss_scale": loss_scale,
-                    "temperature": model_inputs["temperature"],
-                    "use_fused_kernels": model_inputs["use_fused_kernels"],
-                    "calculate_entropy": model_inputs["calculate_entropy"],
-                    "_verl_micro_batch": micro_batch,
-                    "_verl_inputs": model_inputs,
-                }
+                (
+                    self._make_runtime_batch(micro_batch),
+                    self._make_runtime_loss_context(micro_batch, loss_scale=loss_scale),
+                )
             )
 
         runtime_loss_fn = None
@@ -624,43 +683,41 @@ def _forward_backward_batch_with_runtime(
             "metrics": {key: [value] for key, value in metrics.items()},
         }
 
-    def _make_model_inputs(self, micro_batch: TensorDict) -> dict[str, torch.Tensor]:
+    def _make_runtime_batch(self, micro_batch: TensorDict) -> PackedBatch:
+        """Flatten a jagged no-padding batch to a model-agnostic ``PackedBatch``.
+
+        No CP split, no padding, no ``PackedSeqParams`` here: each model's
+        protocol owns its pack/unpack pair (zigzag vs contiguous). ``labels`` are
+        the unrolled tokens; the protocol rolls them while packing.
+        """
         input_ids = micro_batch["input_ids"]
         if not getattr(input_ids, "is_nested", False):
             raise NotImplementedError(
                 "MegatronLiteEngine supports only nested no-padding THD batches."
             )
-
-        ps = self.handle._parallel_state
         loss_mask = self._loss_mask_for_packing(micro_batch, input_ids)
-        packed_batch = pack_nested_thd(
-            input_ids,
-            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,
-            labels=input_ids,
-            roll_labels=True,
-            loss_mask=loss_mask,
-            roll_loss_mask=True,
-        )
-        use_fused_kernels = tu.get_non_tensor_data(
-            data=micro_batch, key="use_fused_kernels", default=self.engine_config.use_fused_kernels
+        return PackedBatch(
+            input_ids=input_ids.values().contiguous(),
+            labels=input_ids.values().contiguous(),
+            loss_mask=None if loss_mask is None else loss_mask.values().contiguous().float(),
+            seq_lens=input_ids.offsets().diff().to(dtype=torch.int64),
         )
 
-        return {
-            "input_ids": packed_batch.input_ids,
-            "labels": packed_batch.labels,
-            "loss_mask": packed_batch.loss_mask,
-            "position_ids": packed_batch.position_ids,
-            "packed_seq_params": packed_batch.packed_seq_params,
-            "packed_batch": packed_batch,
-            "temperature": self._scalar_temperature(micro_batch),
-            "use_fused_kernels": use_fused_kernels,
-            "calculate_entropy": tu.get_non_tensor_data(
-                data=micro_batch, key="calculate_entropy", default=False
+    def _make_runtime_loss_context(
+        self,
+        micro_batch: TensorDict,
+        *,
+        loss_scale: float,
+    ) -> LossContext:
+        return LossContext(
+            temperature=float(self._scalar_temperature(micro_batch)),
+            calculate_entropy=bool(
+                tu.get_non_tensor_data(data=micro_batch, key="calculate_entropy", default=False)
             ),
-        }
+            return_log_probs=True,
+            loss_scale=loss_scale,
+            source_batch=micro_batch,
+        )
 
     @staticmethod
     def _loss_mask_for_packing(
@@ -691,26 +748,32 @@ def _build_verl_model_output(
         self,
         *,
         raw_output: dict[str, torch.Tensor],
-        micro_batch: TensorDict,
-        inputs: dict[str, torch.Tensor],
+        runtime_batch: PackedBatch,
     ) -> dict[str, torch.Tensor]:
-        del micro_batch
         log_probs = raw_output.get("log_probs")
         if log_probs is None:
             raise ValueError("Megatron Lite THD model output must contain token log_probs.")
-        nested_log_probs = unpack_packed_thd_to_nested(log_probs, inputs["packed_batch"])
-        output = {"log_probs": nested_log_probs}
+        proto = self.handle._extras.get("protocol")
+        unpack = getattr(proto, "unpack_forward_output", None)
+        if unpack is None:
+            raise ValueError(
+                "Model protocol must expose unpack_forward_output to reverse THD outputs."
+            )
+        output = {"log_probs": unpack(self.module, runtime_batch, log_probs)}
         entropy = raw_output.get("entropy")
         if entropy is not None:
-            output["entropy"] = unpack_packed_thd_to_nested(entropy, inputs["packed_batch"])
+            output["entropy"] = unpack(self.module, runtime_batch, entropy)
         return output
 
     def _make_runtime_loss_fn(self, loss_function, *, forward_only: bool):
-        def _loss_fn(raw_output: dict[str, torch.Tensor], runtime_batch: dict[str, Any]):
-            micro_batch = runtime_batch["_verl_micro_batch"]
-            inputs = runtime_batch["_verl_inputs"]
+        def _loss_fn(
+            raw_output: dict[str, torch.Tensor],
+            runtime_batch: PackedBatch,
+            loss_context: LossContext,
+        ):
+            micro_batch = loss_context.source_batch
             model_output = self._build_verl_model_output(
-                raw_output=raw_output, micro_batch=micro_batch, inputs=inputs
+                raw_output=raw_output, runtime_batch=runtime_batch
             )
             raw_output["_verl_model_output"] = model_output
             if loss_function is not None:
diff --git a/experimental/lite/examples/verl/verl_mlite/launch.py b/experimental/lite/examples/verl/verl_mlite/launch.py
new file mode 100644
index 00000000000..c862affb172
--- /dev/null
+++ b/experimental/lite/examples/verl/verl_mlite/launch.py
@@ -0,0 +1,26 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Launch a VERL module after applying MLite compatibility patches."""
+
+from __future__ import annotations
+
+import runpy
+import sys
+
+from verl_mlite.compat import apply_runtime_patches
+
+
+def main() -> None:
+    if len(sys.argv) < 2:
+        raise SystemExit("Usage: python -m verl_mlite.launch <module> [args...]")
+    module = sys.argv[1]
+    sys.argv = [module, *sys.argv[2:]]
+    apply_runtime_patches()
+    # Import the engine so its EngineRegistry.register decorator runs before the
+    # verl trainer resolves the "mlite" backend.
+    import verl_mlite.engine  # noqa: F401
+
+    runpy.run_module(module, run_name="__main__", alter_sys=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/experimental/lite/megatron/lite/model/qwen3_5/lite/protocol.py b/experimental/lite/megatron/lite/model/qwen3_5/lite/protocol.py
index a3cbd558bbb..bf8a8a2f135 100644
--- a/experimental/lite/megatron/lite/model/qwen3_5/lite/protocol.py
+++ b/experimental/lite/megatron/lite/model/qwen3_5/lite/protocol.py
@@ -10,15 +10,23 @@
 
 import torch
 import torch.nn as nn
-
+from megatron.lite.model.protocol_utils import (
+    add_cross_entropy_fusion,
+    add_loss_context_kwargs,
+    pack_thd_forward_kwargs,
+    set_cross_entropy_fusion,
+    unpack_thd_forward_output,
+)
 from megatron.lite.model.qwen3_5.config import Qwen35Config
 from megatron.lite.model.qwen3_5.lite.checkpoint import EXPERT_CLASSIFIER, PLACEMENT_FN
 from megatron.lite.model.qwen3_5.lite.checkpoint import export_hf_weights as _export_hf_weights_impl
 from megatron.lite.model.qwen3_5.lite.checkpoint import load_hf_weights as _load_hf_weights_impl
+from megatron.lite.model.qwen3_5.lite.checkpoint import save_hf_weights as _save_hf_weights_impl
 from megatron.lite.primitive.bundle import ModelBundle
 from megatron.lite.primitive.parallel import ParallelState, init_parallel
 from megatron.lite.primitive.recompute import apply_recompute, parse_recompute_spec
 from megatron.lite.runtime.contracts import OptimizerConfig, ParallelConfig
+from megatron.lite.runtime.contracts.data import PackedBatch
 
 __all__ = [
     "EXPERT_CLASSIFIER",
@@ -28,6 +36,7 @@
     "build_model_config",
     "export_hf_weights",
     "load_hf_weights",
+    "save_hf_weights",
     "vocab_size",
 ]
 
@@ -39,11 +48,12 @@ def is_expert_param(name: str) -> bool:
 @dataclass(frozen=True)
 class ImplConfig:
     parallel: ParallelConfig = field(default_factory=ParallelConfig)
-    optimizer: str | None = "mc_full"
+    optimizer: str | None = "dist_opt"
     recompute: list[str] = field(default_factory=list)
     offload: list[str] = field(default_factory=list)
     use_deepep: bool = False
     use_thd: bool = False
+    cross_entropy_fusion: bool = False
     hf_path: str = ""
     attention_backend_override: str | None = None
     router_aux_loss_coef: float | None = None
@@ -56,6 +66,7 @@ class ImplConfig:
     mtp_loss_scaling_factor: float = 0.1
     mtp_use_repeated_layer: bool | None = None
     mount_vision_model: bool = False
+    gdn_cp_mode: str = "fla_allgather"
 
 
 def _full_attn_module(layer, name: str):
@@ -86,20 +97,31 @@ def build_model_config(source: str | Path | dict, **overrides) -> Qwen35Config:
     return cfg
 
 
-def _forward_step(model: nn.Module, batch: dict) -> dict:
-    kwargs: dict[str, Any] = {"input_ids": batch["input_ids"], "labels": batch["labels"]}
-    if "position_ids" in batch:
-        kwargs["position_ids"] = batch["position_ids"]
-    if "packed_seq_params" in batch:
-        kwargs["packed_seq_params"] = batch["packed_seq_params"]
-    for key in ("loss_mask", "temperature", "use_fused_kernels", "calculate_entropy"):
-        if key in batch:
-            kwargs[key] = batch[key]
-    if kwargs["input_ids"].dim() == 1:
-        kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0)
+def _forward_step(model: nn.Module, batch: PackedBatch) -> dict:
+    kwargs = pack_thd_forward_kwargs(model, batch)
+    add_loss_context_kwargs(kwargs)
+    add_cross_entropy_fusion(kwargs, model)
+    return model(**kwargs)
+
+
+def _forward_step_bshd(model: nn.Module, batch: PackedBatch) -> dict:
+    """Dense [b=1, s] forward for a single packed sequence (no THD packing).
+
+    Used for deterministic parity comparison vs a dense Megatron-Core reference:
+    the THD GatedDeltaNet kernel is non-deterministic, whereas the dense path is
+    deterministic. CP=1 only (single unpadded sequence => dense == THD tokens).
+    """
+    input_ids = batch.input_ids.reshape(1, -1)
+    labels = batch.labels.reshape(1, -1) if batch.labels is not None else None
+    kwargs: dict[str, Any] = {"input_ids": input_ids, "labels": labels, "packed_seq_params": None}
+    add_cross_entropy_fusion(kwargs, model)
     return model(**kwargs)
 
 
+def unpack_forward_output(model: nn.Module, batch: PackedBatch, output) -> Any:
+    return unpack_thd_forward_output(model, batch, output)
+
+
 def _make_aux_loss_hook():
     from megatron.lite.primitive.modules.moe import MoEAuxLossAutoScaler
     from megatron.lite.primitive.modules.mtp import MTPLossAutoScaler
@@ -111,10 +133,14 @@ def hook(scale: torch.Tensor) -> None:
     return hook
 
 
-def _build_mc_optimizer(chunks, model_cfg: Qwen35Config, impl_cfg: ImplConfig, ps: ParallelState):
-    from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_training_optimizer
+def _build_dist_opt_optimizer(
+    chunks, model_cfg: Qwen35Config, impl_cfg: ImplConfig, ps: ParallelState
+):
+    from megatron.lite.primitive.optimizers.megatron_wrap import (
+        build_dist_opt_training_optimizer,
+    )
 
-    return build_mc_training_optimizer(
+    return build_dist_opt_training_optimizer(
         chunks,
         model_cfg=model_cfg,
         impl_cfg=impl_cfg,
@@ -173,6 +199,7 @@ def build_model(model_cfg: Qwen35Config, *, impl_cfg: ImplConfig) -> ModelBundle
         mtp_enable_train=mtp_enable_train,
         mtp_detach_encoder=impl_cfg.mtp_detach_encoder,
         mount_vision_model=impl_cfg.mount_vision_model,
+        gdn_cp_mode=impl_cfg.gdn_cp_mode,
     )
 
     if vpp is None:
@@ -184,6 +211,7 @@ def build_model(model_cfg: Qwen35Config, *, impl_cfg: ImplConfig) -> ModelBundle
             .cuda()
             for i in range(vpp)
         ]
+    set_cross_entropy_fusion(chunks, impl_cfg.cross_entropy_fusion)
 
     if recompute_spec:
         for chunk in chunks:
@@ -199,8 +227,8 @@ def build_model(model_cfg: Qwen35Config, *, impl_cfg: ImplConfig) -> ModelBundle
     finalize_grads = None
     post_model_load_hook = None
     optimizer_backend = "none"
-    if impl_cfg.optimizer in {"mc", "mc_full"}:
-        optimizer, finalize_grads = _build_mc_optimizer(chunks, model_cfg, impl_cfg, ps)
+    if impl_cfg.optimizer == "dist_opt":
+        optimizer, finalize_grads = _build_dist_opt_optimizer(chunks, model_cfg, impl_cfg, ps)
         from megatron.lite.primitive.ckpt import attach_model_sharded_state_dict
         from megatron.lite.runtime.megatron_utils import register_training_hooks
 
@@ -208,7 +236,7 @@ def build_model(model_cfg: Qwen35Config, *, impl_cfg: ImplConfig) -> ModelBundle
             chunks, ps, get_placements=PLACEMENT_FN, is_expert=is_expert_param
         )
         register_training_hooks(chunks, optimizer)
-        optimizer_backend = "distopt"
+        optimizer_backend = "dist_opt"
     elif impl_cfg.optimizer == "fsdp2":
         optimizer_backend = "fsdp2"
 
@@ -238,7 +266,7 @@ def _post_model_load_hook():
         parallel_state=ps,
         optimizer=optimizer,
         finalize_grads=finalize_grads,
-        forward_step=_forward_step,
+        forward_step=_forward_step if impl_cfg.use_thd else _forward_step_bshd,
         extras={
             "model_cfg": model_cfg,
             "optimizer_backend": optimizer_backend,
@@ -262,6 +290,12 @@ def export_hf_weights(
     yield from _export_hf_weights_impl(chunks, model_cfg, ps, **kwargs)
 
 
+def save_hf_weights(
+    chunks: list[nn.Module], path: str, model_cfg: Qwen35Config, ps: ParallelState
+) -> None:
+    _save_hf_weights_impl(chunks, path, model_cfg, ps)
+
+
 def vocab_size(model_cfg) -> int | None:
     cfg = getattr(model_cfg, "text_config", model_cfg)
     return getattr(cfg, "vocab_size", None)
diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py b/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py
index e8de45d63d6..9643d759495 100644
--- a/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py
+++ b/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py
@@ -25,7 +25,13 @@
 
 import torch
 import torch.nn as nn
-
+from megatron.lite.model.protocol_utils import (
+    add_cross_entropy_fusion,
+    add_loss_context_kwargs,
+    pack_thd_forward_kwargs,
+    set_cross_entropy_fusion,
+    unpack_thd_forward_output,
+)
 from megatron.lite.model.qwen3_moe.common import is_expert_param
 from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig
 from megatron.lite.model.qwen3_moe.lite.checkpoint import EXPERT_CLASSIFIER, PLACEMENT_FN
@@ -41,6 +47,7 @@
 from megatron.lite.primitive.parallel import ParallelState, init_parallel
 from megatron.lite.primitive.recompute import apply_recompute, parse_recompute_spec
 from megatron.lite.runtime.contracts import OptimizerConfig, ParallelConfig
+from megatron.lite.runtime.contracts.data import PackedBatch
 
 __all__ = [
     "EXPERT_CLASSIFIER",
@@ -63,11 +70,12 @@ class ImplConfig:
     """Lite impl knobs. Constructed by runtime from user config."""
 
     parallel: ParallelConfig = field(default_factory=ParallelConfig)
-    optimizer: str | None = "mc"  # None = no optimizer (inference)
+    optimizer: str | None = "dist_opt"  # None = no optimizer (inference)
     recompute: list[str] = field(default_factory=list)
     offload: list[str] = field(default_factory=list)
     use_deepep: bool = False
     use_thd: bool = False
+    cross_entropy_fusion: bool = False
     router_aux_loss_coef: float | None = None
     router_bias_rate: float = 0.0
     # User-level OptimizerConfig threaded through the runtime.
@@ -117,26 +125,17 @@ def build_model_config(source: str | Path | dict, **overrides) -> Qwen3MoEConfig
 # ---------------------------------------------------------------------------
 
 
-def _forward_step(model: nn.Module, batch: dict) -> dict:
-    kwargs = {"input_ids": batch["input_ids"], "labels": batch["labels"]}
-    if "packed_seq_params" in batch:
-        kwargs["packed_seq_params"] = batch["packed_seq_params"]
-    if "position_ids" in batch:
-        kwargs["position_ids"] = batch["position_ids"]
-    for key in (
-        "loss_mask",
-        "temperature",
-        "use_fused_kernels",
-        "calculate_entropy",
-        "return_log_probs",
-    ):
-        if key in batch:
-            kwargs[key] = batch[key]
-    if kwargs["input_ids"].dim() == 1:
-        kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0)
+def _forward_step(model: nn.Module, batch: PackedBatch) -> dict:
+    kwargs = pack_thd_forward_kwargs(model, batch)
+    add_loss_context_kwargs(kwargs, include_return_log_probs=True)
+    add_cross_entropy_fusion(kwargs, model)
     return model(**kwargs)
 
 
+def unpack_forward_output(model: nn.Module, batch: PackedBatch, output) -> Any:
+    return unpack_thd_forward_output(model, batch, output)
+
+
 def build_model(model_cfg: Qwen3MoEConfig, *, impl_cfg: ImplConfig) -> ModelBundle:
     """Build lite Qwen3MoE: model, parallel state, optimizer — everything.
 
@@ -193,6 +192,8 @@ def build_model(model_cfg: Qwen3MoEConfig, *, impl_cfg: ImplConfig) -> ModelBund
                 .cuda()
             )
 
+    set_cross_entropy_fusion(chunks, impl_cfg.cross_entropy_fusion)
+
     # ── recompute ──
     if recompute_spec:
         for chunk in chunks:
@@ -217,10 +218,12 @@ def build_model(model_cfg: Qwen3MoEConfig, *, impl_cfg: ImplConfig) -> ModelBund
     optimizer = None
     finalize_grads = None
     post_model_load_hook = None
-    if impl_cfg.optimizer == "mc":
-        from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_training_optimizer
+    if impl_cfg.optimizer == "dist_opt":
+        from megatron.lite.primitive.optimizers.megatron_wrap import (
+            build_dist_opt_training_optimizer,
+        )
 
-        optimizer, finalize_grads = build_mc_training_optimizer(
+        optimizer, finalize_grads = build_dist_opt_training_optimizer(
             chunks,
             model_cfg=model_cfg,
             impl_cfg=impl_cfg,
@@ -234,7 +237,7 @@ def build_model(model_cfg: Qwen3MoEConfig, *, impl_cfg: ImplConfig) -> ModelBund
         attach_model_sharded_state_dict(
             chunks, ps, get_placements=PLACEMENT_FN, is_expert=is_expert_param
         )
-        optimizer_backend = "distopt"
+        optimizer_backend = "dist_opt"
     elif impl_cfg.optimizer == "fsdp2":
         optimizer_backend = "fsdp2"
 
diff --git a/experimental/lite/megatron/lite/primitive/optimizers/megatron_wrap.py b/experimental/lite/megatron/lite/primitive/optimizers/megatron_wrap.py
index ec182609fe1..8545a07d22d 100644
--- a/experimental/lite/megatron/lite/primitive/optimizers/megatron_wrap.py
+++ b/experimental/lite/megatron/lite/primitive/optimizers/megatron_wrap.py
@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from types import SimpleNamespace
 from typing import Any
 
@@ -13,23 +13,22 @@
 from megatron.lite.primitive.protocols import ExpertClassifierFn, default_expert_classifier
 
 
-def validate_mc_config(engine_cfg) -> None:
+def validate_dist_opt_config(engine_cfg) -> None:
     """Validate dist_opt constraints owned by this optimizer primitive."""
     p = engine_cfg.parallel
     if p.vpp > 1 and p.pp == 1:
         raise ValueError("dist_opt requires pp>1 when vpp>1.")
 
 
-# Legacy alias — kept for compat shim path
-validate_mc_session = validate_mc_config
+validate_dist_opt_session = validate_dist_opt_config
 
 
 def _effective_etp(parallel) -> int:
     return int(parallel.etp if parallel.etp is not None else 1)
 
 
-def _ensure_mc_mpu_parallel_state(engine_cfg) -> None:
-    """Initialize Megatron-Core mpu globals when MC fallback groups are used."""
+def _ensure_dist_opt_mpu_parallel_state(engine_cfg) -> None:
+    """Initialize Megatron-Core mpu globals when dist_opt fallback groups are used."""
 
     from megatron.core import parallel_state as mpu  # pyright: ignore[reportMissingImports]
 
@@ -61,8 +60,10 @@ def _ensure_mc_mpu_parallel_state(engine_cfg) -> None:
     )
 
 
-def build_mc_optimizer_config(opt, *, override_optimizer_config: dict[str, Any] | None = None):
-    """Build MC OptimizerConfig from user's OptimizerConfig (duck-typed).
+def build_dist_opt_optimizer_config(
+    opt, *, override_optimizer_config: dict[str, Any] | None = None
+):
+    """Build Megatron-Core OptimizerConfig from user's OptimizerConfig (duck-typed).
 
     Single source of truth for Megatron Lite's Megatron-Core optimizer stack.
 
@@ -70,7 +71,7 @@ def build_mc_optimizer_config(opt, *, override_optimizer_config: dict[str, Any]
     or a `SimpleNamespace` with the same field names (legacy lite path).
     """
     from megatron.core.optimizer.optimizer_config import (
-        OptimizerConfig as MCOptimizerConfig,  # pyright: ignore[reportMissingImports]
+        OptimizerConfig as CoreOptimizerConfig,  # pyright: ignore[reportMissingImports]
     )
 
     offload = getattr(opt, "offload_fraction", None) or 0.0
@@ -100,10 +101,10 @@ def build_mc_optimizer_config(opt, *, override_optimizer_config: dict[str, Any]
         args["decoupled_weight_decay"] = opt.decoupled_weight_decay
     if override_optimizer_config:
         args.update(override_optimizer_config)
-    return MCOptimizerConfig(**args)
+    return CoreOptimizerConfig(**args)
 
 
-def build_mc_stack(
+def build_dist_opt_stack(
     model_chunks: list[nn.Module],
     *,
     model_cfg,
@@ -113,11 +114,11 @@ def build_mc_stack(
     proto=None,
     skip_ddp_wrap: bool = False,
 ):
-    """Wrap ML model chunks with MC DDP and build the matching MC optimizer.
+    """Wrap ML model chunks with Megatron-Core DDP and build the matching dist_opt optimizer.
 
     Args:
         skip_ddp_wrap: when True, ``model_chunks`` are assumed to already be
-            MC ``DistributedDataParallel``-wrapped; we skip our own wrapping
+            Megatron-Core ``DistributedDataParallel``-wrapped; we skip our own wrapping
             and feed them directly to the optimizer. The bucket layout
             influences optimizer master-grad sharding, so callers that prewrap
             chunks own the DDP config compatibility.
@@ -127,13 +128,13 @@ def build_mc_stack(
     from megatron.core.optimizer import get_megatron_optimizer
     from megatron.core.transformer.enums import ModelType
 
-    validate_mc_config(engine_cfg)
+    validate_dist_opt_config(engine_cfg)
 
     p = engine_cfg.parallel
     opt = engine_cfg.optimizer
 
-    mc_transformer_cfg = _build_transformer_config(model_cfg, engine_cfg)
-    mc_transformer_cfg.finalize_model_grads_func = finalize_model_grads
+    dist_opt_transformer_cfg = _build_transformer_config(model_cfg, engine_cfg)
+    dist_opt_transformer_cfg.finalize_model_grads_func = finalize_model_grads
     if is_expert is not None:
         is_expert_param = is_expert
     elif proto is not None and hasattr(proto, "EXPERT_CLASSIFIER"):
@@ -142,12 +143,12 @@ def build_mc_stack(
         is_expert_param = default_expert_classifier
     use_mpu_groups = bool(getattr(engine_cfg, "deterministic", False))
     if use_mpu_groups:
-        _ensure_mc_mpu_parallel_state(engine_cfg)
+        _ensure_dist_opt_mpu_parallel_state(engine_cfg)
     pg_collection = None if use_mpu_groups else _build_pg_collection(ps, engine_cfg)
 
     if skip_ddp_wrap:
         # Caller already wrapped and marked every param. Our helper setting
-        # `param.allreduce` on dense params could clash with MC code paths that
+        # `param.allreduce` on dense params could clash with Megatron-Core code paths that
         # distinguish `hasattr(param,'allreduce')` from `getattr(..., True)`.
         wrapped_chunks = list(model_chunks)
     else:
@@ -157,13 +158,13 @@ def build_mc_stack(
         wrapped_chunks = []
         for chunk_idx, chunk in enumerate(model_chunks):
             chunk.model_type = ModelType.encoder_or_decoder
-            _mark_mc_parallel_attrs(chunk, is_expert_param, tp_size=p.tp)
+            _mark_dist_opt_parallel_attrs(chunk, is_expert_param, tp_size=p.tp)
             ddp_kwargs = {}
             if pg_collection is not None:
                 ddp_kwargs["pg_collection"] = pg_collection
             wrapped_chunks.append(
                 DistributedDataParallel(
-                    mc_transformer_cfg,
+                    dist_opt_transformer_cfg,
                     ddp_config,
                     chunk,
                     disable_bucketing=(chunk_idx > 0),
@@ -173,14 +174,14 @@ def build_mc_stack(
 
     # Single-source-of-truth OptimizerConfig construction for native lite
     # model protocols.
-    opt_config = build_mc_optimizer_config(opt)
+    opt_config = build_dist_opt_optimizer_config(opt)
 
-    # This branch falls back to MC mpu globals for the optimizer's process
+    # This branch falls back to Megatron-Core mpu globals for the optimizer's process
     # groups. Long term, this primitive should always pass its own
     # `pg_collection`.
     if skip_ddp_wrap or use_mpu_groups:
         optimizer = get_megatron_optimizer(config=opt_config, model_chunks=wrapped_chunks)
-        optimizer._mc_pg_collection = None  # pyright: ignore[reportAttributeAccessIssue]
+        optimizer._dist_opt_pg_collection = None  # pyright: ignore[reportAttributeAccessIssue]
     else:
         optimizer = get_megatron_optimizer(
             config=opt_config,
@@ -188,11 +189,13 @@ def build_mc_stack(
             use_gloo_process_groups=False,
             pg_collection=pg_collection,
         )
-        optimizer._mc_pg_collection = pg_collection  # pyright: ignore[reportAttributeAccessIssue]
+        optimizer._dist_opt_pg_collection = (
+            pg_collection  # pyright: ignore[reportAttributeAccessIssue]
+        )
     return wrapped_chunks, optimizer
 
 
-def build_mc_training_optimizer(
+def build_dist_opt_training_optimizer(
     model_chunks: list[nn.Module],
     *,
     model_cfg,
@@ -203,7 +206,7 @@ def build_mc_training_optimizer(
     skip_ddp_wrap: bool = False,
     deterministic: bool | None = None,
 ):
-    """Build the MC DDP+optimizer stack from a Megatron Lite model ImplConfig."""
+    """Build the dist_opt DDP+optimizer stack from a Megatron Lite model ImplConfig."""
 
     opt = impl_cfg.optimizer_config
     if opt is None:
@@ -228,7 +231,7 @@ def build_mc_training_optimizer(
         optimizer=opt,
         deterministic=bool(deterministic),
     )
-    model_chunks[:], optimizer = build_mc_stack(
+    model_chunks[:], optimizer = build_dist_opt_stack(
         model_chunks,
         model_cfg=model_cfg,
         engine_cfg=engine_cfg,
@@ -238,16 +241,16 @@ def build_mc_training_optimizer(
     )
 
     def finalize_grads() -> None:
-        finalize_mc_grads(model_chunks, optimizer)
+        finalize_dist_opt_grads(model_chunks, optimizer)
 
     return optimizer, finalize_grads
 
 
-def finalize_mc_grads(model_chunks: list[nn.Module], optimizer) -> None:
-    """Run MC gradient finalization to match the optimizer's expected contract."""
+def finalize_dist_opt_grads(model_chunks: list[nn.Module], optimizer) -> None:
+    """Run Megatron-Core gradient finalization to match the optimizer's expected contract."""
     from megatron.core.distributed.finalize_model_grads import finalize_model_grads
 
-    finalize_model_grads(model_chunks, pg_collection=optimizer._mc_pg_collection)
+    finalize_model_grads(model_chunks, pg_collection=optimizer._dist_opt_pg_collection)
 
 
 def _build_transformer_config(model_cfg, engine_cfg):
@@ -272,31 +275,31 @@ def _build_transformer_config(model_cfg, engine_cfg):
     )
     if hasattr(model_cfg, "add_bias_linear"):
         kwargs["add_bias_linear"] = bool(model_cfg.add_bias_linear)
-    elif kwargs["num_moe_experts"] is not None and kwargs["expert_tensor_parallel_size"] > 1:
+    elif kwargs["num_moe_experts"] is not None:
         kwargs["add_bias_linear"] = False
     if p.pp > 1:
         kwargs["pipeline_dtype"] = torch.bfloat16
     return TransformerConfig(**kwargs)
 
 
-def _mark_mc_parallel_attrs(
+def _mark_dist_opt_parallel_attrs(
     model: nn.Module, is_expert_param: ExpertClassifierFn, *, tp_size: int
 ) -> None:
-    """Mark per-param MC metadata (allreduce / tensor_model_parallel / sequence_parallel).
+    """Mark per-param optimizer metadata (allreduce / tensor_model_parallel / sequence_parallel).
 
-    IMPORTANT: respect attrs that are already set. Prewrapped MC models may
+    IMPORTANT: respect attrs that are already set. Prewrapped Megatron-Core models may
     mark these correctly per-param (e.g. `moe.router.weight` is 2D but
     TP-replicated, and must NOT have `tensor_model_parallel=True`). Blind
-    override would cause MC grad-norm to over-count replicated params.
+    override would cause dist_opt grad-norm to over-count replicated params.
     """
     sp_param_ids = {id(param) for param in getattr(model, "sp_params", [])}
     for name, param in model.named_parameters():
-        # MC uses `allreduce=False` to route expert params into expert-DP buffers.
+        # Megatron-Core uses `allreduce=False` to route expert params into expert-DP buffers.
         if not hasattr(param, "allreduce"):
             param.allreduce = not is_expert_param(name)
         if tp_size > 1 and id(param) not in sp_param_ids and param.ndim > 1:
             # vision params are replicated across TP (AVG all-reduce, not TP-split).
-            # tensor_model_parallel=True would cause MC to wrong-account their grad-norm.
+            # tensor_model_parallel=True would cause dist_opt to wrong-account their grad-norm.
             if getattr(param, "average_gradients_across_tp_domain", False):
                 continue
             # Skip params already marked sequence_parallel=True: they are TP-replicated
@@ -304,7 +307,7 @@ def _mark_mc_parallel_attrs(
             # Stacking tensor_model_parallel=True on top would cause double all-reduce.
             if getattr(param, "sequence_parallel", False):
                 continue
-            # MC excludes TP replicas from grad-norm accounting via this metadata.
+            # Distopt excludes TP replicas from grad-norm accounting via this metadata.
             if not hasattr(param, "tensor_model_parallel"):
                 param.tensor_model_parallel = True
 
@@ -339,7 +342,7 @@ def _expert_rank(etp_i: int, ep_i: int, edp_i: int, pp_i: int) -> int:
             singleton_group = group
     if singleton_group is None:
         raise RuntimeError(
-            "Failed to construct singleton process group for optional MC reductions."
+            "Failed to construct singleton process group for optional dist_opt reductions."
         )
 
     if engine_cfg.parallel.pp == 1:
@@ -371,9 +374,9 @@ def _expert_rank(etp_i: int, ep_i: int, edp_i: int, pp_i: int) -> int:
                 tp_ep_pp_group = group
 
         if mp_group is None or tp_ep_pp_group is None:
-            raise RuntimeError("Failed to construct mc pipeline-aware process groups.")
+            raise RuntimeError("Failed to construct dist_opt pipeline-aware process groups.")
 
-    return ProcessGroupCollection(
+    pg_kwargs = dict(
         tp=ps.tp_group,
         cp=ps.cp_group,
         pp=ps.pp_group,
@@ -385,15 +388,19 @@ def _expert_rank(etp_i: int, ep_i: int, edp_i: int, pp_i: int) -> int:
         expt_tp=ps.etp_group,
         tp_ep=ps.tp_ep_group,
         tp_ep_pp=tp_ep_pp_group,
-        # For MC distributed optimizer, grad stats are reduced over the full optimizer instance.
+        # For dist_opt, grad stats are reduced over the full optimizer instance.
         # With a single dist-opt instance in this benchmark proof, that is the global world group.
         intra_dist_opt=dist.group.WORLD,
-        # ML models do not expose MC's embedding/position-embedding sharing surface.
-        # Use singleton groups so MC's optional embedding reductions become no-ops
+        # ML models do not expose Megatron-Core's embedding/position-embedding sharing surface.
+        # Use singleton groups so optional embedding reductions become no-ops
         # without falling back to the global MCore embedding group.
         embd=singleton_group,
         pos_embd=singleton_group,
     )
+    supported_fields = {field.name for field in fields(ProcessGroupCollection)}
+    return ProcessGroupCollection(
+        **{key: value for key, value in pg_kwargs.items() if key in supported_fields}
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -402,9 +409,9 @@ def _expert_rank(etp_i: int, ep_i: int, edp_i: int, pp_i: int) -> int:
 
 
 @dataclass(frozen=True, slots=True)
-class MCBackend:
-    name: str = "mc"
-    runtime_backend: str = "mc"
+class DistOptBackend:
+    name: str = "dist_opt"
+    runtime_backend: str = "dist_opt"
 
     def zero_grad(self, optimizer: Any) -> None:
         optimizer.zero_grad()
@@ -431,14 +438,15 @@ def finalize_grads(self, finalize_fn, model_chunks: list[Any], optimizer: Any) -
         finalize_fn(model_chunks, optimizer)
 
 
-BACKEND = MCBackend()
+BACKEND = DistOptBackend()
 
 __all__ = [
     "BACKEND",
-    "MCBackend",
-    "build_mc_stack",
-    "build_mc_training_optimizer",
-    "finalize_mc_grads",
-    "validate_mc_config",
-    "validate_mc_session",
+    "DistOptBackend",
+    "build_dist_opt_optimizer_config",
+    "build_dist_opt_stack",
+    "build_dist_opt_training_optimizer",
+    "finalize_dist_opt_grads",
+    "validate_dist_opt_config",
+    "validate_dist_opt_session",
 ]
diff --git a/experimental/lite/megatron/lite/runtime/backends/bridge/config.py b/experimental/lite/megatron/lite/runtime/backends/bridge/config.py
index e1df53439fa..cdb20c3d684 100644
--- a/experimental/lite/megatron/lite/runtime/backends/bridge/config.py
+++ b/experimental/lite/megatron/lite/runtime/backends/bridge/config.py
@@ -27,6 +27,11 @@ class BridgeConfig:
     load_hf_weights: bool = True
     build_optimizer: bool = True
 
+    # When False, the bridge feeds a dense [b=1, s] forward (no THD packing).
+    # Used for deterministic layout-matched parity vs models whose Megatron-Core
+    # kernel is dense-only (e.g. GatedDeltaNet). Default True keeps THD packing.
+    use_thd: bool = True
+
     override_ddp_config: dict[str, Any] = field(default_factory=dict)
     override_transformer_config: dict[str, Any] = field(default_factory=dict)
     override_optimizer_config: dict[str, Any] = field(default_factory=dict)
diff --git a/experimental/lite/megatron/lite/runtime/backends/bridge/runtime.py b/experimental/lite/megatron/lite/runtime/backends/bridge/runtime.py
index 7f4b7ce886e..2790fdf19ce 100644
--- a/experimental/lite/megatron/lite/runtime/backends/bridge/runtime.py
+++ b/experimental/lite/megatron/lite/runtime/backends/bridge/runtime.py
@@ -11,11 +11,10 @@
 
 import torch
 import torch.distributed as dist
-
-from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_optimizer_config
+from megatron.lite.primitive.optimizers.megatron_wrap import build_dist_opt_optimizer_config
 from megatron.lite.runtime.backends import Runtime as RuntimeBase
 from megatron.lite.runtime.backends.bridge.config import BridgeConfig
-from megatron.lite.runtime.contracts.data import Batch, ForwardResult, ModelOutputs
+from megatron.lite.runtime.contracts.data import Batch, ForwardResult, ModelOutputs, PackedBatch
 from megatron.lite.runtime.contracts.handle import ModelHandle
 from megatron.lite.runtime.megatron_utils import (
     build_sharded_state_dict,
@@ -92,7 +91,12 @@ def _bridge_hf_config(bridge):
 
 def _lower_provider_value(key: str, value: Any) -> Any:
     if key == "attention_backend" and isinstance(value, str):
-        from megatron.core.transformer.enums import AttnBackend
+        try:
+            from megatron.core.transformer.enums import AttnBackend
+        except ModuleNotFoundError as exc:
+            if exc.name != "megatron.core":
+                raise
+            return value
 
         return AttnBackend[value]
     return value
@@ -414,7 +418,7 @@ def _build_optimizer(model_list: list, cfg: BridgeConfig):
     from megatron.core.optimizer import get_megatron_optimizer
 
     return get_megatron_optimizer(
-        config=build_mc_optimizer_config(
+        config=build_dist_opt_optimizer_config(
             cfg.optimizer, override_optimizer_config=cfg.override_optimizer_config
         ),
         model_chunks=model_list,
@@ -491,6 +495,109 @@ def _as_data_iter(data: Any):
     return iter([data])
 
 
+# MLITE_LAYERING_ALLOW_BRIDGE_FORWARD_METADATA_BEGIN
+def _nested_from_packed(tensor: torch.Tensor | None, seq_lens: torch.Tensor):
+    """Split a 1-D packed (true, unpadded) tensor into a jagged nested tensor."""
+    if tensor is None:
+        return None
+    flat = tensor.reshape(-1)
+    pieces = []
+    offset = 0
+    for length_t in seq_lens.tolist():
+        length = int(length_t)
+        pieces.append(flat.narrow(0, offset, length))
+        offset += length
+    if offset != flat.numel():
+        raise ValueError(
+            f"PackedBatch sizes sum to {offset}, tensor has {flat.numel()} tokens."
+        )
+    return torch.nested.as_nested_tensor(pieces, layout=torch.jagged)
+
+
+def _bridge_forward_kwargs_from_packed_batch(
+    batch: PackedBatch,
+    *,
+    tp_size: int = 1,
+    cp_size: int = 1,
+    cp_rank: int = 0,
+    cp_group: Any = None,
+) -> dict[str, Any]:
+    """Render transient Megatron-Core THD kwargs for BridgeRuntime.forward_step only.
+
+    Reuses the same canonical packing the native lite protocols use
+    (:func:`pack_nested_thd` + :func:`prepare_packed_thd_for_context_parallel`):
+    each sequence is padded to the TP/CP zigzag alignment, ``labels`` are rolled
+    one position left, and for ``cp_size > 1`` the token rows / labels / loss
+    mask / position ids are zigzag-split to this CP rank while ``packed_seq_params``
+    keeps full ``cu_seqlens`` plus CP metadata. Identical layout on both backends
+    keeps the mlite-vs-bridge comparison fair and CP-correct.
+    """
+    from megatron.lite.primitive.parallel.thd import (
+        pack_nested_thd,
+        prepare_packed_thd_for_context_parallel,
+    )
+
+    seq_lens = batch.seq_lens
+    has_loss_mask = batch.loss_mask is not None
+    packed = pack_nested_thd(
+        _nested_from_packed(batch.input_ids, seq_lens),
+        tp_size=tp_size,
+        cp_size=cp_size,
+        cp_rank=cp_rank,
+        cp_group=cp_group if 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) if has_loss_mask else None,
+        roll_loss_mask=has_loss_mask,
+    )
+    sample: dict[str, Any] = {
+        "input_ids": packed.input_ids,
+        "labels": packed.labels,
+        "position_ids": packed.position_ids,
+        "packed_seq_params": packed.packed_seq_params,
+    }
+    if packed.loss_mask is not None:
+        sample["loss_mask"] = packed.loss_mask
+
+    if cp_size > 1:
+        tensor_keys = ("input_ids", "labels", "loss_mask", "position_ids")
+        psp, tensors = prepare_packed_thd_for_context_parallel(
+            sample["packed_seq_params"],
+            tuple(sample.get(key) for key in tensor_keys),
+            cp_size=cp_size,
+            cp_rank=cp_rank,
+            cp_group=cp_group,
+        )
+        sample["packed_seq_params"] = psp
+        for key, tensor in zip(tensor_keys, tensors, strict=True):
+            if tensor is not None or key in sample:
+                sample[key] = tensor
+    return sample
+# MLITE_LAYERING_ALLOW_BRIDGE_FORWARD_METADATA_END
+
+
+def _bridge_forward_kwargs_bshd(batch: PackedBatch) -> dict[str, Any]:
+    """Render dense [b=1, s] forward kwargs for a single packed sequence.
+
+    Used when ``use_thd`` is False so the Megatron-Core reference runs its dense
+    kernel (e.g. GatedDeltaNet, whose THD path is unsupported / non-deterministic
+    in the pinned core), giving a deterministic, layout-matched parity vs the
+    native lite dense forward on identical tokens. Carries no THD packing
+    metadata (no ``packed_seq_params`` / ``position_ids``) — single unpadded
+    sequence, so an all-ones attention mask reproduces the full causal sequence.
+    CP=1 only.
+    """
+    total = int(batch.seq_lens.sum().item())
+    return {
+        "input_ids": batch.input_ids.reshape(1, total).contiguous(),
+        "labels": (
+            batch.labels.reshape(1, total).contiguous() if batch.labels is not None else None
+        ),
+        "attention_mask": torch.ones((1, total), dtype=torch.long, device=batch.input_ids.device),
+    }
+
+
 class BridgeRuntime(RuntimeBase):
     """Megatron-Bridge training backend using Megatron-Core optimizer state."""
 
@@ -574,7 +681,7 @@ def build_model(
                 "mpu": mpu,
                 "model_cfg": _bridge_hf_config(bridge),
                 "protocol": _resolve_benchmark_protocol(rt_cfg, bridge),
-                "optimizer_backend": "distopt" if optimizer is not None else "none",
+                "optimizer_backend": "dist_opt" if optimizer is not None else "none",
                 "world_size": dist.get_world_size(),
             },
         )
@@ -597,19 +704,43 @@ def forward_backward(
         model_list = handle._extras["model_list"]
         data_iter = _as_data_iter(data)
         last_loss: list[float | None] = [None]
+        last_output: list[torch.Tensor | None] = [None]
+
+        # CP geometry for the transient PackedBatch -> THD kwargs rendering below.
+        cp_size = mpu.get_context_parallel_world_size()
+        cp_kwargs = {
+            "tp_size": mpu.get_tensor_model_parallel_world_size(),
+            "cp_size": cp_size,
+            "cp_rank": mpu.get_context_parallel_rank(),
+            "cp_group": mpu.get_context_parallel_group() if cp_size > 1 else None,
+        }
 
         def _fwd_step(data_iterator, model):
+            # MLITE_LAYERING_ALLOW_BRIDGE_FORWARD_METADATA_BEGIN
             sample = next(data_iterator)
-            if isinstance(sample, Batch):
+            owns_transient_metadata = False
+            if isinstance(sample, PackedBatch):
+                if self._cfg.use_thd:
+                    sample = _bridge_forward_kwargs_from_packed_batch(sample, **cp_kwargs)
+                    owns_transient_metadata = True
+                else:
+                    sample = _bridge_forward_kwargs_bshd(sample)
+            elif isinstance(sample, Batch):
                 sample = {
                     "input_ids": sample["input_ids"],
                     "labels": sample["labels"],
-                    "position_ids": getattr(sample, "position_ids", None),
                 }
             if not isinstance(sample, dict):
                 raise TypeError(
                     f"BridgeRuntime expected dict or Batch data, got {type(sample).__name__}."
                 )
+            if not owns_transient_metadata:
+                leaked = {"packed_seq_params", "position_ids"}.intersection(sample)
+                if leaked:
+                    raise ValueError(
+                        "BridgeRuntime data must not carry model-internal keys "
+                        f"{sorted(leaked)}; pass a raw PackedBatch instead."
+                    )
 
             output_tensor = model(
                 input_ids=sample["input_ids"],
@@ -620,16 +751,23 @@ def _fwd_step(data_iterator, model):
             )
             if isinstance(output_tensor, tuple):
                 output_tensor = output_tensor[0]
+            last_output[0] = output_tensor
+            loss_sample = {
+                key: value
+                for key, value in sample.items()
+                if key not in {"packed_seq_params", "position_ids"}
+            }
+            # MLITE_LAYERING_ALLOW_BRIDGE_FORWARD_METADATA_END
 
-            def _mc_loss_fn(output_tensor, non_loss_data=False):
+            def _bridge_loss_fn(output_tensor, non_loss_data=False):
                 if loss_fn is not None:
-                    loss, _metrics = loss_fn({"output_tensor": output_tensor}, sample)
+                    loss, _metrics = loss_fn({"output_tensor": output_tensor}, loss_sample)
                 else:
                     loss = output_tensor.mean()
                 last_loss[0] = float(loss.detach().item())
                 return loss, {}
 
-            return output_tensor, _mc_loss_fn
+            return output_tensor, _bridge_loss_fn
 
         vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()
         if vpp_size is not None and vpp_size > 1:
@@ -661,7 +799,7 @@ def _mc_loss_fn(output_tensor, non_loss_data=False):
 
         result_loss = torch.tensor(loss_val or 0.0)
         return ForwardResult(
-            model_output=ModelOutputs(loss=result_loss),
+            model_output=ModelOutputs(loss=result_loss, vocab_parallel_logits=last_output[0]),
             metrics={"loss": loss_val if loss_val is not None else 0.0},
         )
 
diff --git a/experimental/lite/megatron/lite/runtime/backends/mbridge/runtime.py b/experimental/lite/megatron/lite/runtime/backends/mbridge/runtime.py
index faed2738dfa..d2424a69cb4 100644
--- a/experimental/lite/megatron/lite/runtime/backends/mbridge/runtime.py
+++ b/experimental/lite/megatron/lite/runtime/backends/mbridge/runtime.py
@@ -152,7 +152,7 @@ def build_model(
                 "mpu": mpu,
                 "model_cfg": bridge.hf_config,
                 "protocol": _resolve_mbridge_benchmark_protocol(rt_cfg, bridge),
-                "optimizer_backend": "distopt" if optimizer is not None else "none",
+                "optimizer_backend": "dist_opt" if optimizer is not None else "none",
                 "world_size": dist.get_world_size(),
             },
         )
diff --git a/experimental/lite/megatron/lite/runtime/backends/mlite/runtime.py b/experimental/lite/megatron/lite/runtime/backends/mlite/runtime.py
index c197447b70d..fc194f9d4be 100644
--- a/experimental/lite/megatron/lite/runtime/backends/mlite/runtime.py
+++ b/experimental/lite/megatron/lite/runtime/backends/mlite/runtime.py
@@ -12,17 +12,20 @@
 
 import torch
 import torch.distributed as dist
-
 from megatron.lite.runtime.backends import Runtime as RuntimeBase
 from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig
-from megatron.lite.runtime.contracts.data import ForwardResult, ModelOutputs
+from megatron.lite.runtime.contracts.data import ForwardResult, ModelOutputs, PackedBatch
 from megatron.lite.runtime.contracts.handle import ModelHandle
+from megatron.lite.runtime.contracts.loss import split_loss_context
 
 
 def _build_impl_cfg(proto, rt_cfg: MegatronLiteConfig):
     """Construct typed impl config, backfilling hf_path + optimizer_config."""
-    impl_cfg_kwargs = {**rt_cfg.impl_cfg, "parallel": rt_cfg.parallel}
     init_fields = {f.name for f in dc_fields(proto.ImplConfig) if f.init}
+    # Only forward impl_cfg keys this model's ImplConfig declares: the connector
+    # may pass knobs (e.g. cross_entropy_fusion) that some models don't model.
+    impl_cfg_kwargs = {key: value for key, value in rt_cfg.impl_cfg.items() if key in init_fields}
+    impl_cfg_kwargs["parallel"] = rt_cfg.parallel
     if (
         "attention_backend_override" in init_fields
         and impl_cfg_kwargs.get("attention_backend_override") is None
@@ -64,13 +67,13 @@ def _apply_attention_backend_env(backend: str | None, *, tag: str) -> None:
     os.environ["NVTE_UNFUSED_ATTN"] = unfused
 
 
-def _infer_pipeline_tensor_shape(batch: Any, model_cfg: Any, ps) -> tuple[int, int, int]:
+def _infer_pipeline_tensor_shape(batch: PackedBatch, model_cfg: Any, ps) -> tuple[int, int, int]:
     if model_cfg is None or not hasattr(model_cfg, "hidden_size"):
         raise ValueError("Megatron Lite pipeline runtime requires model_cfg.hidden_size.")
-    if not isinstance(batch, dict) or "input_ids" not in batch:
-        raise TypeError("Megatron Lite pipeline runtime requires dict batches with input_ids.")
+    if not isinstance(batch, PackedBatch):
+        raise TypeError("Megatron Lite pipeline runtime requires PackedBatch inputs.")
 
-    input_ids = batch["input_ids"]
+    input_ids = batch.input_ids
     if input_ids.dim() == 1:
         batch_size = 1
         local_seq_len = int(input_ids.size(0))
@@ -84,16 +87,33 @@ def _infer_pipeline_tensor_shape(batch: Any, model_cfg: Any, ps) -> tuple[int, i
         raise ValueError("Pipeline tensor shape requires non-empty sequence.")
 
     tp_size = int(getattr(ps, "tp_size", 1) or 1)
-    if tp_size > 1:
+    cp_size = int(getattr(ps, "cp_size", 1) or 1)
+    # The model scatters THD activations into Megatron sequence-parallel + context-parallel form
+    # before the first layer, so PP-communicated hidden states carry padded_S / (CP * TP) per rank.
+    # The raw input_ids length is UNPADDED and not CP-divided; sizing the P2P buffer from it makes the
+    # receiver wait for CP*TP-too-many elements -> NCCL recv hang. Replicate thd.pack_nested_thd padding
+    # (each sequence padded to align = tp * (2*cp if cp>1 else 1)) then divide by CP*TP.
+    seq_lens = getattr(batch, "seq_lens", None)
+    if seq_lens is not None and cp_size * tp_size > 1:
+        sl = seq_lens if isinstance(seq_lens, torch.Tensor) else torch.as_tensor(seq_lens)
+        sl = sl.to(torch.int64).reshape(-1)
+        align = tp_size * (2 * cp_size if cp_size > 1 else 1)
+        padded = sl + (align - sl % align) % align
+        total_padded = int(padded.sum().item())
+        local_seq_len = total_padded // (cp_size * tp_size)
+    elif tp_size > 1:
         if local_seq_len % tp_size != 0:
             raise ValueError(
                 f"Pipeline tensor sequence length {local_seq_len} is not divisible by TP={tp_size}."
             )
-        # Megatron Lite Qwen3.5 scatters embeddings into Megatron sequence-parallel form
-        # before the first layer, so PP activations carry S / (CP * TP).
         local_seq_len //= tp_size
 
-    return (local_seq_len, batch_size, int(model_cfg.hidden_size))
+    # Models with multi-head hyper-connections (e.g. DeepSeek V4) carry hc_mult
+    # parallel residual streams across pipeline stages. The inter-stage tensor folds
+    # hc_mult into the hidden dim ([B, S, hc_mult * H]); size the P2P buffer to match.
+    # hc_mult defaults to 1, so this is a no-op for every other model.
+    hc_mult = int(getattr(model_cfg, "hc_mult", 1) or 1)
+    return (local_seq_len, batch_size, int(model_cfg.hidden_size) * hc_mult)
 
 
 def _last_loss_output(outputs: list[dict]) -> dict:
@@ -195,8 +215,8 @@ def build_model(
             if callable(reload_model_params):
                 reload_model_params()
 
-        # ── forward_step default ──
-        forward_fn = bundle.forward_step or (lambda m, b: m(**b))
+        if bundle.forward_step is None:
+            raise ValueError("Megatron Lite model bundles must provide a typed forward_step.")
 
         p = rt_cfg.parallel
         model = bundle.chunks[0] if len(bundle.chunks) == 1 else bundle.chunks
@@ -209,7 +229,7 @@ def build_model(
             _extras={
                 "model_chunks": bundle.chunks,
                 "model_cfg": model_cfg,
-                "forward_step": forward_fn,
+                "forward_step": bundle.forward_step,
                 "protocol": proto,
                 "finalize_grads": bundle.finalize_grads,
                 "world_size": dist.get_world_size(),
@@ -391,8 +411,9 @@ def forward_backward(
 
             from megatron.lite.primitive.parallel.pipeline import forward_backward_pipelining
 
-            first_batch = next(data_iter)
-            data_iter = chain([first_batch], data_iter)
+            first_item = next(data_iter)
+            first_batch, _loss_context = split_loss_context(first_item)
+            data_iter = chain([first_item], data_iter)
             tensor_shape = _infer_pipeline_tensor_shape(
                 first_batch, handle._extras.get("model_cfg"), ps
             )
@@ -429,6 +450,7 @@ def forward_backward(
                 dist_opt=not forward_only,
                 pre_forward_hook=handle._extras.get("pre_forward_hook"),
                 loss_fn=loss_fn,
+                forward_only=forward_only,
             )
 
         if not forward_only:
diff --git a/experimental/lite/megatron/lite/runtime/megatron_utils.py b/experimental/lite/megatron/lite/runtime/megatron_utils.py
index 129868f08a5..54b5ec61843 100644
--- a/experimental/lite/megatron/lite/runtime/megatron_utils.py
+++ b/experimental/lite/megatron/lite/runtime/megatron_utils.py
@@ -82,12 +82,24 @@ def register_training_hooks(model_list: list, optimizer) -> None:
 # ======================================================================
 
 
+def _is_megatron_ddp(model_chunk: Any) -> bool:
+    try:
+        from megatron.core.distributed import DistributedDataParallel as DDP
+    except Exception:
+        return False
+
+    return (
+        isinstance(model_chunk, DDP)
+        and hasattr(model_chunk, "buffers")
+        and hasattr(model_chunk, "expert_parallel_buffers")
+        and hasattr(model_chunk, "module")
+    )
+
+
 def offload_model_to_cpu(model_list: list) -> None:
     """Offload DDP model to CPU via buffer-resize (zero-copy on GPU side)."""
-    from megatron.core.distributed import DistributedDataParallel as DDP
-
     for model_chunk in model_list:
-        if isinstance(model_chunk, DDP):
+        if _is_megatron_ddp(model_chunk):
             all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
             for buffers in all_buffers:
                 for buffer in buffers:
@@ -109,10 +121,8 @@ def offload_model_to_cpu(model_list: list) -> None:
 
 def load_model_to_gpu(model_list: list, load_grad: bool = True) -> None:
     """Load DDP model back to GPU from pinned CPU copy."""
-    from megatron.core.distributed import DistributedDataParallel as DDP
-
     for model_chunk in model_list:
-        if isinstance(model_chunk, DDP):
+        if _is_megatron_ddp(model_chunk):
             all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
             for buffers in all_buffers:
                 for buffer in buffers:
diff --git a/experimental/lite/tests/conftest.py b/experimental/lite/tests/conftest.py
index fa2840e238a..ad048bdefbc 100644
--- a/experimental/lite/tests/conftest.py
+++ b/experimental/lite/tests/conftest.py
@@ -49,8 +49,11 @@ def install() -> None:
             import transformer_engine.pytorch  # noqa: F401
 
             return
-        except ModuleNotFoundError as exc:
-            if exc.name not in {"transformer_engine", "transformer_engine.pytorch"}:
+        except (ModuleNotFoundError, OSError) as exc:
+            if isinstance(exc, ModuleNotFoundError) and exc.name not in {
+                "transformer_engine",
+                "transformer_engine.pytorch",
+            }:
                 raise
 
         class _UnavailableTE:
@@ -64,8 +67,38 @@ def __init__(self, *args, **kwargs):
         pytorch.LayerNormLinear = _UnavailableTE
         pytorch.Linear = _UnavailableTE
         pytorch.RMSNorm = _UnavailableTE
+        permutation = types.ModuleType("transformer_engine.pytorch.permutation")
+        router = types.ModuleType("transformer_engine.pytorch.router")
+        cpp_extensions = types.ModuleType("transformer_engine.pytorch.cpp_extensions")
+        module = types.ModuleType("transformer_engine.pytorch.module")
+        module_base = types.ModuleType("transformer_engine.pytorch.module.base")
+
+        def unavailable_kernel(*args, **kwargs):
+            raise RuntimeError("Transformer Engine fused kernel is not installed.")
+
+        permutation.moe_permute = unavailable_kernel
+        permutation.moe_permute_and_pad_with_probs = unavailable_kernel
+        permutation.moe_permute_with_probs = unavailable_kernel
+        permutation.moe_unpermute = unavailable_kernel
+        router.fused_compute_score_for_moe_aux_loss = unavailable_kernel
+        router.fused_moe_aux_loss = unavailable_kernel
+        router.fused_topk_with_score_function = unavailable_kernel
+        cpp_extensions.general_gemm = lambda *args, **kwargs: None
+        module_base.get_workspace = lambda: None
+        module.base = module_base
+        pytorch.permutation = permutation
+        pytorch.router = router
+        pytorch.cpp_extensions = cpp_extensions
+        pytorch.module = module
         root.pytorch = pytorch
         monkeypatch.setitem(sys.modules, "transformer_engine", root)
         monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", pytorch)
+        monkeypatch.setitem(sys.modules, "transformer_engine.pytorch.permutation", permutation)
+        monkeypatch.setitem(sys.modules, "transformer_engine.pytorch.router", router)
+        monkeypatch.setitem(
+            sys.modules, "transformer_engine.pytorch.cpp_extensions", cpp_extensions
+        )
+        monkeypatch.setitem(sys.modules, "transformer_engine.pytorch.module", module)
+        monkeypatch.setitem(sys.modules, "transformer_engine.pytorch.module.base", module_base)
 
     return install
diff --git a/experimental/lite/tests/run_layering_contracts.sh b/experimental/lite/tests/run_layering_contracts.sh
new file mode 100755
index 00000000000..967ef50b4d5
--- /dev/null
+++ b/experimental/lite/tests/run_layering_contracts.sh
@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
+cd "$REPO_ROOT"
+
+export PYTHONPATH="$REPO_ROOT:$REPO_ROOT/experimental/lite${PYTHONPATH:+:$PYTHONPATH}"
+pytest -q experimental/lite/tests/unit/runtime/test_layering_contracts.py "$@"
diff --git a/experimental/lite/tests/smoke/primitive/test_distopt_checkpoint_smoke.py b/experimental/lite/tests/smoke/primitive/test_distopt_checkpoint_smoke.py
index e7732201191..00561ef825b 100644
--- a/experimental/lite/tests/smoke/primitive/test_distopt_checkpoint_smoke.py
+++ b/experimental/lite/tests/smoke/primitive/test_distopt_checkpoint_smoke.py
@@ -6,6 +6,7 @@
 
 import pytest
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.tensor import Replicate, Shard
 
@@ -15,13 +16,12 @@
 from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
 from megatron.core.transformer import TransformerConfig
 from megatron.lite.primitive.ckpt import attach_model_sharded_state_dict
-from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_stack
+from megatron.lite.primitive.optimizers.megatron_wrap import build_dist_opt_stack
 from megatron.lite.primitive.parallel import ParallelState, init_parallel
 from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime
 from megatron.lite.runtime.contracts.config import OptimizerConfig as LiteOptimizerConfig
 from megatron.lite.runtime.contracts.config import ParallelConfig
 from megatron.lite.runtime.contracts.handle import ModelHandle
-from tests.unit_tests.test_utilities import Utils
 
 pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed]
 
@@ -58,18 +58,42 @@ def forward(self, x):
 
 
 @pytest.fixture(scope="module", autouse=True)
-def _single_node_cuda_distopt():
+def _single_node_cuda_dist_opt():
     if not torch.cuda.is_available():
-        pytest.skip("CUDA is required for distopt smoke tests.")
+        pytest.skip("CUDA is required for dist_opt smoke tests.")
     if int(os.environ.get("WORLD_SIZE", "1")) > 8:
         pytest.skip("Megatron Lite smoke tests are capped at single-node 8 GPUs.")
 
-    Utils.set_world_size(
-        int(os.environ.get("WORLD_SIZE", "1")), int(os.environ.get("LOCAL_RANK", "0"))
-    )
-    Utils.initialize_model_parallel()
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29541")
+
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    created_pg = False
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl", init_method="env://")
+        created_pg = True
+
+    from megatron.core import parallel_state as mpu
+
+    created_mpu = False
+    if not mpu.is_initialized():
+        mpu.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            virtual_pipeline_model_parallel_size=None,
+            context_parallel_size=1,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=1,
+        )
+        created_mpu = True
     yield
-    Utils.destroy_model_parallel()
+    if created_mpu:
+        mpu.destroy_model_parallel()
+    if created_pg and dist.is_initialized():
+        dist.destroy_process_group()
 
 
 def _global_tensor(shape: tuple[int, ...], offset: float) -> torch.Tensor:
@@ -93,7 +117,7 @@ def _is_expert_param(name: str) -> bool:
     return "experts" in name
 
 
-def _build_sharded_model_and_distopt(parallel: ParallelConfig):
+def _build_sharded_model_and_dist_opt(parallel: ParallelConfig):
     ps = init_parallel(parallel)
     model = TinyTopologyAwareState(ps)
     model_cfg = SimpleNamespace(
@@ -110,7 +134,7 @@ def _build_sharded_model_and_distopt(parallel: ParallelConfig):
         optimizer=LiteOptimizerConfig(optimizer="adam", lr=1.0e-3, weight_decay=0.0),
         deterministic=False,
     )
-    wrapped_chunks, optimizer = build_mc_stack(
+    wrapped_chunks, optimizer = build_dist_opt_stack(
         [model], model_cfg=model_cfg, engine_cfg=engine_cfg, ps=ps, is_expert=_is_expert_param
     )
     _seed_optimizer_state(optimizer)
@@ -145,7 +169,7 @@ def _inner_optimizers(optimizer):
         yield optimizer
 
 
-def _distopt_handle(wrapped_chunks, optimizer, ps: ParallelState, parallel: ParallelConfig):
+def _dist_opt_handle(wrapped_chunks, optimizer, ps: ParallelState, parallel: ParallelConfig):
     return ModelHandle(
         model=wrapped_chunks,
         optimizer=optimizer,
@@ -160,7 +184,7 @@ def _distopt_handle(wrapped_chunks, optimizer, ps: ParallelState, parallel: Para
     )
 
 
-def _build_model_and_distopt():
+def _build_model_and_dist_opt():
     torch.manual_seed(2468)
     model = TinyDense().bfloat16().cuda()
     ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=True)
@@ -210,10 +234,10 @@ def _assert_model_close(lhs, rhs):
         torch.testing.assert_close(lhs_params[name], rhs_params[name], atol=0.0, rtol=0.0)
 
 
-def test_distopt_checkpoint_load_matches_uninterrupted_training_single_node(tmp_path):
-    model_for_ckpt, optimizer_for_ckpt = _build_model_and_distopt()
-    direct_model, direct_optimizer = _build_model_and_distopt()
-    loaded_model, loaded_optimizer = _build_model_and_distopt()
+def test_dist_opt_checkpoint_load_matches_uninterrupted_training_single_node(tmp_path):
+    model_for_ckpt, optimizer_for_ckpt = _build_model_and_dist_opt()
+    direct_model, direct_optimizer = _build_model_and_dist_opt()
+    loaded_model, loaded_optimizer = _build_model_and_dist_opt()
     runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime)
 
     torch.manual_seed(1357)
@@ -250,9 +274,9 @@ def test_distopt_checkpoint_load_matches_uninterrupted_training_single_node(tmp_
     _assert_model_close(direct_model, loaded_model)
 
 
-def test_distopt_checkpoint_reshards_from_pp_ep_to_tp_pp_ep_etp(tmp_path):
+def test_dist_opt_checkpoint_reshards_from_pp_ep_to_tp_pp_ep_etp(tmp_path):
     if torch.distributed.get_world_size() < 8:
-        pytest.skip("TP2/PP2/EP2/ETP2 distopt reshard smoke requires 8 GPUs.")
+        pytest.skip("TP2/PP2/EP2/ETP2 dist_opt reshard smoke requires 8 GPUs.")
 
     runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime)
     checkpoint_root = _shared_tmp_path(tmp_path)
@@ -261,25 +285,25 @@ def test_distopt_checkpoint_reshards_from_pp_ep_to_tp_pp_ep_etp(tmp_path):
     source_parallel = ParallelConfig(tp=1, ep=2, etp=1, pp=2, cp=1)
     target_parallel = ParallelConfig(tp=2, ep=2, etp=2, pp=2, cp=1)
 
-    source_chunks, source_optimizer, source_ps = _build_sharded_model_and_distopt(source_parallel)
+    source_chunks, source_optimizer, source_ps = _build_sharded_model_and_dist_opt(source_parallel)
     runtime.save_checkpoint(
-        _distopt_handle(source_chunks, source_optimizer, source_ps, source_parallel),
+        _dist_opt_handle(source_chunks, source_optimizer, source_ps, source_parallel),
         source_dir,
         step=3,
         save_rng=False,
     )
 
-    target_chunks, target_optimizer, target_ps = _build_sharded_model_and_distopt(target_parallel)
+    target_chunks, target_optimizer, target_ps = _build_sharded_model_and_dist_opt(target_parallel)
     assert (
         runtime.load_checkpoint(
-            _distopt_handle(target_chunks, target_optimizer, target_ps, target_parallel),
+            _dist_opt_handle(target_chunks, target_optimizer, target_ps, target_parallel),
             source_dir,
             load_rng=False,
         )
         == 3
     )
     runtime.save_checkpoint(
-        _distopt_handle(target_chunks, target_optimizer, target_ps, target_parallel),
+        _dist_opt_handle(target_chunks, target_optimizer, target_ps, target_parallel),
         reserialized_dir,
         step=3,
         save_rng=False,
diff --git a/experimental/lite/tests/smoke/primitive/test_qwen3_moe_distopt_checkpoint_smoke.py b/experimental/lite/tests/smoke/primitive/test_qwen3_moe_distopt_checkpoint_smoke.py
index 38fcde9856b..7cf98b11df0 100644
--- a/experimental/lite/tests/smoke/primitive/test_qwen3_moe_distopt_checkpoint_smoke.py
+++ b/experimental/lite/tests/smoke/primitive/test_qwen3_moe_distopt_checkpoint_smoke.py
@@ -12,6 +12,7 @@
 from megatron.lite.primitive.deterministic import set_deterministic
 from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime
 from megatron.lite.runtime.contracts.config import OptimizerConfig, ParallelConfig
+from megatron.lite.runtime.contracts.data import PackedBatch
 from megatron.lite.runtime.contracts.handle import ModelHandle
 
 pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed]
@@ -20,7 +21,7 @@
 def _qwen3_moe_symbols():
     te = pytest.importorskip(
         "transformer_engine.pytorch",
-        reason="Qwen3MoE distopt checkpoint smoke requires real Transformer Engine.",
+        reason="Qwen3MoE dist_opt checkpoint smoke requires real Transformer Engine.",
     )
     assert hasattr(te, "Linear"), "Qwen3MoE smoke requires real Transformer Engine Linear."
     from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig
@@ -32,7 +33,7 @@ def _qwen3_moe_symbols():
 @pytest.fixture(scope="module", autouse=True)
 def _single_node_cuda_dist():
     if not torch.cuda.is_available():
-        pytest.skip("CUDA is required for Qwen3MoE distopt checkpoint smoke tests.")
+        pytest.skip("CUDA is required for Qwen3MoE dist_opt checkpoint smoke tests.")
     if int(os.environ.get("WORLD_SIZE", "1")) > 8:
         pytest.skip("Megatron Lite smoke tests are capped at single-node 8 GPUs.")
 
@@ -90,7 +91,7 @@ def _build_handle(model_seed: int) -> ModelHandle:
     model_cfg = _tiny_qwen3_moe_config()
     impl_cfg = protocol.ImplConfig(
         parallel=parallel,
-        optimizer="mc",
+        optimizer="dist_opt",
         optimizer_config=OptimizerConfig(
             optimizer="adam", lr=1.0e-3, weight_decay=0.0, clip_grad=1.0
         ),
@@ -122,29 +123,26 @@ def _shared_tmp_path(tmp_path) -> str:
     return payload[0]
 
 
-def _random_batch(vocab_size: int) -> dict[str, torch.Tensor | bool]:
-    return {
-        "input_ids": torch.randint(0, vocab_size, (2, 4), device="cuda"),
-        "labels": torch.randint(0, vocab_size, (2, 4), device="cuda"),
-        "return_log_probs": False,
-    }
+def _random_batch(vocab_size: int) -> PackedBatch:
+    # Raw THD PackedBatch: 1-D packed tokens + true seq lengths (protocol packs).
+    return PackedBatch(
+        input_ids=torch.randint(0, vocab_size, (8,), device="cuda"),
+        labels=torch.randint(0, vocab_size, (8,), device="cuda"),
+        seq_lens=torch.full((2,), 4, dtype=torch.int64, device="cuda"),
+    )
 
 
-def _clone_batch(batch: dict[str, Any]) -> dict[str, Any]:
-    return {
-        key: value.detach().clone() if torch.is_tensor(value) else value
-        for key, value in batch.items()
-    }
+def _clone_batch(batch: PackedBatch) -> PackedBatch:
+    return PackedBatch(
+        input_ids=batch.input_ids.detach().clone(),
+        labels=batch.labels.detach().clone(),
+        seq_lens=batch.seq_lens.detach().clone(),
+    )
 
 
-def _assert_batch_equal(actual: dict[str, Any], expected: dict[str, Any]) -> None:
-    assert actual.keys() == expected.keys()
-    for key, expected_value in expected.items():
-        actual_value = actual[key]
-        if torch.is_tensor(expected_value):
-            assert torch.equal(actual_value, expected_value), key
-        else:
-            assert actual_value == expected_value
+def _assert_batch_equal(actual: PackedBatch, expected: PackedBatch) -> None:
+    assert torch.equal(actual.input_ids, expected.input_ids)
+    assert torch.equal(actual.labels, expected.labels)
 
 
 def _train_step(runtime: MegatronLiteRuntime, handle: ModelHandle, batch: dict[str, Any]) -> None:
@@ -170,9 +168,9 @@ def _assert_params_bitwise_equal(lhs: ModelHandle, rhs: ModelHandle) -> None:
         torch.testing.assert_close(lhs_params[name], rhs_params[name], atol=0.0, rtol=0.0)
 
 
-def test_qwen3_moe_distopt_checkpoint_restores_rng_and_continues_bitwise_tp2_pp2_ep2(tmp_path):
+def test_qwen3_moe_dist_opt_checkpoint_restores_rng_and_continues_bitwise_tp2_pp2_ep2(tmp_path):
     if dist.get_world_size() != 8:
-        pytest.skip("Qwen3MoE tp2/pp2/ep2 distopt checkpoint smoke requires exactly 8 GPUs.")
+        pytest.skip("Qwen3MoE tp2/pp2/ep2 dist_opt checkpoint smoke requires exactly 8 GPUs.")
 
     set_deterministic(2026)
     model_cfg = _tiny_qwen3_moe_config()
diff --git a/experimental/lite/tests/unit/primitive/test_checkpoint_runtime.py b/experimental/lite/tests/unit/primitive/test_checkpoint_runtime.py
index de25fbdae05..c2377b89f62 100644
--- a/experimental/lite/tests/unit/primitive/test_checkpoint_runtime.py
+++ b/experimental/lite/tests/unit/primitive/test_checkpoint_runtime.py
@@ -83,7 +83,7 @@ def test_runtime_local_checkpoint_load_matches_uninterrupted_training(tmp_path):
 
 
 class DistOptLike:
-    """Small optimizer wrapper with the same checkpoint contract as distopt."""
+    """Small optimizer wrapper with the same checkpoint contract as dist_opt."""
 
     def __init__(self, optimizer: torch.optim.Optimizer):
         self.optimizer = optimizer
@@ -101,11 +101,11 @@ def step(self):
 
     def state_dict(self):
         state = self.optimizer.state_dict()
-        state["distopt_like_marker"] = {"load_calls": self.load_calls}
+        state["dist_opt_like_marker"] = {"load_calls": self.load_calls}
         return state
 
     def load_state_dict(self, state):
-        marker = state.pop("distopt_like_marker")
+        marker = state.pop("dist_opt_like_marker")
         self.load_calls = int(marker["load_calls"]) + 1
         self.optimizer.load_state_dict(state)
 
diff --git a/experimental/lite/tests/unit/primitive/test_checkpoint_unit.py b/experimental/lite/tests/unit/primitive/test_checkpoint_unit.py
index 52008febbbc..9f0c816bc56 100644
--- a/experimental/lite/tests/unit/primitive/test_checkpoint_unit.py
+++ b/experimental/lite/tests/unit/primitive/test_checkpoint_unit.py
@@ -74,7 +74,7 @@ def test_runtime_checkpoint_load_matches_uninterrupted_training(tmp_path):
 
 
 class DistOptLike:
-    """Small optimizer wrapper with the same checkpoint contract as distopt."""
+    """Small optimizer wrapper with the same checkpoint contract as dist_opt."""
 
     def __init__(self, optimizer: torch.optim.Optimizer):
         self.optimizer = optimizer
@@ -91,11 +91,11 @@ def step(self):
 
     def state_dict(self):
         state = self.optimizer.state_dict()
-        state["distopt_like_marker"] = {"load_calls": self.load_calls}
+        state["dist_opt_like_marker"] = {"load_calls": self.load_calls}
         return state
 
     def load_state_dict(self, state):
-        marker = state.pop("distopt_like_marker")
+        marker = state.pop("dist_opt_like_marker")
         self.load_calls = int(marker["load_calls"]) + 1
         self.optimizer.load_state_dict(state)
 
diff --git a/experimental/lite/tests/unit/primitive/test_dist_opt_validation.py b/experimental/lite/tests/unit/primitive/test_dist_opt_validation.py
index 5af9fdc4e74..ffe9b112020 100644
--- a/experimental/lite/tests/unit/primitive/test_dist_opt_validation.py
+++ b/experimental/lite/tests/unit/primitive/test_dist_opt_validation.py
@@ -3,7 +3,10 @@
 
 import pytest
 
-from megatron.lite.primitive.optimizers.megatron_wrap import validate_mc_config, validate_mc_session
+from megatron.lite.primitive.optimizers.megatron_wrap import (
+    validate_dist_opt_config,
+    validate_dist_opt_session,
+)
 from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig
 from megatron.lite.runtime.contracts.config import ParallelConfig
 
@@ -13,14 +16,14 @@ def _engine_cfg(*, model_name: str, pp: int = 1, vpp: int = 1) -> MegatronLiteCo
 
 
 def test_dist_opt_validation_accepts_model_agnostic_config():
-    validate_mc_config(_engine_cfg(model_name="synthetic_custom_model", pp=1, vpp=1))
+    validate_dist_opt_config(_engine_cfg(model_name="synthetic_custom_model", pp=1, vpp=1))
 
 
 def test_dist_opt_validation_keeps_vpp_parallel_constraint():
     with pytest.raises(ValueError, match="dist_opt requires pp>1 when vpp>1"):
-        validate_mc_config(_engine_cfg(model_name="synthetic_custom_model", pp=1, vpp=2))
+        validate_dist_opt_config(_engine_cfg(model_name="synthetic_custom_model", pp=1, vpp=2))
 
 
-def test_validate_mc_session_alias_matches_config_validator():
-    assert validate_mc_session is validate_mc_config
-    validate_mc_session(_engine_cfg(model_name="another_synthetic_model", pp=2, vpp=2))
+def test_validate_dist_opt_session_alias_matches_config_validator():
+    assert validate_dist_opt_session is validate_dist_opt_config
+    validate_dist_opt_session(_engine_cfg(model_name="another_synthetic_model", pp=2, vpp=2))
diff --git a/experimental/lite/tests/unit/primitive/test_training_checkpoint.py b/experimental/lite/tests/unit/primitive/test_training_checkpoint.py
index f07cda7858d..e42b2c06254 100644
--- a/experimental/lite/tests/unit/primitive/test_training_checkpoint.py
+++ b/experimental/lite/tests/unit/primitive/test_training_checkpoint.py
@@ -4,9 +4,12 @@
 import copy
 from types import SimpleNamespace
 
+import pytest
 import torch
 from torch.distributed.tensor import Replicate, Shard
 
+pytest.importorskip("megatron.core.dist_checkpointing")
+
 from megatron.core.dist_checkpointing.strategies.torch import (
     _replace_state_dict_keys_with_sharded_keys,
 )
@@ -100,7 +103,7 @@ def load_state_dict(self, *args, **kwargs):
 }
 
 
-def test_distopt_checkpoint_dispatches_to_mcore_distckpt(monkeypatch, tmp_path) -> None:
+def test_dist_opt_checkpoint_dispatches_to_mcore_distckpt(monkeypatch, tmp_path) -> None:
     model = torch.nn.Linear(4, 2)
     optimizer = FakeDistOpt()
     ps = ParallelState(pp_rank=1, tp_rank=2, dp_cp_rank=3)
@@ -128,7 +131,7 @@ def fake_save(state_dict, checkpoint_dir, **kwargs):
     assert not (tmp_path / "step_5" / "optimizer_rank_0.pt").exists()
 
 
-def test_distopt_checkpoint_offsets_cover_tp_pp_ep_etp_topology() -> None:
+def test_dist_opt_checkpoint_offsets_cover_tp_pp_ep_etp_topology() -> None:
     ps = ParallelState(
         pp_size=2,
         pp_rank=1,
@@ -160,7 +163,7 @@ def test_distopt_checkpoint_offsets_cover_tp_pp_ep_etp_topology() -> None:
     assert expert_replica == (0, 0, 0)
 
 
-def test_distopt_replica_id_groups_sharded_axes_by_placement() -> None:
+def test_dist_opt_replica_id_groups_sharded_axes_by_placement() -> None:
     placements = [Replicate(), Replicate(), Replicate(), Shard(0)]
     rank_offsets0, replica_id0 = _rank_offsets_and_replica_id(
         placements, ParallelState(tp_size=2, tp_rank=0), expert=False
@@ -183,7 +186,7 @@ def test_distopt_replica_id_groups_sharded_axes_by_placement() -> None:
     assert expert_replica_id == (0, 0, 0)
 
 
-def test_distopt_replica_id_does_not_treat_pp_as_a_replica_axis() -> None:
+def test_dist_opt_replica_id_does_not_treat_pp_as_a_replica_axis() -> None:
     rank_offsets, replica_id = _rank_offsets_and_replica_id(
         [Replicate(), Replicate(), Replicate(), Shard(0)],
         ParallelState(pp_size=2, pp_rank=1, tp_size=2, tp_rank=1),
@@ -202,7 +205,7 @@ def test_distopt_replica_id_does_not_treat_pp_as_a_replica_axis() -> None:
     assert replica_id == (0, 0, 0)
 
 
-def test_distopt_pp_rank_one_model_keys_survive_torch_dist_main_replica_filter() -> None:
+def test_dist_opt_pp_rank_one_model_keys_survive_torch_dist_main_replica_filter() -> None:
     ps = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True)
     model = torch.nn.Linear(4, 2)
     attach_model_sharded_state_dict([model], ps)
@@ -215,7 +218,7 @@ def test_distopt_pp_rank_one_model_keys_survive_torch_dist_main_replica_filter()
     assert set(filtered_sd) == {"model_pp1.weight", "model_pp1.bias"}
 
 
-def test_distopt_model_state_keys_are_pp_and_vpp_aware() -> None:
+def test_dist_opt_model_state_keys_are_pp_and_vpp_aware() -> None:
     ps = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True)
     single_chunk = torch.nn.Linear(4, 2)
     attach_model_sharded_state_dict([single_chunk], ps)
@@ -240,7 +243,7 @@ def test_distopt_model_state_keys_are_pp_and_vpp_aware() -> None:
     assert _single_or_all_model_state(vpp_sd) is vpp_sd
 
 
-def test_distopt_checkpoint_loads_from_mcore_distckpt(monkeypatch, tmp_path) -> None:
+def test_dist_opt_checkpoint_loads_from_mcore_distckpt(monkeypatch, tmp_path) -> None:
     wrapped_module = torch.nn.Linear(4, 2)
     model = FakeWrapper(wrapped_module)
     optimizer = FakeDistOpt()
@@ -271,7 +274,7 @@ def fake_load(sharded_state_dict, checkpoint_dir, **kwargs):
     assert optimizer.loaded_state == {"loaded": True}
 
 
-def test_distopt_step_sync_traverses_multi_optimizer_chain_without_optimizer_property() -> None:
+def test_dist_opt_step_sync_traverses_multi_optimizer_chain_without_optimizer_property() -> None:
     class FakeTorchOptimizer:
         def __init__(self, steps):
             self.state = {
diff --git a/experimental/lite/tests/unit/runtime/test_layering_contracts.py b/experimental/lite/tests/unit/runtime/test_layering_contracts.py
new file mode 100644
index 00000000000..c336579dd3f
--- /dev/null
+++ b/experimental/lite/tests/unit/runtime/test_layering_contracts.py
@@ -0,0 +1,235 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Static layering guards for MLite public data boundaries and imports.
+
+Five layers — bench, verl_mlite, runtime, model, primitive — have directional
+import boundaries (``test_layer_import_boundaries``). On top of that the bench
+and verl connector layers must hand the runtime only a model-agnostic batch:
+``packed_seq_params`` / ``position_ids`` are *transient* THD metadata that may be
+materialised only at the immediate forward boundary, inside the explicitly marked
+allow range in the bridge runtime.
+
+The guard enforces the contract; it does not police prose. Substring scans ignore
+string/comment content (so a docstring may still document a dependency), and the
+import denylist honours :data:`IMPORT_ALLOWLIST` for vetted, optional cross-layer
+imports whose reason is recorded inline.
+"""
+
+from __future__ import annotations
+
+import ast
+import io
+import tokenize
+from collections.abc import Iterable
+from pathlib import Path
+
+LITE_ROOT = Path(__file__).resolve().parents[3]
+BENCH_ROOT = LITE_ROOT / "examples" / "bench"
+VERL_MLITE_ROOT = LITE_ROOT / "examples" / "verl" / "verl_mlite"
+RUNTIME_ROOT = LITE_ROOT / "megatron" / "lite" / "runtime"
+MODEL_ROOT = LITE_ROOT / "megatron" / "lite" / "model"
+PRIMITIVE_ROOT = LITE_ROOT / "megatron" / "lite" / "primitive"
+BRIDGE_RUNTIME = RUNTIME_ROOT / "backends" / "bridge" / "runtime.py"
+
+ALLOW_BEGIN = "MLITE_LAYERING_ALLOW_BRIDGE_FORWARD_METADATA_BEGIN"
+ALLOW_END = "MLITE_LAYERING_ALLOW_BRIDGE_FORWARD_METADATA_END"
+LAYER_ROOTS = {
+    "bench": BENCH_ROOT,
+    "verl_mlite": VERL_MLITE_ROOT,
+    "runtime": RUNTIME_ROOT,
+    "model": MODEL_ROOT,
+    "primitive": PRIMITIVE_ROOT,
+}
+MODEL_PACKAGE_PREFIXES = (
+    "megatron.lite.model.deepseek_v4",
+    "megatron.lite.model.glm5",
+    "megatron.lite.model.kimi_k2",
+    "megatron.lite.model.qwen3_5",
+    "megatron.lite.model.qwen3_moe",
+)
+MODEL_NAME_TERMS = {"deepseek_v4", "glm5", "kimi_k2", "qwen3", "qwen3_5", "qwen3_moe"}
+DENIED_IMPORT_PREFIXES = {
+    "bench": ("examples.verl", "verl", "verl_mlite", "megatron.lite.model"),
+    "verl_mlite": ("examples.bench", *MODEL_PACKAGE_PREFIXES),
+    "runtime": ("examples", "verl", "verl_mlite", *MODEL_PACKAGE_PREFIXES),
+    "model": ("examples", "verl", "verl_mlite", "megatron.lite.runtime.backends"),
+    "primitive": (
+        "examples",
+        "verl",
+        "verl_mlite",
+        "megatron.lite.model",
+        "megatron.lite.runtime.backends",
+        "megatron.lite.runtime.megatron_utils",
+    ),
+}
+
+# Vetted cross-layer imports that override the denylist. Each entry is a single
+# file -> {allowed module prefix: reason}. Use this ONLY for a sanctioned,
+# self-contained optional dependency — never to paper over a structural leak.
+IMPORT_ALLOWLIST: dict[str, dict[str, str]] = {
+    # The primitive linear-cross-entropy op uses VERL's Triton-backed fused kernel
+    # as an optional CUDA fast path and falls back to the local torch implementation
+    # when the kernel (or verl) is not importable. It is a legitimate optional
+    # dependency on a single leaf op, not a connector back-edge, so it overrides the
+    # primitive `verl` denylist.
+    "megatron/lite/primitive/ops/linear_cross_entropy.py": {
+        "verl.utils.kernel.linear_cross_entropy": (
+            "optional VERL fused linear-cross-entropy kernel fast-path with local fallback"
+        ),
+    },
+}
+
+
+def _python_files(root: Path) -> list[Path]:
+    return sorted(path for path in root.rglob("*.py") if path.is_file())
+
+
+def _matches_prefix(module: str, prefix: str) -> bool:
+    return module == prefix or module.startswith(prefix + ".")
+
+
+def _allowlisted_import(rel_path: str, module: str) -> bool:
+    for allowed in IMPORT_ALLOWLIST.get(rel_path, {}):
+        if _matches_prefix(module, allowed):
+            return True
+    return False
+
+
+def _imported_modules(path: Path) -> list[tuple[int, str]]:
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    imports: list[tuple[int, str]] = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            imports.extend((node.lineno, alias.name) for alias in node.names)
+        elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module:
+            imports.append((node.lineno, node.module))
+    return imports
+
+
+def _code_lines(path: Path) -> list[str]:
+    """Return source lines with string/comment spans blanked out.
+
+    The contract guard targets executable references, not documentation: a
+    docstring or comment may legitimately mention ``packed_seq_params`` or a model
+    name to explain provenance. Blanking string/comment tokens keeps those out of
+    the substring scan while preserving exact line numbers. Falls back to raw lines
+    if the file does not tokenize (fail toward flagging, never toward hiding).
+    """
+    text = path.read_text(encoding="utf-8")
+    lines = text.splitlines()
+    blanked = [list(line) for line in lines]
+    masked_types = {tokenize.STRING, tokenize.COMMENT}
+    fstring_middle = getattr(tokenize, "FSTRING_MIDDLE", None)
+    if fstring_middle is not None:
+        masked_types.add(fstring_middle)
+    try:
+        tokens = list(tokenize.generate_tokens(io.StringIO(text).readline))
+    except (tokenize.TokenError, IndentationError, SyntaxError):
+        return lines
+    for tok in tokens:
+        if tok.type not in masked_types:
+            continue
+        (srow, scol), (erow, ecol) = tok.start, tok.end
+        for row in range(srow, erow + 1):
+            chars = blanked[row - 1]
+            start = scol if row == srow else 0
+            end = ecol if row == erow else len(chars)
+            for col in range(start, min(end, len(chars))):
+                chars[col] = " "
+    return ["".join(chars) for chars in blanked]
+
+
+def _allow_ranges(path: Path) -> list[range]:
+    if path != BRIDGE_RUNTIME:
+        return []
+
+    ranges: list[range] = []
+    start: int | None = None
+    for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
+        if ALLOW_BEGIN in line:
+            if start is not None:
+                raise AssertionError(f"nested allow range in {path}")
+            start = lineno
+        elif ALLOW_END in line:
+            if start is None:
+                raise AssertionError(f"unmatched allow range end in {path}:{lineno}")
+            ranges.append(range(start, lineno + 1))
+            start = None
+    if start is not None:
+        raise AssertionError(f"unclosed allow range in {path}")
+    return ranges
+
+
+def _violations(paths: Iterable[Path], denied_terms: set[str]) -> list[str]:
+    found: list[str] = []
+    for path in paths:
+        ranges = _allow_ranges(path)
+        for lineno, line in enumerate(_code_lines(path), start=1):
+            if any(lineno in allowed for allowed in ranges):
+                continue
+            for term in sorted(denied_terms):
+                if term in line:
+                    rel = path.relative_to(LITE_ROOT)
+                    found.append(f"{rel}:{lineno}: {term}")
+    return found
+
+
+def _import_violations(layer: str) -> list[str]:
+    denied = DENIED_IMPORT_PREFIXES[layer]
+    found: list[str] = []
+    for path in _python_files(LAYER_ROOTS[layer]):
+        rel = path.relative_to(LITE_ROOT).as_posix()
+        for lineno, module in _imported_modules(path):
+            if _allowlisted_import(rel, module):
+                continue
+            for prefix in denied:
+                if _matches_prefix(module, prefix):
+                    found.append(f"{rel}:{lineno}: {module} matches denied {prefix}")
+    return found
+
+
+def test_layer_import_boundaries() -> None:
+    violations = []
+    for layer in LAYER_ROOTS:
+        violations.extend(_import_violations(layer))
+    assert violations == []
+
+
+def test_import_allowlist_entries_are_live() -> None:
+    """Each allowlisted import must still exist — stop the allowlist from rotting."""
+    stale: list[str] = []
+    for rel_path, allowed in IMPORT_ALLOWLIST.items():
+        path = LITE_ROOT / rel_path
+        if not path.is_file():
+            stale.append(f"{rel_path}: file missing")
+            continue
+        modules = {module for _, module in _imported_modules(path)}
+        for allowed_module in allowed:
+            if not any(_matches_prefix(module, allowed_module) for module in modules):
+                stale.append(f"{rel_path}: no import matches allowlisted {allowed_module}")
+    assert stale == []
+
+
+def test_bench_layer_does_not_see_model_internal_batch_fields() -> None:
+    violations = _violations(
+        _python_files(BENCH_ROOT),
+        {"packed_seq_params", "position_ids", "to_bridge_dict"},
+    )
+    assert violations == []
+
+
+def test_verl_mlite_layer_does_not_see_model_internal_batch_fields() -> None:
+    violations = _violations(
+        _python_files(VERL_MLITE_ROOT),
+        {"packed_seq_params", "position_ids", "to_bridge_dict"},
+    )
+    assert violations == []
+
+
+def test_runtime_packed_seq_params_is_bridge_forward_transient_only() -> None:
+    violations = _violations(_python_files(RUNTIME_ROOT), {"packed_seq_params"})
+    assert violations == []
+
+
+def test_primitive_layer_is_model_name_agnostic() -> None:
+    violations = _violations(_python_files(PRIMITIVE_ROOT), MODEL_NAME_TERMS)
+    assert violations == []
diff --git a/experimental/lite/tests/unit/runtime/test_packed_batch_bridge.py b/experimental/lite/tests/unit/runtime/test_packed_batch_bridge.py
new file mode 100644
index 00000000000..0ef4629679b
--- /dev/null
+++ b/experimental/lite/tests/unit/runtime/test_packed_batch_bridge.py
@@ -0,0 +1,142 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Unit coverage for the PackedBatch -> bridge forward metadata boundary.
+
+CPU-only. The public ``PackedBatch`` contract stays free of transient THD metadata
+(``packed_seq_params``); the bridge runtime renders Megatron-Core THD kwargs at the
+forward call using the same canonical packing as the native lite protocols, so the
+mlite-vs-bridge comparison is fair and context-parallel-correct.
+"""
+
+from __future__ import annotations
+
+import torch
+from examples.bench.session import _infinite_packed_batches
+from megatron.lite.primitive.parallel.thd import (
+    pack_nested_thd,
+    reconstruct_packed_from_cp_parts,
+)
+from megatron.lite.runtime.backends.bridge.runtime import (
+    _bridge_forward_kwargs_from_packed_batch,
+    _nested_from_packed,
+)
+from megatron.lite.runtime.contracts.data import Batch, PackedBatch
+
+
+def _packed_batch() -> PackedBatch:
+    return PackedBatch(
+        input_ids=torch.arange(8, dtype=torch.long),
+        labels=torch.arange(100, 108, dtype=torch.long),
+        seq_lens=torch.tensor([3, 5], dtype=torch.int64),
+    )
+
+
+def test_packed_batch_contract_carries_no_transient_thd_metadata() -> None:
+    # The contract keeps its optional extensibility slots (position_ids for custom
+    # layouts, routed_experts for router replay, extras for multimodal) but must
+    # never carry the transient Megatron-Core THD metadata the bridge derives.
+    batch = _packed_batch()
+    assert hasattr(batch, "position_ids")
+    assert hasattr(batch, "routed_experts")
+    assert hasattr(batch, "extras")
+    assert batch.position_ids is None
+    assert not hasattr(batch, "packed_seq_params")
+    assert not hasattr(batch, "to_bridge_dict")
+
+
+def test_packed_batch_is_batch_subclass() -> None:
+    assert issubclass(PackedBatch, Batch)
+
+
+def test_bridge_forward_kwargs_are_transient_bridge_metadata() -> None:
+    batch = _packed_batch()
+    out = _bridge_forward_kwargs_from_packed_batch(batch)
+
+    assert set(out) == {"input_ids", "labels", "position_ids", "packed_seq_params"}
+    assert out["input_ids"].shape == (1, 8)
+    assert out["labels"].shape == (1, 8)
+    assert out["position_ids"].shape == (1, 8)
+    # tp=cp=1 -> no padding, input ids unchanged.
+    assert torch.equal(out["input_ids"].reshape(-1), batch.input_ids)
+    # Labels are rolled one position left per sequence (last token zeroed), exactly
+    # like the native pack_thd_forward_kwargs path, so both backends train on the
+    # same shifted targets.
+    assert torch.equal(
+        out["labels"].reshape(-1),
+        torch.tensor([101, 102, 0, 104, 105, 106, 107, 0]),
+    )
+    assert torch.equal(
+        out["position_ids"].reshape(-1),
+        torch.tensor([0, 1, 2, 0, 1, 2, 3, 4]),
+    )
+
+    psp = out["packed_seq_params"]
+    assert psp.qkv_format == "thd"
+    assert torch.equal(psp.cu_seqlens_q, torch.tensor([0, 3, 8], dtype=torch.int32))
+    assert psp.max_seqlen_q == 5
+
+
+def test_bridge_forward_kwargs_carry_loss_mask_only_inside_bridge() -> None:
+    batch = _packed_batch()
+    batch.loss_mask = torch.tensor([1, 1, 0, 1, 1, 1, 0, 1], dtype=torch.long)
+    out = _bridge_forward_kwargs_from_packed_batch(batch)
+    assert "loss_mask" in out
+    # loss_mask is rolled with the labels so it masks the shifted targets.
+    assert torch.equal(
+        out["loss_mask"].reshape(-1),
+        torch.tensor([1, 0, 0, 1, 1, 0, 1, 0]),
+    )
+
+
+def test_bridge_forward_kwargs_are_context_parallel_correct() -> None:
+    # cp_size=2 must pad to the zigzag (2*cp) alignment, keep full cu_seqlens, and
+    # hand each rank only its half of the tokens — the exact behaviour the previous
+    # hand-rolled implementation lacked.
+    batch = _packed_batch()
+    seq_lens = batch.seq_lens
+
+    # Full CP-aligned reference (padded but not yet CP-split).
+    ref = pack_nested_thd(
+        _nested_from_packed(batch.input_ids, seq_lens),
+        tp_size=1,
+        cp_size=2,
+        cp_rank=0,
+        split_cp=False,
+    )
+    full_padded = int(ref.cu_seqlens_padded[-1].item())
+    assert full_padded == 12  # [3->4, 5->8] padded to 2*cp alignment
+
+    rank0 = _bridge_forward_kwargs_from_packed_batch(batch, cp_size=2, cp_rank=0)
+    rank1 = _bridge_forward_kwargs_from_packed_batch(batch, cp_size=2, cp_rank=1)
+
+    for local in (rank0, rank1):
+        assert local["input_ids"].shape == (1, full_padded // 2)
+        assert local["position_ids"].shape == (1, full_padded // 2)
+        psp = local["packed_seq_params"]
+        # cu_seqlens stays full (CP-aligned), not the unpadded [0, 3, 8].
+        assert torch.equal(psp.cu_seqlens_q, ref.cu_seqlens_padded)
+        assert int(getattr(psp, "local_cp_size", 1)) == 2
+
+    # The two rank-local zigzag shards reconstruct the full padded sequence.
+    recon = reconstruct_packed_from_cp_parts(
+        [rank0["input_ids"][0], rank1["input_ids"][0]],
+        cu_seqlens_padded=ref.cu_seqlens_padded,
+        cp_size=2,
+        dim=0,
+    )
+    assert torch.equal(recon, ref.input_ids[0])
+
+
+def test_infinite_packed_batches_shape_and_determinism() -> None:
+    gen_a = _infinite_packed_batches(vocab_size=32, seq_len=6, device="cpu", seed=7)
+    gen_b = _infinite_packed_batches(vocab_size=32, seq_len=6, device="cpu", seed=7)
+
+    a = next(gen_a)
+    b = next(gen_b)
+    assert isinstance(a, PackedBatch)
+    assert a.input_ids.shape == (6,)
+    assert a.labels.shape == (6,)
+    assert torch.equal(a.seq_lens, torch.tensor([6], dtype=torch.int64))
+    # No transient THD metadata baked into bench data.
+    assert a.position_ids is None
+    assert torch.equal(a.input_ids, b.input_ids)
+    assert torch.equal(a.labels, b.labels)
diff --git a/experimental/lite/tests/unit/runtime/test_runtime_backend_unit.py b/experimental/lite/tests/unit/runtime/test_runtime_backend_unit.py
index c71df0f4bff..c3e2ab888fa 100644
--- a/experimental/lite/tests/unit/runtime/test_runtime_backend_unit.py
+++ b/experimental/lite/tests/unit/runtime/test_runtime_backend_unit.py
@@ -3,6 +3,8 @@
 
 import os
 import subprocess
+import sys
+import types
 from dataclasses import dataclass
 from pathlib import Path
 from unittest.mock import MagicMock, patch
@@ -183,6 +185,151 @@ def test_runtime_to_prefers_optimizer_specific_offload_hooks():
     assert optimizer.calls == ["offload", "load"]
 
 
+class _FakeStorage:
+    def __init__(self, size: int):
+        self._size = size
+        self.resize_calls: list[int] = []
+
+    def size(self):
+        return self._size
+
+    def resize_(self, size: int):
+        self.resize_calls.append(size)
+        self._size = size
+        return self
+
+
+class _FakeBufferData:
+    def __init__(self, size: int):
+        self._storage = _FakeStorage(size)
+        self.cpu_called = False
+        self.pinned = False
+        self.copied_from = None
+        self.copy_non_blocking = None
+        self.zero_calls = 0
+
+    @property
+    def data(self):
+        return self
+
+    def cpu(self):
+        self.cpu_called = True
+        return self
+
+    def pin_memory(self):
+        self.pinned = True
+        return self
+
+    def storage(self):
+        return self._storage
+
+    def copy_(self, other, *, non_blocking: bool):
+        self.copied_from = other
+        self.copy_non_blocking = non_blocking
+        return self
+
+    def zero_(self):
+        self.zero_calls += 1
+        return self
+
+
+class _FakeBuffer:
+    def __init__(self):
+        self.param_data = _FakeBufferData(3)
+        self.grad_data = _FakeBufferData(5)
+
+
+class _FakeModule:
+    def parameters(self):
+        return []
+
+
+class _FakeMegatronDDP:
+    def __init__(self):
+        self.buffer = _FakeBuffer()
+        self.buffers = [self.buffer]
+        self.expert_parallel_buffers = []
+        self.module = _FakeModule()
+        self.to_calls: list[str] = []
+
+    def to(self, device):
+        self.to_calls.append(device)
+        raise AssertionError("DDP model chunks must use the buffer offload path")
+
+
+class _FakeMegatronDDPSubclass(_FakeMegatronDDP):
+    pass
+
+
+class _FakeNativeModel:
+    def __init__(self):
+        self.calls: list[str] = []
+
+    def to(self, device):
+        self.calls.append(device)
+        return self
+
+
+def _install_fake_megatron_ddp(monkeypatch) -> None:
+    core = types.ModuleType("megatron.core")
+    distributed = types.ModuleType("megatron.core.distributed")
+    distributed.DistributedDataParallel = _FakeMegatronDDP
+    core.distributed = distributed
+    monkeypatch.setitem(sys.modules, "megatron.core", core)
+    monkeypatch.setitem(sys.modules, "megatron.core.distributed", distributed)
+
+
+def test_megatron_ddp_detection_accepts_ddp_and_subclasses(monkeypatch):
+    from megatron.lite.runtime.megatron_utils import _is_megatron_ddp
+
+    _install_fake_megatron_ddp(monkeypatch)
+
+    assert _is_megatron_ddp(_FakeMegatronDDP()) is True
+    assert _is_megatron_ddp(_FakeMegatronDDPSubclass()) is True
+    assert _is_megatron_ddp(_FakeNativeModel()) is False
+
+
+@pytest.mark.parametrize("model_cls", [_FakeMegatronDDP, _FakeMegatronDDPSubclass])
+def test_megatron_ddp_model_move_helpers_use_buffer_path(monkeypatch, model_cls):
+    from megatron.lite.runtime.megatron_utils import load_model_to_gpu, offload_model_to_cpu
+
+    _install_fake_megatron_ddp(monkeypatch)
+    model = model_cls()
+    buffer = model.buffer
+
+    offload_model_to_cpu([model])
+
+    assert model.to_calls == []
+    assert buffer.param_data.cpu_called is True
+    assert buffer.param_data.pinned is True
+    assert buffer.param_data_size == 3
+    assert buffer.grad_data_size == 5
+    assert buffer.param_data.storage().size() == 0
+    assert buffer.grad_data.storage().size() == 0
+
+    load_model_to_gpu([model])
+
+    assert model.to_calls == []
+    assert buffer.param_data.storage().size() == 3
+    assert buffer.grad_data.storage().size() == 5
+    assert buffer.param_data.copied_from is buffer.param_data.cpu_data
+    assert buffer.param_data.copy_non_blocking is True
+    assert buffer.grad_data.zero_calls == 1
+
+
+def test_native_model_move_helpers_do_not_require_megatron_core(monkeypatch):
+    from megatron.lite.runtime.megatron_utils import load_model_to_gpu, offload_model_to_cpu
+
+    monkeypatch.setitem(sys.modules, "megatron.core", None)
+    monkeypatch.setitem(sys.modules, "megatron.core.distributed", None)
+    model = _FakeNativeModel()
+
+    offload_model_to_cpu([model])
+    load_model_to_gpu([model])
+
+    assert model.calls == ["cpu", "cuda"]
+
+
 def test_model_handle_dp_defaults():
     handle = ModelHandle(model=MagicMock())
 
diff --git a/experimental/lite/tests/unit/verl/test_mlite_engine_config.py b/experimental/lite/tests/unit/verl/test_mlite_engine_config.py
index fa6a9f71442..f90fd7f1aff 100644
--- a/experimental/lite/tests/unit/verl/test_mlite_engine_config.py
+++ b/experimental/lite/tests/unit/verl/test_mlite_engine_config.py
@@ -1,8 +1,10 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 from types import SimpleNamespace
 
+import pytest
+
 from verl_mlite.engine.config import MegatronLiteEngineConfig
-from verl_mlite.engine.mlite_engine import MegatronLiteEngine
+from verl_mlite.engine.mlite_engine import MegatronLiteEngine, _build_lr_scheduler
 
 
 def _optimizer_config(**override_optimizer_config) -> SimpleNamespace:
@@ -111,3 +113,36 @@ def test_mlite_config_threads_rl_parallel_and_impl_settings() -> None:
     assert config.attention_backend_override == "flash"
     assert config.impl_cfg["use_thd"] is True
     assert config.impl_cfg["deterministic"] is False
+
+
+def test_local_lr_scheduler_warmup_decay_and_state_roundtrip() -> None:
+    optimizer = SimpleNamespace(param_groups=[{"lr": 0.0, "weight_decay": 0.1}])
+    opt = SimpleNamespace(
+        total_training_steps=4,
+        lr_warmup_steps=1,
+        lr_warmup_steps_ratio=0.0,
+        lr_warmup_init=0.0,
+        lr=1.0,
+        min_lr=0.1,
+        lr_decay_steps=4,
+        lr_decay_style="linear",
+        weight_decay=0.1,
+        weight_decay_incr_style="constant",
+        lr_wsd_decay_steps=None,
+        lr_wsd_decay_style="exponential",
+    )
+
+    scheduler = _build_lr_scheduler(optimizer, opt)
+
+    assert optimizer.param_groups[0]["lr"] == 0.0
+    scheduler.step(1)
+    assert optimizer.param_groups[0]["lr"] == 1.0
+    scheduler.step(1)
+    assert optimizer.param_groups[0]["lr"] == pytest.approx(0.7)
+
+    state = scheduler.state_dict()
+    scheduler.step(10)
+    scheduler.load_state_dict(state)
+
+    assert scheduler.state_dict() == state
+    assert optimizer.param_groups[0]["lr"] == pytest.approx(0.7)

From 5af7fa139ae5153ccb4328f6a99ee3eac2a84f3f Mon Sep 17 00:00:00 2001
From: Yan Bai <bayan@nvidia.com>
Date: Wed, 24 Jun 2026 07:51:35 -0700
Subject: [PATCH 3/4] [dev] Megatron Lite (3/4): fused kernels

Signed-off-by: Yan Bai <bayan@nvidia.com>
---
 .../lite/primitive/kernels/dsa_kernels.py     | 1146 +++++++++++++++++
 .../lite/primitive/kernels/grouped_gemm.py    |   26 +
 2 files changed, 1172 insertions(+)
 create mode 100644 experimental/lite/megatron/lite/primitive/kernels/dsa_kernels.py
 create mode 100644 experimental/lite/megatron/lite/primitive/kernels/grouped_gemm.py

diff --git a/experimental/lite/megatron/lite/primitive/kernels/dsa_kernels.py b/experimental/lite/megatron/lite/primitive/kernels/dsa_kernels.py
new file mode 100644
index 00000000000..7e56ec30388
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/kernels/dsa_kernels.py
@@ -0,0 +1,1146 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+"""
+DSA kernel wrappers for Megatron's DSv4 sparse attention.
+
+Mirrors the three integration paths of the old standalone ``dsa_kernels``
+package, but built on top of
+
+* :mod:`cudnn.deepseek_sparse_attention` (a.k.a. ``DSA``) — CuTe-DSL backward
+  + indexer score kernels + TRT-LLM radix top-K, shipped as part of
+  cuDNN Frontend.
+* :mod:`flash_mla` — production sparse-attention forward kernel, expected to
+  be available as a separate PyPI package.
+
+Public API (same shape as the old ``dsa_kernels`` package):
+
+* ``build_flat_topk_idxs`` / ``local_to_global_flat`` — index helpers.
+* ``dsa_sparse_attn`` — Path A / Path C step 2, differentiable sparse attention.
+* ``indexer_topk`` — Path C inference indexer scoring + top-K.
+* ``fused_indexer_sparse_attn`` — Path B training, fused indexer loss +
+  sparse attention with shared backward.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Callable, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+# ---------------------------------------------------------------------------
+# Lazy kernel imports
+# ---------------------------------------------------------------------------
+
+
+_flash_mla_sparse_fwd = None
+_DSA = None
+_indexer_fwd_sm90: Optional[Callable] = None
+_indexer_fwd_sm100: Optional[Callable] = None
+
+
+def _ensure_flash_mla():
+    """Lazily import the FlashMLA sparse-forward kernel.
+
+    FlashMLA ships ``flash_mla_sparse_fwd`` with a multi-head-KV signature;
+    :func:`_dsa_fwd_flash_mla` below is a thin adapter that unbatches the
+    DSA-shape inputs and pads ``TopK`` to the alignment expected by
+    FlashMLA's SM90 / SM100 kernels.
+    """
+    global _flash_mla_sparse_fwd
+    if _flash_mla_sparse_fwd is not None:
+        return
+
+    try:
+        from flash_mla import flash_mla_sparse_fwd as _fwd
+    except ImportError as e:
+        raise ImportError(
+            "FlashMLA is required for DSA sparse attention forward. "
+            "Install from https://github.com/deepseek-ai/FlashMLA/tree/nv_dev "
+            "so that `from flash_mla import flash_mla_sparse_fwd` succeeds."
+        ) from e
+    _flash_mla_sparse_fwd = _fwd
+
+
+def _get_topk_alignment() -> int:
+    """Minimum ``TopK`` alignment required by the current GPU architecture.
+
+    * SM90 : dual-warpgroup loop steps by 2 blocks → ``2 * B_TOPK = 128``
+    * SM100: single-pipeline loop steps by 1 block → ``B_TOPK`` (64 for
+      head64, 128 for head128). DSA uses ``D = 512`` which maps to the
+      head64 kernel path → 64.
+    """
+    sm = torch.cuda.get_device_capability()
+    if sm[0] >= 10:
+        return 64
+    return 128
+
+
+def _dsa_fwd_flash_mla(
+    q: Tensor,
+    kv: Tensor,
+    topk_idxs: Tensor,
+    softmax_scale: float,
+    d_v: int = 512,
+    attn_sink: Optional[Tensor] = None,
+    topk_length: Optional[Tensor] = None,
+    indexer_topk: int = 0,
+) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+    """DSA-shaped adapter around :func:`flash_mla.flash_mla_sparse_fwd`.
+
+    Accepts flat (unbatched) tensors with global indices; pads ``TopK`` to
+    the GPU-specific alignment; returns ``(out, lse, lse_indexer)``.
+    """
+    assert not (
+        indexer_topk > 0 and topk_length is not None
+    ), "indexer_topk > 0 requires non-compact mode (topk_length must be None)"
+    _ensure_flash_mla()
+
+    _total_S_q, _H, _D = q.shape
+    TopK = topk_idxs.shape[-1]
+    topk_align = _get_topk_alignment()
+    TopK_padded = (TopK + topk_align - 1) // topk_align * topk_align
+    if TopK_padded != TopK:
+        pad_width = TopK_padded - TopK
+        topk_idxs = torch.nn.functional.pad(topk_idxs, (0, pad_width), value=-1)
+
+    kv_3d = kv.unsqueeze(1)  # (total_S_kv, 1, D)  h_kv=1
+    indices = topk_idxs.unsqueeze(1)  # (total_S_q, 1, TopK_padded) h_kv=1
+
+    with torch.cuda.nvtx.range("flash_mla_sparse_fwd"):
+        res = _flash_mla_sparse_fwd(
+            q,
+            kv_3d,
+            indices,
+            softmax_scale,
+            d_v=d_v,
+            attn_sink=attn_sink,
+            topk_length=topk_length,
+            indexer_topk=indexer_topk,
+        )
+        if indexer_topk > 0:
+            out, _max_logits, lse, lse_indexer = res
+        else:
+            out, _max_logits, lse = res
+            lse_indexer = None
+
+    if indexer_topk > 0:
+        # When indexer_topk == total TopK, lse_indexer should equal lse but
+        # the kernel may not snapshot correctly; fall back to lse.
+        if indexer_topk >= TopK:
+            return out, lse, lse.clone()
+        return out, lse, lse_indexer
+    return out, lse, None
+
+
+def _ensure_dsa_namespace():
+    """Lazily import the cudnn-frontend DSA namespace."""
+    global _DSA
+    if _DSA is not None:
+        return
+    try:
+        from cudnn import DSA as _ns
+    except ImportError as e:
+        try:
+            from cudnn.deepseek_sparse_attention import DSA as _ns
+        except ImportError:
+            raise ImportError(
+                "cudnn-frontend DSA namespace not available. Install with "
+                "`pip install nvidia-cudnn-frontend[cutedsl]`; newer "
+                "versions expose it as `cudnn.deepseek_sparse_attention.DSA`."
+            ) from e
+    _DSA = _ns
+
+
+def _load_indexer_fwd_sm90():
+    """Load the H100 SM90 indexer forward entry only when it is selected."""
+    global _indexer_fwd_sm90
+    if _indexer_fwd_sm90 is None:
+        try:
+            module = import_module(
+                "cudnn.deepseek_sparse_attention.indexer_forward._interface_sm90"
+            )
+            _indexer_fwd_sm90 = module.indexer_fwd
+        except (AttributeError, ImportError) as exc:
+            raise ImportError(
+                "H100 DSA indexer forward requires the SM90 cudnn route "
+                "`cudnn.deepseek_sparse_attention.indexer_forward._interface_sm90.indexer_fwd`."
+            ) from exc
+    return _indexer_fwd_sm90
+
+
+def _load_indexer_fwd_sm100():
+    """Load the Blackwell SM100 indexer forward entry only when it is selected."""
+    global _indexer_fwd_sm100
+    if _indexer_fwd_sm100 is None:
+        try:
+            module = import_module("cudnn.deepseek_sparse_attention.indexer_forward._interface")
+            _indexer_fwd_sm100 = module.indexer_fwd
+        except (AttributeError, ImportError) as exc:
+            raise ImportError(
+                "Blackwell DSA indexer forward requires the SM100 cudnn route "
+                "`cudnn.deepseek_sparse_attention.indexer_forward._interface.indexer_fwd` "
+                "(exported by cudnn-frontend as indexer_fwd_sm100)."
+            ) from exc
+    return _indexer_fwd_sm100
+
+
+def _select_indexer_forward(device):
+    major, _minor = torch.cuda.get_device_capability(device)
+    if major == 9:
+        return _load_indexer_fwd_sm90()
+    if major >= 10:
+        return _load_indexer_fwd_sm100()
+    return None
+
+
+def _dsa_indexer_forward_wrapper(
+    q: Tensor,
+    k: Tensor,
+    w: Tensor,
+    *,
+    ratio: int = 4,
+    qhead_per_kv_head: Optional[int] = None,
+    sm_scale: float = 1.0,
+    cu_seqlens_q: Optional[Tensor] = None,
+    cu_seqlens_k: Optional[Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+):
+    """Route indexer forward to architecture-specific CuTe-DSL backends."""
+    if q.is_cuda:
+        indexer_fwd = _select_indexer_forward(q.device)
+        if indexer_fwd is not None:
+            return {
+                "scores": indexer_fwd(
+                    q,
+                    k,
+                    w,
+                    ratio=ratio,
+                    qhead_per_kv_head=qhead_per_kv_head,
+                    sm_scale=sm_scale,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                )
+            }
+    _ensure_dsa_namespace()
+    return _DSA.indexer_forward_wrapper(
+        q,
+        k,
+        w,
+        ratio=ratio,
+        qhead_per_kv_head=qhead_per_kv_head,
+        sm_scale=sm_scale,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Index helpers
+# ---------------------------------------------------------------------------
+
+
+def local_to_global_flat(local_idxs: Tensor, batch_size: int, seqlen_kv: int) -> Tensor:
+    """Convert local per-batch indices to global flat indices.
+
+    Follows the convention used by FlashMLA / SparseAttentionBackward:
+    flat row order is SBHD ``row[s * B + b]``; global index is
+    ``local * B + b`` for valid entries and ``-1`` otherwise.
+
+    Args:
+        local_idxs: ``(b, sq, topk)`` int, values in ``[0, seqlen_kv)`` or -1.
+        batch_size: ``B``.
+        seqlen_kv: KV sequence length per batch (used for shape assertions
+            only; callers compute the values).
+
+    Returns:
+        ``(sq*b, topk)`` int32.
+    """
+    b, sq, topk = local_idxs.shape
+    assert b == batch_size
+
+    idxs_sb = local_idxs.permute(1, 0, 2).reshape(sq * b, topk)
+    valid = idxs_sb >= 0
+    batch_ids = torch.arange(sq * b, device=local_idxs.device) % b
+    batch_ids_exp = batch_ids.unsqueeze(1).expand_as(idxs_sb)
+    idxs_sb = torch.where(valid, idxs_sb * b + batch_ids_exp, idxs_sb)
+    return idxs_sb.int()
+
+
+def build_flat_topk_idxs(
+    *idx_groups: Tensor, batch_size: int, seqlen_kv: int, compact: bool = False
+) -> Tuple[Tensor, Optional[Tensor]]:
+    """Combine local per-batch index groups and convert to flat global form.
+
+    Each *idx_group* is ``(b, sq, topk_i)`` with local per-batch KV indices
+    (already in ``kv_full`` index space, i.e. with any compressed-position
+    offset applied). ``-1`` marks invalid positions.
+
+    Args:
+        *idx_groups: one or more ``(b, sq, topk_i)`` int tensors.
+        batch_size: ``B``.
+        seqlen_kv: total KV sequence length per batch.
+        compact: if True, pack valid entries to the front of each row and
+            additionally return ``topk_length``; if False, leave as-is and
+            return ``None``.
+
+    Returns:
+        ``(topk_idxs, topk_length)`` where
+        ``topk_idxs`` is ``(sq*b, total_topk)`` int32 (flat global) and
+        ``topk_length`` is ``(sq*b,)`` int32 when ``compact``, else ``None``.
+    """
+    combined = torch.cat(idx_groups, dim=-1)  # (b, sq, total_topk)
+    b, sq, total_topk = combined.shape
+
+    # Globalize first, compact second. Both ops are element-wise + (-1)-preserving,
+    # so swapping the order is a no-op for correctness; the win is that the
+    # global indices come out already in (sq*b, total_topk) flat layout, which is
+    # exactly the row order the cuDNN compactify kernel returns its per-row
+    # ``length`` in — no extra permute on the length tensor.
+    global_idxs = local_to_global_flat(combined, b, seqlen_kv)
+
+    topk_length_flat = None
+    if compact:
+        if global_idxs.is_cuda:
+            # Fast path: single warp-per-row CuTe DSL kernel from cuDNN's DSA
+            # namespace. Replaces a stable argsort + gather + sum + permute
+            # chain with one global-load + global-store per element.
+            _ensure_dsa_namespace()
+            res = _DSA.compactify_wrapper(global_idxs)
+            global_idxs, topk_length_flat = res["indices"], res["topk_length"]
+        else:
+            # CPU fallback so the unit tests that exercise this helper without
+            # CUDA still work. Production callers always go through the CUDA
+            # path above.
+            valid_mask = global_idxs >= 0
+            sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True)
+            global_idxs = global_idxs.gather(-1, sorted_indices)
+            topk_length_flat = valid_mask.sum(dim=-1).int()
+
+    return global_idxs, topk_length_flat
+
+
+# ---------------------------------------------------------------------------
+# Path A + Path C step 2: differentiable sparse attention
+# ---------------------------------------------------------------------------
+
+
+class SparseAttnFunc(torch.autograd.Function):
+    """SM100 sparse attention fwd + bwd on flat tensors.
+
+    Forward uses :mod:`flash_mla`; backward uses cuDNN Frontend's
+    :attr:`cudnn.DSA.sparse_attention_backward_wrapper`.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        q: Tensor,  # (total_sq, H, D) bf16
+        kv: Tensor,  # (total_skv, D) bf16
+        attn_sink: Tensor,  # (H,) f32
+        topk_idxs: Tensor,  # (total_sq, TopK) int32 global
+        topk_length: Optional[Tensor],  # (total_sq,) int32 or None
+        softmax_scale: float,
+        indexer_topk: int,
+        value_dim: Optional[int],
+    ) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
+        """Run FlashMLA sparse-attention forward and save tensors for backward."""
+        out, lse, lse_indexer = _dsa_fwd_flash_mla(
+            q,
+            kv,
+            topk_idxs,
+            softmax_scale,
+            attn_sink=attn_sink,
+            topk_length=topk_length,
+            indexer_topk=indexer_topk,
+            d_v=512 if value_dim is None else value_dim,
+        )
+
+        ctx.save_for_backward(q, kv, attn_sink, topk_idxs, out, lse)
+        ctx.softmax_scale = softmax_scale
+        ctx.topk_length = topk_length
+        return out, lse, lse_indexer
+
+    @staticmethod
+    def backward(ctx, dO, d_lse, d_lse_indexer):
+        """Compute sparse-attention backward via cuDNN DSA wrapper."""
+        _ensure_dsa_namespace()
+
+        q, kv, attn_sink, topk_idxs, out, lse = ctx.saved_tensors
+
+        result = _DSA.sparse_attention_backward_wrapper(
+            q,
+            kv,
+            out,
+            dO,
+            lse,
+            attn_sink,
+            topk_idxs,
+            softmax_scale=ctx.softmax_scale,
+            topk_length=ctx.topk_length,
+        )
+        dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"]
+        return dq, dkv, d_sink, None, None, None, None, None
+
+
+def dsa_sparse_attn(
+    query: Tensor,
+    kv: Tensor,
+    attn_sink: Tensor,
+    topk_idxs: Tensor,
+    softmax_scale: float,
+    topk_length: Optional[Tensor] = None,
+    indexer_topk: int = 0,
+    value_dim: Optional[int] = None,
+) -> Tensor:
+    """Sparse attention (Path A / Path C step 2).
+
+    Args:
+        query: ``(sq, b, np, d)`` bf16 SBHD.
+        kv:    ``(skv, b, d)`` bf16 SBD (K=V).
+        attn_sink: ``(np,)`` f32.
+        topk_idxs: ``(sq*b, topk)`` int32 — **flat global** indices produced
+            by :func:`build_flat_topk_idxs`.
+        softmax_scale: scalar float.
+        topk_length: ``(sq*b,)`` int32 — optional compact fast-path. Must be
+            ``None`` when ``indexer_topk > 0`` (FlashMLA constraint).
+        indexer_topk: int; ``0`` for Paths A/C, positive for Path B to enable
+            FlashMLA's ``lse_indexer`` output.
+        value_dim: FlashMLA value dimension. Defaults to ``512`` to preserve
+            the existing DSA wrapper behavior.
+
+    Returns:
+        ``(sq, b, np * d_v)`` bf16 output.
+    """
+    sq, b, np_, d = query.shape
+    skv = kv.shape[0]
+
+    q_flat = query.reshape(sq * b, np_, d)
+    kv_flat = kv.reshape(skv * b, d)
+
+    out_flat, _lse, _lse_indexer = SparseAttnFunc.apply(
+        q_flat, kv_flat, attn_sink, topk_idxs, topk_length, softmax_scale, indexer_topk, value_dim
+    )
+
+    d_v = out_flat.shape[-1]
+    return out_flat.reshape(sq, b, np_, d_v).reshape(sq, b, np_ * d_v)
+
+
+# ---------------------------------------------------------------------------
+# Path C inference: indexer scoring + top-K
+# ---------------------------------------------------------------------------
+
+
+def _indexer_topk_bshd(
+    q_bshd: Tensor, k_bsd: Tensor, w_bsh: Tensor, topk: int, ratio: int = 4
+) -> Tuple[Tensor, Tensor, Tensor]:
+    """BSHD-layout core for :func:`indexer_topk`.
+
+    Internal entry point used by both the public SBHD wrapper and Path B's
+    ``FusedIndexerSparseAttnFunc.forward`` so the SBHD→BSHD permute can be
+    performed once at the call site and reused across both the indexer
+    forward and the score-backward kernels (predict / target).
+
+    Args:
+        q_bshd: ``(b, sq, idx_nh, idx_hd)`` bf16, C-contiguous.
+        k_bsd:  ``(b, sk, idx_hd)`` bf16, C-contiguous.
+        w_bsh:  ``(b, sq, idx_nh)`` bf16, C-contiguous, **already
+            ``indexer_softmax_scale``-scaled** by the caller.
+        topk:   number of top-K indices to return per query.
+        ratio:  compression ratio for the kernel's causal mask.
+
+    Returns:
+        ``(topk_indices, topk_length, scores)`` where:
+
+        * ``topk_indices``: ``(b, sq, topk)`` int32, invalid slots ``-1``.
+        * ``topk_length``:  ``(b, sq)`` int32, per-row valid count.
+        * ``scores``: ``(b, sq, sk)`` fp32, raw scores from
+          :attr:`cudnn.DSA.indexer_forward_wrapper` with ``-inf`` on
+          causally-masked positions.
+    """
+    _ensure_dsa_namespace()
+
+    b, sq, _idx_nh, _idx_hd = q_bshd.shape
+    sk = k_bsd.shape[1]
+    device = q_bshd.device
+
+    k_bshd = k_bsd.unsqueeze(2)  # (b, sk, 1, idx_hd)
+
+    scores = _dsa_indexer_forward_wrapper(q_bshd, k_bshd, w_bsh, ratio=ratio)[
+        "scores"
+    ]  # (b, sq, sk) fp32, -inf on masked positions
+
+    # Top-K selection via the TRT-LLM CuTe-DSL radix kernel.
+    n_rows = b * sq
+    scores_flat = scores.reshape(n_rows, sk).contiguous()
+    q_idx = torch.arange(sq, device=device)
+    valid_per_q = ((q_idx + 1) // ratio).clamp(max=sk).to(torch.int32)  # (sq,)
+    seq_lens = valid_per_q.repeat(b)  # (b*sq,), row-major over (b, sq)
+
+    topk_k = min(topk, sk)
+    tk_result = _DSA.indexer_top_k_wrapper(
+        scores_flat, seq_lens, top_k=topk_k, next_n=1, return_val=False
+    )
+    topk_indices = tk_result["indices"].view(b, sq, topk_k)
+
+    if topk_k < topk:
+        pad = torch.full((b, sq, topk - topk_k), -1, dtype=torch.int32, device=device)
+        topk_indices = torch.cat([topk_indices, pad], dim=-1)
+
+    topk_length = (topk_indices >= 0).sum(dim=-1).int()  # (b, sq)
+    return topk_indices.int(), topk_length, scores
+
+
+def _sbhd_to_bshd_indexer_inputs(
+    q_indexer: Tensor, k_indexer: Tensor, weights: Tensor, indexer_softmax_scale: float
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """Permute the indexer inputs SBHD→BSHD once, returning both the raw
+    BSHD weights and (when needed) a separate scaled copy.
+
+    The ``relu(c·x) = c·relu(x)`` trick lets us push the indexer softmax
+    scale onto ``W`` (``(B, S_q, H)``, small) instead of the score tensor
+    (``(B, S_q, S_k)``, big). The raw ``w_bsh`` is preserved for the
+    backward GEMM path, which takes ``sm_scale`` directly. When
+    ``indexer_softmax_scale == 1.0`` the two views alias each other.
+
+    Returns ``(q_bshd, k_bsd, w_bsh, w_bsh_scaled)``.
+    """
+    q_bshd = q_indexer.permute(1, 0, 2, 3).contiguous()
+    k_bsd = k_indexer.permute(1, 0, 2).contiguous()
+    w_bsh = weights.permute(1, 0, 2).contiguous()
+
+    if indexer_softmax_scale != 1.0:
+        w_bsh_scaled = (w_bsh.float() * indexer_softmax_scale).to(w_bsh.dtype)
+    else:
+        w_bsh_scaled = w_bsh
+
+    return q_bshd, k_bsd, w_bsh, w_bsh_scaled
+
+
+def indexer_topk(
+    q_indexer: Tensor,
+    k_indexer: Tensor,
+    weights: Tensor,
+    topk: int,
+    ratio: int = 4,
+    indexer_softmax_scale: float = 1.0,
+) -> Tuple[Tensor, Tensor]:
+    """Score + top-K selection for inference (no KL loss, no backward).
+
+    Built on cuDNN Frontend's CuTe-DSL indexer forward kernel followed by
+    TRT-LLM's radix top-K kernel.
+
+    Args:
+        q_indexer: ``(sq, b, idx_nh, idx_hd)`` bf16 SBHD.
+        k_indexer: ``(sk, b, idx_hd)`` bf16 SBD.
+        weights:   ``(sq, b, idx_nh)`` bf16 SBH — raw (unscaled) weights.
+        topk: number of top-K indices to select.
+        ratio: compression ratio for the causal mask.
+        indexer_softmax_scale: scale applied to the indexer ``Q @ K^T``
+            scores (typically ``idx_hd ** -0.5``). Applied internally via
+            the weights-scaling trick (``relu(c·x) = c·relu(x)`` for
+            ``c > 0``) so the caller passes raw weights. Default ``1.0``
+            means weights are treated as already-scaled.
+
+    Returns:
+        topk_indices: ``(b, sq, topk)`` int32 — local per-batch indices into
+            ``k_indexer``; invalid positions are ``-1``.
+        topk_length:  ``(b, sq)`` int32 — per-query valid count.
+    """
+    q_bshd, k_bsd, _w_bsh_raw, w_bsh_scaled = _sbhd_to_bshd_indexer_inputs(
+        q_indexer, k_indexer, weights, indexer_softmax_scale
+    )
+    topk_indices, topk_length, _ = _indexer_topk_bshd(q_bshd, k_bsd, w_bsh_scaled, topk, ratio)
+    return topk_indices, topk_length
+
+
+# ---------------------------------------------------------------------------
+# Path B: fused indexer + sparse attention (training)
+# ---------------------------------------------------------------------------
+
+
+_CLIP_PROB_MIN = torch.finfo(torch.float32).tiny  # kept compatible w/ cudnn kernel
+
+
+def _compute_indexer_predict(
+    q_indexer_bshd: Tensor,
+    k_indexer_bsd: Tensor,
+    weights_bsh: Tensor,
+    topk_indices: Tensor,
+    qhead_per_kv_head: int,
+) -> Tensor:
+    """Compute ``predict`` distribution (softmax over top-K of indexer scores).
+
+    Wraps :attr:`cudnn.DSA.sparse_indexer_score_recompute_wrapper`.
+
+    Args:
+        q_indexer_bshd: ``(B, S_q, H_q, D)`` bf16.
+        k_indexer_bsd:  ``(B, S_k, D)`` bf16.
+        weights_bsh:    ``(B, S_q, H_q)`` bf16.
+        topk_indices:   ``(B, S_q, topk)`` int32.
+        qhead_per_kv_head: ``H_q`` (MQA).
+
+    Returns:
+        predict: ``(B, S_q, topk)`` fp32, softmax over the top-K axis.
+    """
+    _ensure_dsa_namespace()
+    result = _DSA.sparse_indexer_score_recompute_wrapper(
+        q_indexer_bshd,
+        k_indexer_bsd,
+        weights_bsh,
+        topk_indices,
+        qhead_per_kv_head=qhead_per_kv_head,
+    )
+    return result["predict"]
+
+
+def _compute_attn_target(
+    q_attn_bshd: Tensor,
+    k_attn_bsd: Tensor,
+    lse: Tensor,
+    topk_indices: Tensor,
+    softmax_scale: float,
+    qhead_per_kv_head: int,
+) -> Tensor:
+    """Compute ``target`` distribution (L1-normalised head-sum softmax).
+
+    Wraps :attr:`cudnn.DSA.sparse_attn_score_recompute_wrapper`.
+
+    Shapes match :func:`_compute_indexer_predict`; ``lse`` is
+    ``(B, S_q, H_q)`` FP32 (comes from the attention forward pass).
+    """
+    _ensure_dsa_namespace()
+    result = _DSA.sparse_attn_score_recompute_wrapper(
+        q_attn_bshd,
+        k_attn_bsd,
+        lse,
+        topk_indices,
+        softmax_scale,
+        qhead_per_kv_head=qhead_per_kv_head,
+    )
+    return result["target"]
+
+
+def _kl_loss_from_target_predict(
+    target: Tensor,
+    predict: Tensor,
+    topk_indices: Tensor,
+    loss_coeff: float,
+    calculate_per_token_loss: bool = False,
+) -> Tensor:
+    """KL(target || predict) reduced over ``(B, S_q)`` and scaled by loss_coeff.
+
+    Rows with no valid top-K positions (early query rows with ratio causal
+    masking) contribute 0 to the loss — the sparse score kernels produce
+    garbage for those rows, mirroring ``compute_dsa_indexer_loss``'s
+    ``row_valid`` handling. The default mean is taken over all ``(B, S_q)``
+    positions. Per-token-loss mode returns a raw local sum so finalize can
+    apply the global token divisor.
+    """
+    eps = _CLIP_PROB_MIN
+    t = target.clamp(min=eps)
+    p = predict.clamp(min=eps)
+    kl_per_row = (t * (torch.log(t) - torch.log(p))).sum(dim=-1)  # (B, S_q)
+
+    row_valid = (topk_indices >= 0).any(dim=-1)  # (B, S_q)
+    kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row))
+    loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean()
+    return loss_coeff * loss
+
+
+# ---------------------------------------------------------------------------
+# Dense path (``sparse_loss=False``) — full-KV indexer loss
+# ---------------------------------------------------------------------------
+
+
+def _compute_dense_indexer_score(
+    q_indexer_bshd: Tensor,
+    k_indexer_bshd: Tensor,
+    weights_bsh: Tensor,
+    qhead_per_kv_head: int,
+    indexer_softmax_scale: float,
+    ratio: int,
+) -> Tuple[Tensor, Tensor]:
+    """Dense indexer score forward over the full ``S_k`` axis.
+
+    Wraps :attr:`cudnn.DSA.dense_indexer_score_recompute_wrapper`. Returns
+    ``(out, denom)`` where
+
+    * ``out``    : ``(B, S_q, S_k)`` fp32, the raw head-reduced score
+      ``S[b,q,k] = indexer_softmax_scale * sum_h ReLU(Q_h · K_k^T) · W_{b,q,h}``
+      with the kernel's ``ratio``-causal mask applied to invalid columns.
+    * ``denom``  : ``(B, S_q)`` fp32, the LSE denom of ``out`` along
+      ``S_k`` — i.e. ``predict = exp(out - denom[..., None])`` is the
+      indexer softmax distribution over the full KV.
+
+    Both outputs are forwarded into :func:`_kl_loss_from_dense_scores`
+    *and* saved for the dense-path backward, where the dense indexer-grad
+    kernel consumes them directly.
+    """
+    _ensure_dsa_namespace()
+    result = _DSA.dense_indexer_score_recompute_wrapper(
+        q_indexer_bshd,
+        k_indexer_bshd,
+        weights_bsh,
+        qhead_per_kv_head=qhead_per_kv_head,
+        sm_scale=indexer_softmax_scale,
+        ratio=ratio,
+    )
+    return result["out"], result["denom"]
+
+
+def _compute_dense_attn_score(
+    q_attn_bshd: Tensor,
+    k_attn_bshd: Tensor,
+    lse: Tensor,
+    qhead_per_kv_head: int,
+    softmax_scale: float,
+    ratio: int,
+) -> Tuple[Tensor, Tensor]:
+    """Dense attention score forward over the full ``S_k`` axis.
+
+    Wraps :attr:`cudnn.DSA.dense_attn_score_recompute_wrapper`. Returns
+    ``(out, denom)`` where
+
+    * ``out``   : ``(B, S_q, S_k)`` fp32, the head-summed unnormalized
+      attention probability ``S[b,q,k] = sum_h exp(Q_h · K_k^T · scale - LSE[b,q,h])``
+      with ``ratio`` causal mask applied.
+    * ``denom`` : ``(B, S_q)`` fp32, the L1-norm denom ``sum_k S[b,q,:]``.
+      ``target = out / denom[..., None]`` is the L1-normalized
+      head-summed attention distribution.
+    """
+    _ensure_dsa_namespace()
+    result = _DSA.dense_attn_score_recompute_wrapper(
+        q_attn_bshd,
+        k_attn_bshd,
+        lse,
+        softmax_scale,
+        qhead_per_kv_head=qhead_per_kv_head,
+        ratio=ratio,
+    )
+    return result["out"], result["denom"]
+
+
+def _kl_loss_from_dense_scores(
+    attn_score: Tensor,
+    attn_l1norm: Tensor,
+    index_score: Tensor,
+    index_lse: Tensor,
+    loss_coeff: float,
+    calculate_per_token_loss: bool = False,
+) -> Tensor:
+    """KL(target || predict) over the **full** KV axis, averaged over ``(B, S_q)``.
+
+    Derives ``target = attn_score / attn_l1norm`` (L1-normalised, matches
+    ``compute_dsa_indexer_loss``'s ``attention_scores / sum`` step) and
+    ``log_predict = index_score - index_lse`` (LSE-normalised log-softmax),
+    then computes ``KL = sum_k target * (log target - log predict)`` and
+    scales by ``loss_coeff``.
+
+    Rows where the kernel's ``ratio`` causal mask leaves no valid KV
+    position have ``attn_l1norm <= 0`` (L1) or ``index_lse == -inf``
+    (LSE); those rows contribute 0 to the loss — the same ``row_valid``
+    semantics as the reference ``compute_dsa_indexer_loss``.
+    """
+    eps = _CLIP_PROB_MIN
+    # row_valid: rows with at least one un-masked KV position.
+    row_valid = (attn_l1norm > eps) & torch.isfinite(index_lse)
+
+    # Safe denoms: replace invalid rows with a finite value so target /
+    # log-predict don't produce NaN; the row mask zeroes their KL below.
+    safe_l1 = attn_l1norm.clamp(min=eps)
+    safe_lse = torch.where(row_valid, index_lse, torch.zeros_like(index_lse))
+
+    target = attn_score / safe_l1.unsqueeze(-1)
+    target_clamped = target.clamp(min=eps)
+    # Per-position validity: the indexer-score kernel emits -inf at
+    # ratio-masked positions; those contribute 0 to KL by the
+    # ``0 · log(0/p) = 0`` convention. Without this gate, the eps-clamp
+    # on target makes the term ``eps · (log eps - (-inf)) = +inf``.
+    position_valid = torch.isfinite(index_score)
+    safe_index_score = torch.where(position_valid, index_score, torch.zeros_like(index_score))
+    log_predict = safe_index_score - safe_lse.unsqueeze(-1)
+
+    kl_terms = target_clamped * (torch.log(target_clamped) - log_predict)
+    kl_terms = torch.where(position_valid, kl_terms, torch.zeros_like(kl_terms))
+    kl_per_row = kl_terms.sum(dim=-1)  # (B, S_q)
+    kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row))
+    loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean()
+    return loss_coeff * loss
+
+
+class FusedIndexerSparseAttnFunc(torch.autograd.Function):
+    """Path B: fused indexer (+KL loss) + sparse attention in one autograd.
+
+    Differentiable w.r.t. ``query``, ``kv_full``, ``attn_sink``,
+    ``q_indexer``, ``k_indexer``, ``weights``.
+
+    Two indexer-loss variants, selected by the ``sparse_loss`` argument
+    (matches ``compute_dsa_indexer_loss`` in the reference ``dsa.py``):
+
+    * **Sparse loss** (``sparse_loss=True``) — KL is computed only over
+      the top-K KV positions the indexer has selected.
+    * **Dense loss** (``sparse_loss=False``, the default) — KL is
+      computed over *all* causally valid KV positions.
+
+    Both variants share the FlashMLA sparse-attention forward + the
+    cuDNN sparse-attn backward; only the indexer-loss path branches.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        # Sparse attn inputs (differentiable)
+        query: Tensor,  # (sq, b, np, d) bf16
+        kv_full: Tensor,  # (skv, b, d) bf16
+        attn_sink: Tensor,  # (np,) f32
+        # Window indices (not differentiable)
+        window_idxs: Tensor,  # (b, sq, win_topk) int32
+        # Indexer inputs (differentiable)
+        q_indexer: Tensor,  # (sq, b, idx_nh, idx_hd) bf16
+        k_indexer: Tensor,  # (n_comp, b, idx_hd) bf16
+        weights: Tensor,  # (sq, b, idx_nh) bf16 — raw (unscaled)
+        # Scalars
+        indexer_topk: int,
+        ratio: int,
+        softmax_scale: float,
+        indexer_softmax_scale: float,
+        loss_coeff: float,
+        sparse_loss: bool,
+        kv_offset: int,
+        calculate_per_token_loss: bool,
+        value_dim: Optional[int],
+    ) -> Tuple[Tensor, Tensor]:
+        """Fused forward: indexer scoring, sparse attention, KL loss, and indexer backward."""
+        _ensure_dsa_namespace()
+
+        sq, b, np_, d = query.shape
+        skv = kv_full.shape[0]
+        n_comp = k_indexer.shape[0]
+        idx_nh, idx_hd = q_indexer.shape[2], q_indexer.shape[3]
+
+        requested_topk = indexer_topk
+        effective_topk = min(requested_topk, n_comp)
+
+        # ---- 1. Permute indexer inputs SBHD->BSHD ONCE. -------------------
+        q_idx_bshd, k_idx_bsd, w_bsh, w_bsh_scaled = _sbhd_to_bshd_indexer_inputs(
+            q_indexer, k_indexer, weights, indexer_softmax_scale
+        )
+
+        # ---- 2. Indexer scoring + top-K (with scores retained). -------------
+        topk_indices_cmp, _, indexer_scores = _indexer_topk_bshd(
+            q_idx_bshd, k_idx_bsd, w_bsh_scaled, effective_topk, ratio
+        )  # topk_indices_cmp: (b, sq, effective_topk) int32; indexer_scores: (b, sq, n_comp) fp32
+
+        # ---- 3. Combine indices (indexer first, then window). --------------
+        compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + kv_offset, -1)
+        if requested_topk > effective_topk:
+            pad = torch.full(
+                (b, sq, requested_topk - effective_topk),
+                -1,
+                device=compress_topk_idxs.device,
+                dtype=compress_topk_idxs.dtype,
+            )
+            compress_topk_idxs = torch.cat([compress_topk_idxs, pad], dim=-1)
+        combined_local = torch.cat([compress_topk_idxs, window_idxs], dim=-1)
+        global_idxs = local_to_global_flat(combined_local, b, skv)
+
+        # ---- 4. FlashMLA forward (non-compact, indexer_topk > 0). ---------
+        q_flat = query.reshape(sq * b, np_, d)
+        kv_flat = kv_full.reshape(skv * b, d)
+        out_flat, lse, lse_indexer = _dsa_fwd_flash_mla(
+            q_flat,
+            kv_flat,
+            global_idxs,
+            softmax_scale,
+            attn_sink=attn_sink,
+            topk_length=None,
+            indexer_topk=requested_topk,
+            d_v=512 if value_dim is None else value_dim,
+        )
+
+        # ---- 5. Derive predict from indexer_scores, compute target. --------
+        # Attention-path tensors (detached — loss is not differentiable through them).
+        q_attn_bshd = query.detach().permute(1, 0, 2, 3).contiguous()
+        k_attn_compressed_bsd = kv_full[kv_offset:].detach().permute(1, 0, 2).contiguous()
+        lse_indexer_bsqh = lse_indexer.reshape(sq, b, np_).permute(1, 0, 2)
+
+        if sparse_loss:
+            # Derive predict: gather topk scores from indexer_scores → softmax.
+            safe_indices = topk_indices_cmp.clamp(min=0).long()
+            gathered_scores = torch.gather(indexer_scores, dim=2, index=safe_indices)
+            gathered_scores = torch.where(
+                topk_indices_cmp >= 0, gathered_scores, torch.finfo(torch.float32).min
+            )
+            predict = torch.softmax(gathered_scores, dim=-1)  # (b, sq, topk) fp32
+
+            target = _compute_attn_target(
+                q_attn_bshd,
+                k_attn_compressed_bsd,
+                lse_indexer_bsqh,
+                topk_indices_cmp,
+                softmax_scale,
+                qhead_per_kv_head=np_,
+            )
+
+            if loss_coeff > 0:
+                indexer_loss = _kl_loss_from_target_predict(
+                    target, predict, topk_indices_cmp, loss_coeff, calculate_per_token_loss
+                )
+            else:
+                indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32)
+        else:
+            # Dense: use full indexer_scores directly + logsumexp.
+            index_score = indexer_scores  # (b, sq, n_comp) fp32
+            index_lse = torch.logsumexp(indexer_scores, dim=-1)  # (b, sq) fp32
+
+            attn_score, attn_l1norm = _compute_dense_attn_score(
+                q_attn_bshd,
+                k_attn_compressed_bsd.unsqueeze(2),
+                lse_indexer_bsqh,
+                qhead_per_kv_head=np_,
+                softmax_scale=softmax_scale,
+                ratio=ratio,
+            )
+
+            if loss_coeff > 0:
+                indexer_loss = _kl_loss_from_dense_scores(
+                    attn_score,
+                    attn_l1norm,
+                    index_score,
+                    index_lse,
+                    loss_coeff,
+                    calculate_per_token_loss,
+                )
+            else:
+                indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32)
+
+        # ---- 6. Eagerly compute indexer backward (grad_loss=1). ------------
+        # The actual grad_loss scaling is deferred to backward (when
+        # DSAIndexerLossAutoScaler provides the correct scale).
+        indexer_loss_coeff = loss_coeff
+        if calculate_per_token_loss:
+            indexer_loss_coeff = loss_coeff * (b * sq)
+
+        unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32)
+
+        if loss_coeff > 0:
+            if sparse_loss:
+                attn_score_for_bwd = target.clone()
+                index_score_for_bwd = predict.clone()
+                ig = _DSA.indexer_backward_wrapper(
+                    q_idx_bshd,
+                    w_bsh,
+                    k_idx_bsd,
+                    attn_score_for_bwd,
+                    index_score_for_bwd,
+                    topk_indices_cmp,
+                    sm_scale=indexer_softmax_scale,
+                    loss_coeff=indexer_loss_coeff,
+                    grad_loss=unit_grad_loss,
+                    block_I=128,
+                )
+            else:
+                attn_score_for_bwd = attn_score.clone()
+                index_score_for_bwd = index_score.clone()
+                ig = _DSA.dense_indexer_backward_wrapper(
+                    q_idx_bshd,
+                    w_bsh,
+                    k_idx_bsd,
+                    attn_score_for_bwd,
+                    attn_l1norm,
+                    index_score_for_bwd,
+                    index_lse,
+                    sm_scale=indexer_softmax_scale,
+                    loss_coeff=indexer_loss_coeff,
+                    grad_loss=unit_grad_loss,
+                    ratio=ratio,
+                    block_I=128,
+                )
+            # BSHD -> SBHD (match input layout).
+            precomputed_grad_q_indexer = ig["d_index_q"].permute(1, 0, 2, 3).contiguous()
+            precomputed_grad_k_indexer = ig["d_index_k"].permute(1, 0, 2).contiguous()
+            precomputed_grad_weights = ig["d_weights"].permute(1, 0, 2).contiguous()
+        else:
+            precomputed_grad_q_indexer = torch.zeros_like(q_indexer)
+            precomputed_grad_k_indexer = torch.zeros_like(k_indexer)
+            precomputed_grad_weights = torch.zeros_like(weights)
+
+        # ---- 7. Save context (only sparse-attn bwd tensors + indexer grads).
+        ctx.save_for_backward(
+            q_flat,
+            kv_flat,
+            attn_sink,
+            global_idxs,
+            out_flat,
+            lse,
+            precomputed_grad_q_indexer,
+            precomputed_grad_k_indexer,
+            precomputed_grad_weights,
+        )
+        ctx.softmax_scale = softmax_scale
+        ctx.sq = sq
+        ctx.b = b
+        ctx.np_ = np_
+        ctx.d = d
+        ctx.skv = skv
+
+        # ---- 8. Return. ---------------------------------------------------
+        d_v = out_flat.shape[-1]
+        output = out_flat.reshape(sq, b, np_, d_v).reshape(sq, b, np_ * d_v)
+        return output, indexer_loss
+
+    @staticmethod
+    def backward(ctx, grad_output, grad_loss):
+        """Backward: sparse attention bwd + scale pre-computed indexer grads."""
+        (
+            q_flat,
+            kv_flat,
+            attn_sink,
+            global_idxs,
+            out_flat,
+            lse,
+            precomputed_grad_q_indexer,
+            precomputed_grad_k_indexer,
+            precomputed_grad_weights,
+        ) = ctx.saved_tensors
+
+        sq, b, np_, d = ctx.sq, ctx.b, ctx.np_, ctx.d
+        skv = ctx.skv
+
+        # ---- 1. Sparse attn backward. -------------------------------------
+        d_v = out_flat.shape[-1]
+        dO_flat = grad_output.reshape(sq * b, np_, d_v)
+
+        attn_bwd = _DSA.sparse_attention_backward_wrapper(
+            q_flat,
+            kv_flat,
+            out_flat,
+            dO_flat,
+            lse,
+            attn_sink,
+            global_idxs,
+            softmax_scale=ctx.softmax_scale,
+            topk_length=None,
+        )
+        grad_query = attn_bwd["dq"].reshape(sq, b, np_, d)
+        grad_kv_full = attn_bwd["dkv"].reshape(skv, b, d)
+        d_sink = attn_bwd["d_sink"]
+
+        # ---- 2. Scale pre-computed indexer grads by grad_loss. -------------
+        grad_q_indexer = precomputed_grad_q_indexer * grad_loss
+        grad_k_indexer = precomputed_grad_k_indexer * grad_loss
+        grad_weights = precomputed_grad_weights * grad_loss
+
+        # Grads: query, kv_full, attn_sink, window_idxs, q_indexer, k_indexer,
+        #   weights, indexer_topk, ratio, softmax_scale, indexer_softmax_scale,
+        #   loss_coeff, sparse_loss, kv_offset, calculate_per_token_loss, value_dim
+        return (
+            grad_query,
+            grad_kv_full,
+            d_sink,
+            None,
+            grad_q_indexer,
+            grad_k_indexer,
+            grad_weights,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+def fused_indexer_sparse_attn(
+    query: Tensor,
+    kv_full: Tensor,
+    attn_sink: Tensor,
+    window_idxs: Tensor,
+    q_indexer: Tensor,
+    k_indexer: Tensor,
+    weights: Tensor,
+    indexer_topk: int,
+    ratio: int,
+    softmax_scale: float,
+    indexer_softmax_scale: float = 1.0,
+    loss_coeff: float = 0.0,
+    sparse_loss: bool = False,
+    kv_offset: int = 0,
+    calculate_per_token_loss: bool = False,
+    value_dim: Optional[int] = None,
+) -> Tuple[Tensor, Tensor]:
+    """Path B (training): fused indexer (+KL loss) + sparse attention.
+
+    See :class:`FusedIndexerSparseAttnFunc` for the detailed data flow.
+
+    Args:
+        query:        ``(sq, b, np, d)`` bf16 SBHD — attention query.
+        kv_full:      ``(skv, b, d)`` bf16 SBD — original + compressed KV.
+        attn_sink:    ``(np,)`` f32 — learnable sink per head.
+        window_idxs:  ``(b, sq, win_topk)`` int32 — local window indices.
+        q_indexer:    ``(sq, b, idx_nh, idx_hd)`` bf16 — indexer query.
+        k_indexer:    ``(n_comp, b, idx_hd)`` bf16 — indexer key (compressed).
+        weights:      ``(sq, b, idx_nh)`` bf16 — raw indexer weights.
+        indexer_topk: number of top-K compressed positions to select.
+        ratio:        compression ratio used for the causal mask.
+        softmax_scale: attention ``Q @ K^T`` scale, typically
+            ``1/sqrt(v_head_dim)``.
+        indexer_softmax_scale: indexer ``Q @ K^T`` scale, typically
+            ``1/sqrt(idx_hd)``. Applied internally — caller passes raw
+            (unscaled) ``weights``.
+        loss_coeff:   coefficient scaling the KL divergence loss.
+        sparse_loss:  if ``True``, KL is computed only over the top-K
+            positions (cheap, less informative); if ``False`` (the
+            default, matches ``transformer_config.dsa_indexer_use_sparse_loss``),
+            KL is computed over the full causally-valid KV (more
+            informative, matches the DeepSeek-V3.2 paper, larger
+            intermediate-tensor footprint). See
+            :class:`FusedIndexerSparseAttnFunc` for the full data flow
+            of each variant.
+        kv_offset:    start of compressed region within ``kv_full``.
+        calculate_per_token_loss: if True, report raw local KL sum and
+            compensate the cuDNN backward wrappers' local averaging.
+        value_dim: FlashMLA value dimension. Defaults to ``512`` to preserve
+            the existing DSA wrapper behavior.
+
+    Returns:
+        ``(output, indexer_loss)`` where ``output`` is ``(sq, b, np * d_v)``
+        bf16 and ``indexer_loss`` is a scalar f32.
+    """
+    return FusedIndexerSparseAttnFunc.apply(
+        query,
+        kv_full,
+        attn_sink,
+        window_idxs,
+        q_indexer,
+        k_indexer,
+        weights,
+        indexer_topk,
+        ratio,
+        softmax_scale,
+        indexer_softmax_scale,
+        loss_coeff,
+        sparse_loss,
+        kv_offset,
+        calculate_per_token_loss,
+        value_dim,
+    )
+
+
+__all__ = [
+    "build_flat_topk_idxs",
+    "local_to_global_flat",
+    "dsa_sparse_attn",
+    "indexer_topk",
+    "fused_indexer_sparse_attn",
+]
diff --git a/experimental/lite/megatron/lite/primitive/kernels/grouped_gemm.py b/experimental/lite/megatron/lite/primitive/kernels/grouped_gemm.py
new file mode 100644
index 00000000000..957eac77fdb
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/kernels/grouped_gemm.py
@@ -0,0 +1,26 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Optional grouped GEMM kernel accessors."""
+
+from __future__ import annotations
+
+try:
+    import grouped_gemm  # pyright: ignore[reportMissingImports]
+except Exception:  # pragma: no cover - optional fused kernel or missing shared libraries.
+    grouped_gemm = None  # type: ignore[assignment]
+
+
+def grouped_gemm_is_available() -> bool:
+    return grouped_gemm is not None
+
+
+def assert_grouped_gemm_is_available() -> None:
+    if grouped_gemm is None:
+        raise AssertionError(
+            "Grouped GEMM is not available. Please install the grouped_gemm package."
+        )
+
+
+ops = grouped_gemm.ops if grouped_gemm_is_available() else None
+
+
+__all__ = ["assert_grouped_gemm_is_available", "grouped_gemm_is_available", "ops"]

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

Signed-off-by: Yan Bai <bayan@nvidia.com>
---
 .../primitive/modules/attention/__init__.py   |  16 +
 .../lite/primitive/modules/attention/cp.py    |  93 +++
 .../lite/primitive/modules/attention/csa.py   | 714 ++++++++++++++++++
 .../lite/primitive/modules/attention/dsa.py   | 697 +++++++++++++++++
 .../lite/primitive/modules/attention/hca.py   |  66 ++
 .../lite/primitive/modules/attention/mhc.py   |  27 +
 .../lite/primitive/modules/attention/mla.py   | 418 ++++++++++
 7 files changed, 2031 insertions(+)
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/attention/__init__.py
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/attention/cp.py
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/attention/csa.py
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/attention/dsa.py
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/attention/hca.py
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/attention/mhc.py
 create mode 100644 experimental/lite/megatron/lite/primitive/modules/attention/mla.py

diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/__init__.py b/experimental/lite/megatron/lite/primitive/modules/attention/__init__.py
new file mode 100644
index 00000000000..f13e7670125
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/attention/__init__.py
@@ -0,0 +1,16 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+from megatron.lite.primitive.modules.attention.dsa import (
+    DynamicSparseAttention,
+    RMSNorm,
+    build_rope_cache,
+    build_rotary_embeddings,
+)
+from megatron.lite.primitive.modules.attention.mla import MultiLatentAttention
+
+__all__ = [
+    "DynamicSparseAttention",
+    "MultiLatentAttention",
+    "RMSNorm",
+    "build_rope_cache",
+    "build_rotary_embeddings",
+]
diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/cp.py b/experimental/lite/megatron/lite/primitive/modules/attention/cp.py
new file mode 100644
index 00000000000..957e2e60d9a
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/attention/cp.py
@@ -0,0 +1,93 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from torch.distributed.nn.functional import all_gather
+
+
+def _all_gather_cp(tensor: torch.Tensor, group: dist.ProcessGroup) -> list[torch.Tensor]:
+    return list(all_gather(tensor.contiguous(), group=group))
+
+
+def iter_cp_sources(tensor, position_ids, *, cp_rank, cp_size, cp_group):
+    if cp_size <= 1:
+        yield cp_rank, tensor, position_ids
+        return
+    if cp_group is None:
+        raise RuntimeError("CP source iteration requires a context-parallel process group.")
+    tensor_parts = _all_gather_cp(tensor, cp_group)
+    position_parts = _all_gather_cp(position_ids.to(dtype=torch.long), cp_group)
+    for rank, (source_tensor, source_positions) in enumerate(zip(tensor_parts, position_parts)):
+        yield rank, source_tensor, source_positions
+
+
+def _gather_contiguous_tail(tensor, *, tail_len, cp_size, cp_group, seq_dim):
+    if cp_size <= 1 or tail_len <= 0:
+        return None
+    if cp_group is None:
+        raise RuntimeError("CP chunk-tail gather requires a context-parallel process group.")
+    if tensor.size(seq_dim) < tail_len:
+        raise ValueError(f"CP chunk tail needs len >= {tail_len}, got {tensor.size(seq_dim)}.")
+    tail = tensor.narrow(seq_dim, tensor.size(seq_dim) - tail_len, tail_len)
+    return _all_gather_cp(tail.contiguous(), cp_group)
+
+
+def compress_contiguous_chunks_for_cp(
+    compressor,
+    tensor,
+    *,
+    position_ids,
+    cp_rank,
+    cp_size,
+    cp_group,
+    compress_kwargs: dict[str, Any] | None = None,
+    seq_dim=1,
+    compressed_seq_dim=2,
+):
+    kwargs = compress_kwargs or {}
+    compress_ratio = int(compressor.compress_ratio)
+    if cp_size <= 1:
+        compressed = compressor(tensor, position_ids=position_ids, **kwargs)
+        if compressed is None:
+            return None
+        cutoff = (tensor.size(seq_dim) // compress_ratio) * compress_ratio
+        comp_pos = position_ids[:, :cutoff:compress_ratio]
+        return compressed, comp_pos
+
+    drop_prefix = 0
+    tail_parts = None
+    if compressor.overlap:
+        tail_parts = _gather_contiguous_tail(
+            tensor,
+            tail_len=compress_ratio,
+            cp_size=cp_size,
+            cp_group=cp_group,
+            seq_dim=seq_dim,
+        )
+        zero_tail = tensor.new_zeros(())
+        for tail in tail_parts:
+            zero_tail = zero_tail + tail.to(dtype=tensor.dtype).sum() * 0.0
+        tensor = tensor + zero_tail
+    if tail_parts is not None and cp_rank > 0:
+        prefix = tail_parts[cp_rank - 1].to(device=tensor.device, dtype=tensor.dtype)
+        prefix_pos = position_ids[:, :compress_ratio] - compress_ratio
+        tensor = torch.cat([prefix, tensor], dim=seq_dim)
+        position_ids = torch.cat([prefix_pos, position_ids], dim=1)
+        drop_prefix = 1
+
+    compressed = compressor(tensor, position_ids=position_ids, **kwargs)
+    if compressed is None:
+        return None
+    cutoff = (tensor.size(seq_dim) // compress_ratio) * compress_ratio
+    comp_pos = position_ids[:, :cutoff:compress_ratio]
+    if drop_prefix:
+        compressed = compressed.narrow(
+            compressed_seq_dim,
+            drop_prefix,
+            compressed.size(compressed_seq_dim) - drop_prefix,
+        )
+        comp_pos = comp_pos[:, drop_prefix:]
+    if compressed.size(compressed_seq_dim) == 0:
+        return None
+    return compressed, comp_pos
diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/csa.py b/experimental/lite/megatron/lite/primitive/modules/attention/csa.py
new file mode 100644
index 00000000000..4cee8ee7528
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/attention/csa.py
@@ -0,0 +1,714 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import transformer_engine.pytorch as te
+from megatron.lite.primitive.modules.attention.dsa import rotate_activation
+from megatron.lite.primitive.modules.attention.cp import (
+    compress_contiguous_chunks_for_cp,
+    iter_cp_sources,
+)
+from megatron.lite.primitive.parallel.state import ParallelState
+from megatron.lite.primitive.utils.rotary import (
+    _yarn_find_correction_range,
+    _yarn_linear_ramp_mask,
+)
+
+
+class GroupedLinear(nn.Module):
+    def __init__(self, in_features_per_group: int, out_features: int, n_groups: int):
+        super().__init__()
+        self.in_features_per_group = in_features_per_group
+        self.out_features = out_features
+        self.n_groups = n_groups
+        self.weight = nn.Parameter(torch.empty(out_features, in_features_per_group))
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out_per_group = self.out_features // self.n_groups
+        weight = self.weight.view(self.n_groups, out_per_group, self.in_features_per_group)
+        return torch.einsum("...gd,god->...go", x, weight)
+
+
+def build_rope_cos_sin(
+    position_ids: torch.Tensor,
+    rope_head_dim: int,
+    rope_theta: float,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (
+        rope_theta
+        ** (torch.arange(0, rope_head_dim, 2, device=device, dtype=torch.float32) / rope_head_dim)
+    )
+    freqs = torch.einsum("bs,d->bsd", position_ids.to(torch.float32), inv_freq)
+    emb = torch.cat([freqs, freqs], dim=-1)
+    return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)
+
+
+def build_yarn_rope_cos_sin(
+    position_ids: torch.Tensor,
+    rope_head_dim: int,
+    rope_theta: float,
+    *,
+    config: Any,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    dim = rope_head_dim
+    inv_freq_extra = 1.0 / (
+        rope_theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)
+    )
+    inv_freq_inter = 1.0 / (
+        config.rotary_scaling_factor
+        * rope_theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)
+    )
+    low, high = _yarn_find_correction_range(
+        config.beta_fast,
+        config.beta_slow,
+        dim,
+        rope_theta,
+        config.original_max_position_embeddings,
+    )
+    inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, dim // 2, device)
+    inv_freq = inv_freq_inter * (1 - inv_freq_mask) + inv_freq_extra * inv_freq_mask
+    freqs = torch.einsum("bs,d->bsd", position_ids.to(torch.float32), inv_freq)
+    emb = torch.cat([freqs, freqs], dim=-1)
+    return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)
+
+
+def build_compressed_rope_cos_sin(
+    position_ids: torch.Tensor,
+    rope_head_dim: int,
+    rope_theta: float,
+    *,
+    config: Any,
+    use_yarn: bool,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if use_yarn:
+        return build_yarn_rope_cos_sin(
+            position_ids,
+            rope_head_dim,
+            rope_theta,
+            config=config,
+            device=device,
+            dtype=dtype,
+        )
+    return build_rope_cos_sin(position_ids, rope_head_dim, rope_theta, device=device, dtype=dtype)
+
+
+def apply_partial_rope(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, rope_head_dim: int
+) -> torch.Tensor:
+    if rope_head_dim == 0:
+        return x
+    rope = x[..., -rope_head_dim:]
+    tail = x[..., :-rope_head_dim]
+    rope_pairs = rope.unflatten(-1, (-1, 2))
+    a, b = rope_pairs[..., 0], rope_pairs[..., 1]
+    c = cos[..., : rope_head_dim // 2]
+    s = sin[..., : rope_head_dim // 2]
+    while c.ndim < a.ndim:
+        c = c.unsqueeze(1)
+        s = s.unsqueeze(1)
+    out_a = a * c - b * s
+    out_b = a * s + b * c
+    rope_out = torch.stack([out_a, out_b], dim=-1).flatten(-2)
+    return torch.cat([tail, rope_out], dim=-1)
+
+
+class CompressedSequenceCompressor(nn.Module):
+    def __init__(self, config: Any, compress_ratio: int, head_dim: int, *, rotate: bool = False):
+        super().__init__()
+        self.config = config
+        self.compress_ratio = compress_ratio
+        self.head_dim = head_dim
+        self.rope_head_dim = min(config.qk_rope_head_dim, head_dim)
+        self.overlap = compress_ratio == 4
+        self.coff = 2 if self.overlap else 1
+        self.rotate = rotate
+        self.wkv = nn.Linear(config.hidden_size, self.coff * head_dim, bias=False)
+        self.wgate = nn.Linear(config.hidden_size, self.coff * head_dim, bias=False)
+        self.ape = nn.Parameter(
+            torch.empty(compress_ratio, self.coff * head_dim, dtype=torch.float32)
+        )
+        self.norm = te.RMSNorm(head_dim, eps=config.rms_norm_eps)
+        nn.init.normal_(self.ape, mean=0.0, std=config.initializer_range)
+
+    def _overlap_transform(self, tensor: torch.Tensor, fill_value: float) -> torch.Tensor:
+        bsz, n_blocks, ratio, _, head_dim = tensor.shape
+        out = tensor.new_full((bsz, n_blocks, 2 * ratio, head_dim), fill_value)
+        out[:, :, ratio:] = tensor[:, :, :, 1]
+        out[:, 1:, :ratio] = tensor[:, :-1, :, 0]
+        return out
+
+    def forward(
+        self, x: torch.Tensor, *, position_ids: torch.Tensor, rope_theta: float
+    ) -> torch.Tensor | None:
+        bsz, seq_len, _ = x.shape
+        ratio = self.compress_ratio
+        n_blocks = seq_len // ratio
+        if n_blocks == 0:
+            return None
+        cutoff = n_blocks * ratio
+        content = self.wkv(x[:, :cutoff])
+        gate = self.wgate(x[:, :cutoff])
+        content = content.view(bsz, n_blocks, ratio, self.coff, self.head_dim)
+        gate = gate.view_as(content)
+        gate = gate + self.ape.view(1, 1, ratio, self.coff, self.head_dim).to(gate.device)
+        if self.overlap:
+            content = self._overlap_transform(content, 0.0)
+            gate = self._overlap_transform(gate, float("-inf"))
+        else:
+            content = content.squeeze(3)
+            gate = gate.squeeze(3)
+        weights = torch.softmax(gate.float(), dim=2).to(dtype=content.dtype)
+        compressed = self.norm((content * weights).sum(dim=2)).unsqueeze(1)
+        compressed_positions = position_ids[:, :cutoff:ratio]
+        cos, sin = build_compressed_rope_cos_sin(
+            compressed_positions,
+            self.rope_head_dim,
+            rope_theta,
+            config=self.config,
+            use_yarn=self.compress_ratio > 1,
+            device=x.device,
+            dtype=compressed.dtype,
+        )
+        compressed = apply_partial_rope(compressed, cos, sin, self.rope_head_dim)
+        return rotate_activation(compressed) if self.rotate else compressed
+
+
+def _source_scores_mask(
+    q_positions: torch.Tensor, k_positions: torch.Tensor, *, sliding_window: int
+) -> torch.Tensor:
+    q_pos = q_positions.unsqueeze(-1)
+    k_pos = k_positions.unsqueeze(1)
+    return (k_pos <= q_pos) & (k_pos >= q_pos - sliding_window + 1)
+
+
+def _compressed_scores_mask(
+    q_positions: torch.Tensor, comp_positions: torch.Tensor, *, ratio: int
+) -> torch.Tensor:
+    visible = (q_positions + 1) // ratio
+    comp_ids = comp_positions // ratio
+    return comp_ids.unsqueeze(1) < visible.unsqueeze(-1)
+
+
+def _window_topk_indices(
+    batch: int, seq_len: int, window: int, *, device: torch.device
+) -> torch.Tensor:
+    topk = max(1, min(int(window), seq_len))
+    query_pos = torch.arange(seq_len, device=device).view(seq_len, 1)
+    offsets = torch.arange(topk, device=device)
+    indices = (query_pos - topk + 1).clamp(min=0) + offsets
+    indices = torch.where(indices > query_pos, -1, indices)
+    return indices.unsqueeze(0).expand(batch, -1, -1).to(torch.int32)
+
+
+def _load_dsa_kernels():
+    from megatron.lite.primitive.kernels import dsa_kernels
+
+    return dsa_kernels
+
+
+class CompressedSparseAttentionIndexer(nn.Module):
+    def __init__(self, config, compress_ratio: int):
+        super().__init__()
+        self.config = config
+        self.index_n_heads = config.index_n_heads
+        self.index_head_dim = config.index_head_dim
+        self.index_topk = config.index_topk
+        self.rope_head_dim = min(config.qk_rope_head_dim, config.index_head_dim)
+        self.softmax_scale = self.index_head_dim**-0.5
+        self.wq_b = nn.Linear(
+            config.q_lora_rank, config.index_n_heads * config.index_head_dim, bias=False
+        )
+        self.weights_proj = nn.Linear(config.hidden_size, config.index_n_heads, bias=False)
+        self.compressor = CompressedSequenceCompressor(
+            config, compress_ratio, config.index_head_dim, rotate=True
+        )
+
+
+class CompressedSparseAttention(nn.Module):
+    def __init__(
+        self,
+        config,
+        *,
+        layer_idx: int,
+        ps: ParallelState,
+    ):
+        super().__init__()
+        self.config = config
+        self.ps = ps
+        self.attention_backend = "torch"
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.rope_head_dim = config.qk_rope_head_dim
+        self.num_heads_per_group = config.num_attention_heads // config.o_groups
+        # MTP layers use layer_idx == num_hidden_layers (+i), which is past the
+        # per-decoder-layer compress_ratios list (length num_hidden_layers); fall
+        # back to the last real layer's ratio so the MTP CSA still builds.
+        if config.compress_ratios:
+            _cr_idx = min(layer_idx, len(config.compress_ratios) - 1)
+            self.compress_ratio = config.compress_ratios[_cr_idx]
+        else:
+            self.compress_ratio = 0
+        self.wq_a = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False)
+        self.q_norm = te.RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps)
+        self.wq_b = nn.Linear(config.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+        self.wkv = nn.Linear(config.hidden_size, self.head_dim, bias=False)
+        self.kv_norm = te.RMSNorm(config.head_dim, eps=config.rms_norm_eps)
+        self.wo_a = GroupedLinear(
+            self.num_heads_per_group * self.head_dim,
+            config.o_groups * config.o_lora_rank,
+            config.o_groups,
+        )
+        self.wo_b = nn.Linear(config.o_groups * config.o_lora_rank, config.hidden_size, bias=False)
+        self.sinks = nn.Parameter(torch.zeros(self.num_heads))
+        self.compressor = (
+            CompressedSequenceCompressor(config, self.compress_ratio, self.head_dim)
+            if self.compress_ratio > 1
+            else None
+        )
+        self.indexer = (
+            CompressedSparseAttentionIndexer(config, self.compress_ratio)
+            if self.compress_ratio == 4
+            else None
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        *,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.ps.cp_size > 1 and attention_mask is not None:
+            raise ValueError("CP expects attention_mask=None; masks are derived from position_ids.")
+        batch, seq_len, _ = x.shape
+        attention_rope_theta = (
+            self.config.compress_rope_theta if self.compress_ratio > 1 else self.config.rope_theta
+        )
+        cos, sin = build_compressed_rope_cos_sin(
+            position_ids,
+            self.rope_head_dim,
+            attention_rope_theta,
+            config=self.config,
+            use_yarn=self.compress_ratio > 1,
+            device=x.device,
+            dtype=x.dtype,
+        )
+        q_low = self.q_norm(self.wq_a(x))
+        q = self.wq_b(q_low).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        q = q * torch.rsqrt(
+            q.float().pow(2).mean(dim=-1, keepdim=True) + self.config.rms_norm_eps
+        ).to(dtype=q.dtype)
+        kv = self.kv_norm(self.wkv(x)).view(batch, seq_len, 1, self.head_dim).transpose(1, 2)
+        q = apply_partial_rope(q, cos, sin, self.rope_head_dim)
+        kv = apply_partial_rope(kv, cos, sin, self.rope_head_dim)
+        use_sparse_backend = self.attention_backend not in {"local", "eager", "torch"}
+        if (
+            use_sparse_backend
+            and self.compress_ratio == 4
+            and self.ps.cp_size == 1
+            and self.compressor is not None
+            and self.indexer is not None
+        ):
+            return self._forward_fused_dsa_cp1(
+                x,
+                q,
+                q_low,
+                kv,
+                position_ids=position_ids,
+                cos=cos,
+                sin=sin,
+                attention_mask=attention_mask,
+            )
+        if use_sparse_backend and self.ps.cp_size == 1 and attention_mask is None:
+            return self._forward_fused_sparse_no_indexer_cp1(
+                x,
+                q,
+                kv,
+                position_ids=position_ids,
+                cos=cos,
+                sin=sin,
+            )
+
+        dense_score_parts = []
+        dense_value_parts = []
+        for _source_rank, source_kv, source_pos in iter_cp_sources(
+            kv,
+            position_ids,
+            cp_rank=self.ps.cp_rank,
+            cp_size=self.ps.cp_size,
+            cp_group=self.ps.cp_group,
+        ):
+            source_heads = source_kv.expand(-1, self.num_heads, -1, -1)
+            scores = torch.matmul(q.float(), source_heads.float().transpose(-1, -2)) / (
+                self.head_dim**0.5
+            )
+            source_mask = _source_scores_mask(
+                position_ids,
+                source_pos,
+                sliding_window=self.config.sliding_window,
+            ).unsqueeze(1)
+            scores = scores.masked_fill(~source_mask, -float("inf"))
+            dense_score_parts.append(scores)
+            dense_value_parts.append(source_heads)
+        dense_scores = torch.cat(dense_score_parts, dim=-1)
+        if attention_mask is not None:
+            dense_scores = dense_scores + attention_mask.to(dtype=dense_scores.dtype)
+        score_parts = [dense_scores]
+        value_parts = [torch.cat(dense_value_parts, dim=2)]
+
+        if self.compressor is not None:
+            compressed_pack = compress_contiguous_chunks_for_cp(
+                self.compressor,
+                x,
+                position_ids=position_ids,
+                cp_rank=self.ps.cp_rank,
+                cp_size=self.ps.cp_size,
+                cp_group=self.ps.cp_group,
+                compress_kwargs={"rope_theta": self.config.compress_rope_theta},
+            )
+        else:
+            compressed_pack = None
+        if compressed_pack is not None:
+            compressed, compressed_pos = compressed_pack
+            compressed, compressed_pos = self._gather_cp_sources(
+                compressed, compressed_pos, seq_dim=2
+            )
+            compressed_values = compressed.expand(-1, self.num_heads, -1, -1)
+            compressed_scores = torch.matmul(
+                q.float(), compressed_values.float().transpose(-1, -2)
+            ) / (self.head_dim**0.5)
+            compressed_valid = _compressed_scores_mask(
+                position_ids,
+                compressed_pos,
+                ratio=self.compress_ratio,
+            ).unsqueeze(1)
+            compressed_scores = compressed_scores.masked_fill(~compressed_valid, -float("inf"))
+
+            if self.indexer is not None:
+                index_comp_pack = compress_contiguous_chunks_for_cp(
+                    self.indexer.compressor,
+                    x,
+                    position_ids=position_ids,
+                    cp_rank=self.ps.cp_rank,
+                    cp_size=self.ps.cp_size,
+                    cp_group=self.ps.cp_group,
+                    compress_kwargs={"rope_theta": self.config.compress_rope_theta},
+                )
+                if index_comp_pack is not None:
+                    index_comp, index_pos = index_comp_pack
+                    index_comp, index_pos = self._gather_cp_sources(
+                        index_comp, index_pos, seq_dim=2
+                    )
+                    idx_cos, idx_sin = build_compressed_rope_cos_sin(
+                        position_ids,
+                        self.indexer.rope_head_dim,
+                        self.config.compress_rope_theta,
+                        config=self.config,
+                        use_yarn=self.compress_ratio > 1,
+                        device=x.device,
+                        dtype=x.dtype,
+                    )
+                    q_idx = (
+                        self.indexer.wq_b(q_low)
+                        .view(
+                            batch, seq_len, self.indexer.index_n_heads, self.indexer.index_head_dim
+                        )
+                        .transpose(1, 2)
+                    )
+                    q_idx = apply_partial_rope(q_idx, idx_cos, idx_sin, self.indexer.rope_head_dim)
+                    index_weights = (
+                        self.indexer.weights_proj(x).float()
+                        * (self.indexer.index_n_heads**-0.5)
+                        * self.indexer.softmax_scale
+                    )
+                    k_idx = index_comp.squeeze(1)
+                    index_scores = torch.einsum(
+                        "bhsd,btd->bsht", q_idx.float(), k_idx.float()
+                    ).relu()
+                    index_scores = (index_scores * index_weights.unsqueeze(-1)).sum(dim=2)
+                    index_valid = _compressed_scores_mask(
+                        position_ids,
+                        index_pos,
+                        ratio=self.compress_ratio,
+                    )
+                    index_scores = index_scores.masked_fill(~index_valid, -float("inf"))
+                    topk = min(self.indexer.index_topk, index_scores.size(-1))
+                    topk_indices = index_scores.topk(topk, dim=-1).indices
+                    topk_mask = torch.zeros_like(index_scores, dtype=torch.bool)
+                    topk_mask.scatter_(-1, topk_indices, True)
+                    compressed_scores = compressed_scores + index_scores.unsqueeze(1)
+                    compressed_scores = compressed_scores.masked_fill(
+                        ~topk_mask.unsqueeze(1), -float("inf")
+                    )
+            score_parts.append(compressed_scores)
+            value_parts.append(compressed_values)
+
+        scores = torch.cat(score_parts, dim=-1)
+        sink = (
+            self.sinks.view(1, self.num_heads, 1, 1)
+            .expand(batch, -1, seq_len, -1)
+            .to(dtype=scores.dtype)
+        )
+        combined_scores = torch.cat([scores, sink], dim=-1)
+        combined_scores = combined_scores - combined_scores.max(dim=-1, keepdim=True).values
+        probs = torch.softmax(combined_scores, dim=-1, dtype=torch.float32).to(dtype=q.dtype)
+
+        context = q.new_zeros(batch, self.num_heads, seq_len, self.head_dim)
+        offset = 0
+        for values in value_parts:
+            next_offset = offset + values.size(2)
+            partial = torch.matmul(probs[..., offset:next_offset], values)
+            context = context + partial
+            offset = next_offset
+
+        return self._project_context(context, cos, sin)
+
+    def _project_context(
+        self, context: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        context = apply_partial_rope(context, cos, -sin, self.rope_head_dim).transpose(1, 2)
+        batch, seq_len = context.shape[:2]
+        grouped = context.reshape(
+            batch, seq_len, self.config.o_groups, self.num_heads_per_group * self.head_dim
+        )
+        return self.wo_b(self.wo_a(grouped).flatten(2))
+
+    def _gather_cp_sources(
+        self, tensor: torch.Tensor, position_ids: torch.Tensor, *, seq_dim: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        parts, pos_parts = [], []
+        for _rank, source, source_pos in iter_cp_sources(
+            tensor,
+            position_ids,
+            cp_rank=self.ps.cp_rank,
+            cp_size=self.ps.cp_size,
+            cp_group=self.ps.cp_group,
+        ):
+            parts.append(source)
+            pos_parts.append(source_pos)
+        pos = torch.cat(pos_parts, dim=1)
+        return torch.cat(parts, dim=seq_dim), pos
+
+    def _forward_fused_sparse_no_indexer_cp1(
+        self,
+        x: torch.Tensor,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        *,
+        position_ids: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        dsa_kernels = _load_dsa_kernels()
+        batch, seq_len, _ = x.shape
+        query = q.transpose(1, 2).transpose(0, 1).contiguous()
+        kv_full = kv.squeeze(1)
+        kv_full = kv_full.transpose(0, 1).contiguous()
+        window_idxs = _window_topk_indices(
+            batch,
+            seq_len,
+            self.config.sliding_window,
+            device=x.device,
+        )
+
+        compressed = None
+        if self.compressor is not None and self.compress_ratio > 1:
+            compressed = self.compressor(
+                x,
+                position_ids=position_ids,
+                rope_theta=self.config.compress_rope_theta,
+            )
+            if compressed is not None:
+                compressed_kv = compressed.squeeze(1)
+                kv_full = torch.cat([kv_full, compressed_kv.transpose(0, 1).contiguous()], dim=0)
+
+        if compressed is not None:
+            n_compressed = compressed.size(2)
+            comp_idx = torch.arange(n_compressed, device=x.device).view(1, n_compressed)
+            valid_per_pos = (
+                torch.arange(1, seq_len + 1, device=x.device) // self.compress_ratio
+            ).view(seq_len, 1)
+            compress_topk_idxs = torch.where(
+                comp_idx < valid_per_pos,
+                comp_idx + seq_len,
+                torch.full_like(comp_idx, -1),
+            )
+            compress_topk_idxs = (
+                compress_topk_idxs.unsqueeze(0).expand(batch, -1, -1).to(torch.int32)
+            )
+            flat_idxs, _flat_tlen = dsa_kernels.build_flat_topk_idxs(
+                window_idxs,
+                compress_topk_idxs,
+                batch_size=batch,
+                seqlen_kv=kv_full.size(0),
+            )
+        else:
+            flat_idxs, _flat_tlen = dsa_kernels.build_flat_topk_idxs(
+                window_idxs,
+                batch_size=batch,
+                seqlen_kv=kv_full.size(0),
+            )
+
+        out = dsa_kernels.dsa_sparse_attn(
+            query,
+            kv_full,
+            self.sinks.float(),
+            flat_idxs,
+            self.head_dim**-0.5,
+        )
+        context = (
+            out.view(seq_len, batch, self.num_heads, self.head_dim).permute(1, 2, 0, 3).contiguous()
+        )
+        return self._project_context(context, cos, sin)
+
+    def _forward_fused_dsa_cp1(
+        self,
+        x: torch.Tensor,
+        q: torch.Tensor,
+        q_low: torch.Tensor,
+        kv: torch.Tensor,
+        *,
+        position_ids: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if self.ps.cp_size != 1:
+            raise NotImplementedError("DeepSeek V4 fused DSA path currently supports CP=1 only.")
+        if attention_mask is not None:
+            raise NotImplementedError(
+                "DeepSeek V4 fused DSA path currently supports causal masking only."
+            )
+        dsa_kernels = _load_dsa_kernels()
+        # The cuDNN SM90 indexer requires seqlen_q <= seqlen_k * ratio, but the
+        # compressor floors to seq_len // ratio blocks, so a seq_len that is not a
+        # multiple of ratio leaves the last query token(s) without a compressed key
+        # block. Right-pad to a multiple of ratio (the causal tail attends only real
+        # tokens and is sliced off the output) so seqlen_k * ratio == seqlen_q.
+        orig_seq_len = x.shape[1]
+        ratio = self.compress_ratio
+        pad = (-orig_seq_len) % ratio
+        if pad:
+            pos_tail = position_ids[:, -1:] + torch.arange(
+                1, pad + 1, device=position_ids.device, dtype=position_ids.dtype
+            )
+            position_ids = torch.cat([position_ids, pos_tail], dim=1)
+            x = torch.nn.functional.pad(x, (0, 0, 0, pad))
+            q = torch.nn.functional.pad(q, (0, 0, 0, pad))
+            q_low = torch.nn.functional.pad(q_low, (0, 0, 0, pad))
+            kv = torch.nn.functional.pad(kv, (0, 0, 0, pad))
+        batch, seq_len, _ = x.shape
+        compressed = self.compressor(
+            x,
+            position_ids=position_ids,
+            rope_theta=self.config.compress_rope_theta,
+        )
+        index_comp = self.indexer.compressor(
+            x,
+            position_ids=position_ids,
+            rope_theta=self.config.compress_rope_theta,
+        )
+        if compressed is None or index_comp is None:
+            raise RuntimeError("DeepSeek V4 fused DSA requires at least one compressed KV entry.")
+        compressed_kv = compressed.squeeze(1)
+        index_k = index_comp.squeeze(1).transpose(0, 1).contiguous()
+        kv_full = torch.cat([kv.squeeze(1), compressed_kv], dim=1)
+        kv_full = kv_full.transpose(0, 1).contiguous()
+
+        idx_cos, idx_sin = build_compressed_rope_cos_sin(
+            position_ids,
+            self.indexer.rope_head_dim,
+            self.config.compress_rope_theta,
+            config=self.config,
+            use_yarn=self.compress_ratio > 1,
+            device=x.device,
+            dtype=x.dtype,
+        )
+        q_indexer = self.indexer.wq_b(q_low).view(
+            batch, seq_len, self.indexer.index_n_heads, self.indexer.index_head_dim
+        )
+        q_indexer = q_indexer.transpose(1, 2)
+        q_indexer = apply_partial_rope(q_indexer, idx_cos, idx_sin, self.indexer.rope_head_dim)
+        q_indexer = rotate_activation(q_indexer)
+        q_indexer = q_indexer.transpose(1, 2).transpose(0, 1).contiguous()
+        weights_indexer = (
+            (self.indexer.weights_proj(x).to(dtype=x.dtype) * (self.indexer.index_n_heads**-0.5))
+            .transpose(0, 1)
+            .contiguous()
+        )
+        indexer_topk = int(self.indexer.index_topk)
+        if indexer_topk <= 0:
+            raise RuntimeError("DeepSeek V4 fused DSA requires positive indexer_topk.")
+        window_idxs = _window_topk_indices(
+            batch,
+            seq_len,
+            self.config.sliding_window,
+            device=x.device,
+        )
+        query = q.transpose(1, 2).transpose(0, 1).contiguous()
+        sink = self.sinks.float()
+
+        if self.training and torch.is_grad_enabled():
+            out, _indexer_loss = dsa_kernels.fused_indexer_sparse_attn(
+                query,
+                kv_full,
+                sink,
+                window_idxs,
+                q_indexer,
+                index_k,
+                weights_indexer,
+                indexer_topk,
+                self.compress_ratio,
+                self.head_dim**-0.5,
+                self.indexer.softmax_scale,
+                0.0,
+                sparse_loss=False,
+                kv_offset=seq_len,
+                calculate_per_token_loss=False,
+            )
+        else:
+            topk_indices, _topk_length = dsa_kernels.indexer_topk(
+                q_indexer,
+                index_k,
+                weights_indexer,
+                indexer_topk,
+                self.compress_ratio,
+                indexer_softmax_scale=self.indexer.softmax_scale,
+            )
+            topk_indices = torch.where(
+                topk_indices >= 0,
+                topk_indices + seq_len,
+                topk_indices,
+            ).to(torch.int32)
+            flat_idxs, flat_tlen = dsa_kernels.build_flat_topk_idxs(
+                window_idxs,
+                topk_indices,
+                batch_size=batch,
+                seqlen_kv=kv_full.size(0),
+                compact=True,
+            )
+            out = dsa_kernels.dsa_sparse_attn(
+                query,
+                kv_full,
+                sink,
+                flat_idxs,
+                self.head_dim**-0.5,
+                topk_length=flat_tlen,
+            )
+
+        context = (
+            out.view(seq_len, batch, self.num_heads, self.head_dim).permute(1, 2, 0, 3).contiguous()
+        )
+        if pad:
+            context = context[:, :, :orig_seq_len, :].contiguous()
+        return self._project_context(context, cos, sin)
diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/dsa.py b/experimental/lite/megatron/lite/primitive/modules/attention/dsa.py
new file mode 100644
index 00000000000..8956d4a9f01
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/attention/dsa.py
@@ -0,0 +1,697 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Dynamic Sparse Attention.
+
+The module is model-agnostic: callers pass architecture dimensions directly and
+keep model config classes out of the primitive layer.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import transformer_engine.pytorch as te
+
+from megatron.lite.primitive.parallel.cp import (
+    zigzag_reconstruct_from_cp_parts,
+    zigzag_slice_for_cp,
+)
+from megatron.lite.primitive.parallel.thd import (
+    reconstruct_packed_from_cp_parts,
+    split_packed_to_cp_local,
+)
+
+from megatron.lite.primitive.kernels import dsa_kernels as _dsa_kernels
+
+if TYPE_CHECKING:
+    from megatron.lite.primitive.modules.attention.mla import MultiLatentAttention
+
+
+def _fused_indexer_sparse_attn(*args, value_dim: int | None = None, **kwargs):
+    try:
+        return _dsa_kernels.fused_indexer_sparse_attn(*args, value_dim=value_dim, **kwargs)
+    except TypeError as exc:
+        if "value_dim" not in str(exc):
+            raise
+        return _dsa_kernels.fused_indexer_sparse_attn(*args, **kwargs)
+
+
+class DSAIndexerLossAutoScaler(torch.autograd.Function):
+    """Attach the DSA indexer loss to the output without changing forward values."""
+
+    main_loss_backward_scale: torch.Tensor | None = None
+
+    @staticmethod
+    def forward(ctx, output: torch.Tensor, indexer_loss: torch.Tensor) -> torch.Tensor:
+        ctx.save_for_backward(indexer_loss)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        (indexer_loss,) = ctx.saved_tensors
+        if DSAIndexerLossAutoScaler.main_loss_backward_scale is None:
+            DSAIndexerLossAutoScaler.main_loss_backward_scale = torch.tensor(
+                1.0, device=indexer_loss.device
+            )
+        indexer_loss_backward_scale = DSAIndexerLossAutoScaler.main_loss_backward_scale
+        scaled_indexer_loss_grad = torch.ones_like(indexer_loss) * indexer_loss_backward_scale
+        return grad_output, scaled_indexer_loss_grad
+
+    @staticmethod
+    def set_loss_scale(scale: torch.Tensor) -> None:
+        if DSAIndexerLossAutoScaler.main_loss_backward_scale is None:
+            DSAIndexerLossAutoScaler.main_loss_backward_scale = scale
+        else:
+            DSAIndexerLossAutoScaler.main_loss_backward_scale.copy_(scale)
+
+
+RMSNorm = te.RMSNorm
+
+
+def _hadamard_transform_torch(x: torch.Tensor, scale: float) -> torch.Tensor:
+    n = x.shape[-1]
+    if n <= 0 or n & (n - 1):
+        raise ValueError(f"Hadamard rotation requires power-of-two dim, got {n}")
+    original_shape = x.shape
+    y = x.reshape(-1, n)
+    h = 1
+    while h < n:
+        y = y.reshape(-1, n // (h * 2), h * 2)
+        left = y[..., :h]
+        right = y[..., h:]
+        y = torch.cat([left + right, left - right], dim=-1)
+        h *= 2
+    return y.reshape(original_shape) * scale
+
+
+try:
+    from fast_hadamard_transform import hadamard_transform as _fast_hadamard_transform
+except Exception:  # pragma: no cover - optional CUDA extension
+    _fast_hadamard_transform = None
+
+
+def rotate_activation(x: torch.Tensor) -> torch.Tensor:
+    x = x.to(torch.bfloat16) if x.dtype != torch.bfloat16 else x
+    scale = x.shape[-1] ** -0.5
+    if _fast_hadamard_transform is not None and x.is_cuda:
+        return _fast_hadamard_transform(x, scale=scale)
+    return _hadamard_transform_torch(x, scale=scale)
+
+
+def build_rope_cache(
+    *, dim: int, max_position_embeddings: int, rope_theta: float, device: torch.device | None = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (
+        rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+    )
+    positions = torch.arange(max_position_embeddings, dtype=torch.float32, device=device)
+    freqs = torch.outer(positions, inv_freq)
+    return freqs.cos(), freqs.sin()
+
+
+def build_rotary_embeddings(
+    *, position_ids: torch.Tensor, dim: int, rope_theta: float, dtype: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    device = position_ids.device
+    inv_freq = 1.0 / (
+        rope_theta
+        ** (torch.arange(0, dim, 2, dtype=torch.int64, device=device).to(torch.float32) / dim)
+    )
+    inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+    position_ids_expanded = position_ids[:, None, :].float()
+    device_type = device.type if isinstance(device.type, str) and device.type != "mps" else "cpu"
+    with torch.autocast(device_type=device_type, enabled=False):
+        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos()
+        sin = emb.sin()
+    return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    x1 = x[..., :half]
+    x2 = x[..., half:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, *, unsqueeze_dim: int
+) -> torch.Tensor:
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+def apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    position_ids: torch.Tensor,
+    *,
+    interleaved: bool = True,
+) -> torch.Tensor:
+    if position_ids.dim() == 3:
+        position_ids = position_ids[0]
+    if position_ids.dim() == 1:
+        position_ids = position_ids.unsqueeze(0)
+
+    input_dtype = x.dtype
+    x = x.float()
+    cos = cos.to(device=x.device)[position_ids].float().unsqueeze(2)
+    sin = sin.to(device=x.device)[position_ids].float().unsqueeze(2)
+    if interleaved:
+        x_even = x[..., 0::2]
+        x_odd = x[..., 1::2]
+        out = torch.empty_like(x)
+        out[..., 0::2] = x_even * cos - x_odd * sin
+        out[..., 1::2] = x_even * sin + x_odd * cos
+        return out.to(input_dtype)
+
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).to(input_dtype)
+
+
+def _rotary_embeddings_from_cache(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    position_ids: torch.Tensor,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if cos.dim() == 3 and sin.dim() == 3:
+        return cos.to(device=device, dtype=dtype), sin.to(device=device, dtype=dtype)
+
+    if position_ids.dim() == 3:
+        position_ids = position_ids[0]
+    if position_ids.dim() == 1:
+        position_ids = position_ids.unsqueeze(0)
+    cos = cos.to(device=device)[position_ids].float()
+    sin = sin.to(device=device)[position_ids].float()
+    if cos.shape[-1] * 2 == dim:
+        cos = torch.cat((cos, cos), dim=-1)
+        sin = torch.cat((sin, sin), dim=-1)
+    return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+def _all_gather_cp(tensor: torch.Tensor, *, cp_size: int, cp_group) -> list[torch.Tensor]:
+    if cp_size <= 1:
+        return [tensor]
+    if cp_group is None:
+        raise RuntimeError("CP>1 requires a context-parallel process group.")
+    from torch.distributed.nn.functional import all_gather
+
+    return list(all_gather(tensor.contiguous(), group=cp_group))
+
+
+class DSAIndexer(nn.Module):
+    """Compute per-token top-k key indices for Dynamic Sparse Attention."""
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        q_lora_rank: int,
+        qk_rope_head_dim: int,
+        index_n_heads: int,
+        index_head_dim: int,
+        index_topk: int,
+        rope_interleaved: bool = True,
+        layer_norm_eps: float = 1e-5,
+        rope_first: bool = False,
+        use_hadamard: bool = True,
+    ):
+        super().__init__()
+        if index_head_dim < qk_rope_head_dim:
+            raise ValueError("index_head_dim must be >= qk_rope_head_dim")
+        self.num_heads = index_n_heads
+        self.head_dim = index_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_nope_head_dim = index_head_dim - qk_rope_head_dim
+        self.index_topk = index_topk
+        self.rope_interleaved = rope_interleaved
+        self.rope_first = rope_first
+        self.use_hadamard = use_hadamard
+
+        self.wq_b = nn.Linear(q_lora_rank, index_n_heads * index_head_dim, bias=False)
+        self.wk = nn.Linear(hidden_size, index_head_dim, bias=False)
+        self.k_norm = nn.LayerNorm(index_head_dim, eps=layer_norm_eps)
+        self.weights_proj = nn.Linear(hidden_size, index_n_heads, bias=False)
+        self.softmax_scale = index_head_dim**-0.5
+
+    def forward_before_topk(
+        self,
+        x: torch.Tensor,
+        q_resid: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Project GLM5 indexer inputs for Megatron's fused DSA kernels."""
+        if attention_mask is not None:
+            raise NotImplementedError(
+                "GLM5 fused DSA indexer only supports causal masking; custom "
+                "attention_mask is not supported."
+            )
+        batch, seq_len, _ = x.shape
+        cos, sin = _rotary_embeddings_from_cache(
+            cos, sin, position_ids, device=x.device, dtype=x.dtype, dim=self.qk_rope_head_dim
+        )
+
+        q = self.wq_b(q_resid).view(batch, seq_len, self.num_heads, self.head_dim)
+
+        k = self.k_norm(self.wk(x))
+        if self.rope_first:
+            q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.qk_nope_head_dim], dim=-1)
+            k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.qk_nope_head_dim], dim=-1)
+        else:
+            q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+            k_nope, k_pe = torch.split(k, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=2)
+        k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, unsqueeze_dim=2)
+        k_pe = k_pe.squeeze(2)
+
+        if self.rope_first:
+            q = torch.cat([q_pe, q_nope], dim=-1)
+            k = torch.cat([k_pe, k_nope], dim=-1)
+        else:
+            q = torch.cat([q_nope, q_pe], dim=-1)
+            k = torch.cat([k_nope, k_pe], dim=-1)
+        if self.use_hadamard:
+            q = rotate_activation(q)
+            k = rotate_activation(k)
+
+        weights = self.weights_proj(x) * (self.num_heads**-0.5)
+        return (
+            q.transpose(0, 1).contiguous(),
+            k.transpose(0, 1).contiguous(),
+            weights.transpose(0, 1).contiguous(),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        q_resid: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        q, k, weights = self.forward_before_topk(
+            x, q_resid, cos, sin, position_ids, attention_mask=attention_mask
+        )
+        topk_indices, _ = _dsa_kernels.indexer_topk(
+            q,
+            k,
+            weights,
+            min(self.index_topk, k.shape[0]),
+            1,
+            indexer_softmax_scale=self.softmax_scale,
+        )
+        return topk_indices
+
+
+class DynamicSparseAttention(nn.Module):
+    """Correctness-first DSA attention path."""
+
+    @staticmethod
+    def dense_attention_cls() -> type[MultiLatentAttention]:
+        from megatron.lite.primitive.modules.attention.mla import MultiLatentAttention
+
+        return MultiLatentAttention
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_attention_heads: int,
+        q_lora_rank: int,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        index_n_heads: int,
+        index_head_dim: int,
+        index_topk: int,
+        rms_norm_eps: float,
+        rope_interleaved: bool = True,
+        latent_rms_norm_eps: float | None = None,
+        indexer_layer_norm_eps: float = 1e-5,
+        indexer_rope_interleaved: bool | None = None,
+        indexer_rope_first: bool = False,
+        indexer_use_hadamard: bool = True,
+        indexer_loss_coeff: float = 0.0,
+        indexer_use_sparse_loss: bool = False,
+        calculate_per_token_loss: bool = False,
+        cp_size: int = 1,
+        cp_rank: int = 0,
+        cp_group=None,
+    ):
+        super().__init__()
+        if cp_size < 1:
+            raise ValueError(f"cp_size must be >= 1, got {cp_size}")
+        if not 0 <= cp_rank < cp_size:
+            raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}")
+        self.num_heads = num_attention_heads
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.rope_interleaved = rope_interleaved
+        self.softmax_scale = self.qk_head_dim**-0.5
+        self.indexer_loss_coeff = indexer_loss_coeff
+        self.indexer_use_sparse_loss = indexer_use_sparse_loss
+        self.calculate_per_token_loss = calculate_per_token_loss
+        self.cp_size = cp_size
+        self.cp_rank = cp_rank
+        self.cp_group = cp_group
+        latent_rms_norm_eps = rms_norm_eps if latent_rms_norm_eps is None else latent_rms_norm_eps
+        indexer_rope_interleaved = (
+            rope_interleaved if indexer_rope_interleaved is None else indexer_rope_interleaved
+        )
+
+        self.q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False)
+        self.q_a_layernorm = RMSNorm(q_lora_rank, eps=latent_rms_norm_eps)
+        self.q_b_proj = nn.Linear(q_lora_rank, num_attention_heads * self.qk_head_dim, bias=False)
+        self.kv_a_proj_with_mqa = nn.Linear(
+            hidden_size, kv_lora_rank + qk_rope_head_dim, bias=False
+        )
+        self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=latent_rms_norm_eps)
+        self.kv_b_proj = nn.Linear(
+            kv_lora_rank, num_attention_heads * (qk_nope_head_dim + v_head_dim), bias=False
+        )
+        self.o_proj = nn.Linear(num_attention_heads * v_head_dim, hidden_size, bias=False)
+        self.indexer = DSAIndexer(
+            hidden_size=hidden_size,
+            q_lora_rank=q_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            index_n_heads=index_n_heads,
+            index_head_dim=index_head_dim,
+            index_topk=index_topk,
+            rope_interleaved=indexer_rope_interleaved,
+            layer_norm_eps=indexer_layer_norm_eps,
+            rope_first=indexer_rope_first,
+            use_hadamard=indexer_use_hadamard,
+        )
+        self.register_buffer(
+            "attn_sink",
+            torch.full((num_attention_heads,), -1.0e20, dtype=torch.float32),
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        *,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        packed_seq_params=None,
+    ) -> torch.Tensor:
+        if attention_mask is not None:
+            raise NotImplementedError(
+                "GLM5 fused DSA only supports causal masking; custom attention_mask "
+                "is not supported."
+            )
+        if packed_seq_params is not None:
+            if self.cp_size > 1:
+                x, position_ids = self._gather_packed_cp_inputs(x, position_ids, packed_seq_params)
+                cos, sin = self._gather_packed_cp_rotary(cos, sin, packed_seq_params, x.device)
+            out = self._forward_packed_full(x, cos, sin, position_ids, packed_seq_params)
+            if self.cp_size > 1:
+                out = split_packed_to_cp_local(
+                    out,
+                    cu_seqlens_padded=self._packed_cu_seqlens(packed_seq_params, x.device),
+                    cp_size=self.cp_size,
+                    cp_rank=self.cp_rank,
+                    dim=1,
+                )
+            return out
+
+        cp_restore = self.cp_size > 1
+        if cp_restore:
+            x, position_ids, attention_mask = self._gather_cp_inputs(
+                x, position_ids, attention_mask
+            )
+
+        out = self._forward_dense_full(x, cos, sin, position_ids)
+        if cp_restore:
+            out = zigzag_slice_for_cp(out, self.cp_rank, self.cp_size, seq_dim=1)
+        return out
+
+    def _forward_packed_full(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        position_ids: torch.Tensor,
+        packed_seq_params,
+    ) -> torch.Tensor:
+        cu_seqlens = self._packed_cu_seqlens(packed_seq_params, x.device)
+        if position_ids.dim() == 1:
+            position_ids = position_ids.unsqueeze(0)
+        if position_ids.shape[-1] != x.shape[1]:
+            raise ValueError(
+                "GLM5 packed DynamicSparseAttention position_ids must cover the reconstructed packed tokens, "
+                f"got {tuple(position_ids.shape)} for packed length {x.shape[1]}."
+            )
+        pieces = []
+        for idx in range(int(cu_seqlens.numel()) - 1):
+            start = int(cu_seqlens[idx].item())
+            end = int(cu_seqlens[idx + 1].item())
+            if end <= start:
+                continue
+            seg_cos, seg_sin = self._slice_rotary_cache(cos, sin, start, end)
+            pieces.append(
+                self._forward_dense_full(
+                    x[:, start:end, :],
+                    seg_cos,
+                    seg_sin,
+                    position_ids[:, start:end],
+                )
+            )
+        if pieces:
+            return torch.cat(pieces, dim=1)
+        return x.new_empty(x.shape[0], 0, self.o_proj.out_features)
+
+    def _forward_dense_full(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+
+        batch, seq_len, _ = x.shape
+        q_resid = self.q_a_layernorm(self.q_a_proj(x))
+        q = self.q_b_proj(q_resid).view(batch, seq_len, self.num_heads, self.qk_head_dim)
+        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        cos, sin = _rotary_embeddings_from_cache(
+            cos, sin, position_ids, device=x.device, dtype=x.dtype, dim=self.qk_rope_head_dim
+        )
+
+        q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=2)
+        k_up_weight, v_up_weight = self._split_kv_b_weights()
+        q_nope = torch.einsum("bshd,hdr->bshr", q_nope, k_up_weight)
+        query_states = torch.cat([q_nope, q_pe], dim=-1).transpose(0, 1).contiguous()
+
+        kv_latent, k_pe = torch.split(
+            self.kv_a_proj_with_mqa(x), [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+        )
+        kv_latent = self.kv_a_layernorm(kv_latent)
+        k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, unsqueeze_dim=2).squeeze(2)
+        kv_full = torch.cat([kv_latent, k_pe], dim=-1).transpose(0, 1).contiguous()
+
+        q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk(
+            x.detach(), q_resid.detach(), cos, sin, position_ids
+        )
+        effective_indexer_topk = min(self.indexer.index_topk, seq_len)
+
+        if self.training and torch.is_grad_enabled():
+            window_idxs = torch.empty(batch, seq_len, 0, device=x.device, dtype=torch.int32)
+            out, indexer_loss = _fused_indexer_sparse_attn(
+                query_states,
+                kv_full,
+                self.attn_sink.float(),
+                window_idxs,
+                q_indexer,
+                k_indexer,
+                weights_indexer,
+                self.indexer.index_topk,
+                1,
+                self.softmax_scale,
+                self.indexer.softmax_scale,
+                self.indexer_loss_coeff,
+                sparse_loss=self.indexer_use_sparse_loss,
+                kv_offset=0,
+                calculate_per_token_loss=self.calculate_per_token_loss,
+                value_dim=self.kv_lora_rank,
+            )
+            if self.indexer_loss_coeff > 0:
+                out = DSAIndexerLossAutoScaler.apply(out, indexer_loss)
+        else:
+            topk_indices, _ = _dsa_kernels.indexer_topk(
+                q_indexer,
+                k_indexer,
+                weights_indexer,
+                effective_indexer_topk,
+                1,
+                indexer_softmax_scale=self.indexer.softmax_scale,
+            )
+            flat_idxs, flat_tlen = _dsa_kernels.build_flat_topk_idxs(
+                topk_indices, batch_size=batch, seqlen_kv=seq_len, compact=True
+            )
+            out = _dsa_kernels.dsa_sparse_attn(
+                query_states,
+                kv_full,
+                self.attn_sink.float(),
+                flat_idxs,
+                self.softmax_scale,
+                topk_length=flat_tlen,
+                value_dim=self.kv_lora_rank,
+            )
+
+        out = out.view(seq_len, batch, self.num_heads, self.kv_lora_rank)
+        out = out.permute(1, 0, 2, 3).contiguous()
+        out = torch.einsum("bshr,hvr->bshv", out, v_up_weight)
+        out = out.reshape(batch, seq_len, self.num_heads * self.v_head_dim)
+        return self.o_proj(out)
+
+    def _split_kv_b_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
+        kv_b = self.kv_b_proj.weight.view(
+            self.num_heads, self.qk_nope_head_dim + self.v_head_dim, self.kv_lora_rank
+        )
+        return (kv_b[:, : self.qk_nope_head_dim, :], kv_b[:, self.qk_nope_head_dim :, :])
+
+    def _gather_cp_inputs(
+        self, x: torch.Tensor, position_ids: torch.Tensor, attention_mask: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        local_batch, local_seq = x.shape[:2]
+        x_parts = _all_gather_cp(x, cp_size=self.cp_size, cp_group=self.cp_group)
+        full_x = zigzag_reconstruct_from_cp_parts(x_parts, seq_dim=1)
+        full_seq = full_x.shape[1]
+
+        full_position_ids = self._full_cp_position_ids(
+            position_ids, batch=local_batch, local_seq=local_seq, full_seq=full_seq, device=x.device
+        )
+        if attention_mask is not None:
+            expected = (full_seq, full_seq)
+            if tuple(attention_mask.shape[-2:]) != expected:
+                raise NotImplementedError(
+                    "GLM5 DynamicSparseAttention CP attention_mask must already cover the reconstructed "
+                    f"full sequence {expected}, got {tuple(attention_mask.shape)}."
+                )
+        return full_x, full_position_ids, attention_mask
+
+    def _full_cp_position_ids(
+        self,
+        position_ids: torch.Tensor,
+        *,
+        batch: int,
+        local_seq: int,
+        full_seq: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if position_ids.dim() == 3:
+            position_ids = position_ids[0]
+        if position_ids.dim() == 1:
+            position_ids = position_ids.unsqueeze(0).expand(batch, -1)
+        if position_ids.shape[-1] == full_seq:
+            return position_ids.to(device=device, dtype=torch.long)
+        if position_ids.shape[-1] != local_seq:
+            raise ValueError(
+                "GLM5 DynamicSparseAttention CP position_ids must be either local or full sequence length, "
+                f"got {tuple(position_ids.shape)} for local_seq={local_seq}, full_seq={full_seq}."
+            )
+
+        pos_parts = _all_gather_cp(
+            position_ids.to(device=device, dtype=torch.long),
+            cp_size=self.cp_size,
+            cp_group=self.cp_group,
+        )
+        return zigzag_reconstruct_from_cp_parts(pos_parts, seq_dim=1)
+
+    def _gather_packed_cp_inputs(
+        self, x: torch.Tensor, position_ids: torch.Tensor, packed_seq_params
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        local_seq = x.shape[1]
+        cu_seqlens = self._packed_cu_seqlens(packed_seq_params, x.device)
+        full_seq = int(cu_seqlens[-1].item())
+        x_parts = _all_gather_cp(x, cp_size=self.cp_size, cp_group=self.cp_group)
+        full_x = reconstruct_packed_from_cp_parts(
+            x_parts, cu_seqlens_padded=cu_seqlens, cp_size=self.cp_size, dim=1
+        )
+
+        if position_ids.dim() == 1:
+            position_ids = position_ids.unsqueeze(0)
+        position_ids = position_ids.to(device=x.device, dtype=torch.long)
+        if position_ids.shape[-1] == full_seq:
+            return full_x, position_ids
+        if position_ids.shape[-1] != local_seq:
+            raise ValueError(
+                "GLM5 packed DynamicSparseAttention CP position_ids must be either local or full packed length, "
+                f"got {tuple(position_ids.shape)} for local_seq={local_seq}, full_seq={full_seq}."
+            )
+        pos_parts = _all_gather_cp(position_ids, cp_size=self.cp_size, cp_group=self.cp_group)
+        full_position_ids = reconstruct_packed_from_cp_parts(
+            pos_parts, cu_seqlens_padded=cu_seqlens, cp_size=self.cp_size, dim=1
+        )
+        return full_x, full_position_ids
+
+    def _gather_packed_cp_rotary(
+        self, cos: torch.Tensor, sin: torch.Tensor, packed_seq_params, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if cos.dim() != 3 or sin.dim() != 3:
+            return cos, sin
+        cu_seqlens = self._packed_cu_seqlens(packed_seq_params, device)
+        full_seq = int(cu_seqlens[-1].item())
+        if cos.shape[1] == full_seq and sin.shape[1] == full_seq:
+            return cos, sin
+        cos_parts = _all_gather_cp(cos, cp_size=self.cp_size, cp_group=self.cp_group)
+        sin_parts = _all_gather_cp(sin, cp_size=self.cp_size, cp_group=self.cp_group)
+        full_cos = reconstruct_packed_from_cp_parts(
+            cos_parts, cu_seqlens_padded=cu_seqlens, cp_size=self.cp_size, dim=1
+        )
+        full_sin = reconstruct_packed_from_cp_parts(
+            sin_parts, cu_seqlens_padded=cu_seqlens, cp_size=self.cp_size, dim=1
+        )
+        return full_cos, full_sin
+
+    @staticmethod
+    def _packed_cu_seqlens(packed_seq_params, device: torch.device) -> torch.Tensor:
+        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)
+        if cu_seqlens is None:
+            raise ValueError("GLM5 packed DynamicSparseAttention requires packed cu_seqlens.")
+        return cu_seqlens.to(device=device, dtype=torch.int32)
+
+    @staticmethod
+    def _slice_rotary_cache(
+        cos: torch.Tensor, sin: torch.Tensor, start: int, end: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if cos.dim() == 3 and sin.dim() == 3:
+            return cos[:, start:end, :], sin[:, start:end, :]
+        return cos, sin
+
+
+__all__ = [
+    "DSAIndexer",
+    "DSAIndexerLossAutoScaler",
+    "DynamicSparseAttention",
+    "RMSNorm",
+    "apply_rotary_emb",
+    "apply_rotary_pos_emb",
+    "build_rope_cache",
+    "build_rotary_embeddings",
+    "rotate_activation",
+    "rotate_half",
+]
diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/hca.py b/experimental/lite/megatron/lite/primitive/modules/attention/hca.py
new file mode 100644
index 00000000000..7593ee46505
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/attention/hca.py
@@ -0,0 +1,66 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def split_sinkhorn(
+    mixes: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    hc_mult: int,
+    iters: int,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    split_sizes = [hc_mult, hc_mult, hc_mult * hc_mult]
+    pre_mix, post_mix, comb_mix = mixes.split(split_sizes, dim=-1)
+    base_pre, base_post, base_comb = hc_base.to(dtype=mixes.dtype, device=mixes.device).split(
+        split_sizes, dim=-1
+    )
+    scale = hc_scale.to(dtype=mixes.dtype, device=mixes.device)
+    pre = torch.sigmoid(pre_mix * scale[0] + base_pre)
+    post = 2 * torch.sigmoid(post_mix * scale[1] + base_post)
+    comb_logits = (comb_mix * scale[2] + base_comb).view(*comb_mix.shape[:-1], hc_mult, hc_mult)
+    comb = torch.exp(comb_logits - comb_logits.max(dim=-1, keepdim=True).values)
+    for _ in range(iters):
+        comb = comb / comb.sum(dim=-1, keepdim=True).clamp(min=eps)
+        comb = comb / comb.sum(dim=-2, keepdim=True).clamp(min=eps)
+    return pre, post, comb
+
+
+class HyperConnection(nn.Module):
+    def __init__(self, hidden_size: int, hc_mult: int, sinkhorn_iters: int, eps: float):
+        super().__init__()
+        mix = (2 + hc_mult) * hc_mult
+        self.hidden_size = hidden_size
+        self.hc_mult = hc_mult
+        self.sinkhorn_iters = sinkhorn_iters
+        self.eps = eps
+        self.fn = nn.Parameter(torch.empty(mix, hc_mult * hidden_size, dtype=torch.float32))
+        self.base = nn.Parameter(torch.zeros(mix, dtype=torch.float32))
+        self.scale = nn.Parameter(torch.ones(3, dtype=torch.float32))
+        nn.init.xavier_uniform_(self.fn)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if x.dim() == 3:
+            x = x.unsqueeze(2).expand(*x.shape[:2], self.hc_mult, x.size(-1))
+        shape, dtype = x.shape, x.dtype
+        xf = x.flatten(2)
+        rms_inv = 1.0 / (xf.norm(dim=-1, keepdim=True) / math.sqrt(xf.shape[-1]) + self.eps)
+        mixes = F.linear(xf, self.fn.to(device=x.device, dtype=dtype)) * rms_inv
+        pre, post, comb = split_sinkhorn(
+            mixes, self.scale, self.base, self.hc_mult, self.sinkhorn_iters, self.eps
+        )
+        y = torch.sum(pre.unsqueeze(-1) * xf.view(shape), dim=2)
+        return y.to(dtype), post, comb
+
+    @staticmethod
+    def post(
+        x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor
+    ) -> torch.Tensor:
+        dtype = x.dtype
+        placed = post.to(dtype).unsqueeze(-1) * x.unsqueeze(-2)
+        mixed = torch.matmul(comb.to(dtype), residual.to(dtype))
+        return placed + mixed
diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/mhc.py b/experimental/lite/megatron/lite/primitive/modules/attention/mhc.py
new file mode 100644
index 00000000000..dba32bd4d8e
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/attention/mhc.py
@@ -0,0 +1,27 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MultiHeadHyperConnectionHead(nn.Module):
+    def __init__(self, hidden_size: int, hc_mult: int, eps: float):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.hc_mult = hc_mult
+        self.eps = eps
+        self.hc_fn = nn.Parameter(torch.empty(hc_mult, hc_mult * hidden_size, dtype=torch.float32))
+        self.hc_base = nn.Parameter(torch.zeros(hc_mult, dtype=torch.float32))
+        self.hc_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
+        nn.init.xavier_uniform_(self.hc_fn)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 3:
+            return x
+        shape, dtype = x.shape, x.dtype
+        xf = x.flatten(2).float()
+        rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + self.eps)
+        mixes = F.linear(xf, self.hc_fn.float()) * rsqrt
+        pre = torch.sigmoid(mixes * self.hc_scale.float() + self.hc_base.float()) + self.eps
+        y = torch.sum(pre.unsqueeze(-1) * xf.view(shape), dim=2)
+        return y.to(dtype)
diff --git a/experimental/lite/megatron/lite/primitive/modules/attention/mla.py b/experimental/lite/megatron/lite/primitive/modules/attention/mla.py
new file mode 100644
index 00000000000..792c27471a3
--- /dev/null
+++ b/experimental/lite/megatron/lite/primitive/modules/attention/mla.py
@@ -0,0 +1,418 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Shared Multi-Latent Attention primitive."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformer_engine.pytorch as te
+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,
+    YarnRotaryEmbedding,
+    _yarn_get_mscale,
+)
+
+from megatron.lite.primitive.parallel import (
+    ColumnParallelLinear,
+    ParallelState,
+    RowParallelLinear,
+    gather_from_sequence_parallel,
+)
+from megatron.lite.primitive.parallel.cp import (
+    zigzag_reconstruct_from_cp_parts,
+    zigzag_slice_for_cp,
+)
+from megatron.lite.primitive.parallel.thd import (
+    reconstruct_packed_from_cp_parts,
+    split_packed_to_cp_local,
+)
+
+_KEPT_PSP_FIELDS = (
+    "qkv_format",
+    "cu_seqlens_q",
+    "cu_seqlens_kv",
+    "cu_seqlens_q_padded",
+    "cu_seqlens_kv_padded",
+    "max_seqlen_q",
+    "max_seqlen_kv",
+)
+
+
+def _apply_mla_rope_bshd(t: torch.Tensor, freqs: torch.Tensor, *, mscale: float) -> torch.Tensor:
+    return _apply_rotary_pos_emb_bshd(
+        t, freqs, rotary_interleaved=False, mscale=mscale, mla_rotary_interleaved=True
+    )
+
+
+def _apply_mla_rope_thd(
+    t: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    freqs: torch.Tensor,
+    *,
+    mscale: float,
+    cp_group,
+) -> torch.Tensor:
+    return _apply_rotary_pos_emb_thd(
+        t,
+        cu_seqlens,
+        freqs,
+        rotary_interleaved=False,
+        mscale=mscale,
+        cp_group=cp_group,
+        mla_rotary_interleaved=True,
+    )
+
+
+class MultiLatentAttention(nn.Module):
+    """Native MLA composition using lite parallel linears and TE core attention."""
+
+    _cp_stream: torch.cuda.Stream | None = None
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_attention_heads: int,
+        q_lora_rank: int,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        ps: ParallelState,
+        rms_norm_eps: float = 1e-6,
+        rope_theta: float = 10_000.0,
+        rope_scaling: dict | None = None,
+        use_thd: bool = False,
+    ):
+        super().__init__()
+        if num_attention_heads % ps.tp_size != 0:
+            raise ValueError("num_attention_heads must be divisible by tensor parallel size")
+        self.ps = ps
+        self.num_heads_local = num_attention_heads // ps.tp_size
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.q_head_dim = qk_nope_head_dim + qk_rope_head_dim
+
+        self.linear_proj = RowParallelLinear(
+            num_attention_heads * v_head_dim,
+            hidden_size,
+            ps,
+            bias=False,
+        )
+        self.linear_q_down_proj = nn.Linear(hidden_size, q_lora_rank, bias=False)
+        self.linear_q_up_proj = ColumnParallelLinear(
+            q_lora_rank,
+            num_attention_heads * self.q_head_dim,
+            ps,
+            bias=False,
+            normalization="RMSNorm",
+            eps=rms_norm_eps,
+        )
+        self.linear_kv_down_proj = nn.Linear(
+            hidden_size,
+            kv_lora_rank + qk_rope_head_dim,
+            bias=False,
+        )
+        self.linear_kv_up_proj = ColumnParallelLinear(
+            kv_lora_rank,
+            num_attention_heads * (qk_nope_head_dim + v_head_dim),
+            ps,
+            bias=False,
+            normalization="RMSNorm",
+            eps=rms_norm_eps,
+        )
+
+        rope_scaling = dict(rope_scaling or {})
+        rope_type = rope_scaling.get("type", "rope")
+        if rope_type == "yarn":
+            factor = float(rope_scaling.get("factor", 1.0))
+            self.rotary = YarnRotaryEmbedding(
+                qk_rope_head_dim,
+                rotary_base=rope_theta,
+                scaling_factor=factor,
+                original_max_position_embeddings=int(
+                    rope_scaling.get("original_max_position_embeddings", 4096)
+                ),
+                beta_fast=float(rope_scaling.get("beta_fast", 32.0)),
+                beta_slow=float(rope_scaling.get("beta_slow", 1.0)),
+                mscale=float(rope_scaling.get("mscale", 1.0)),
+                mscale_all_dim=float(rope_scaling.get("mscale_all_dim", 1.0)),
+                cp_group=ps.cp_group if ps.cp_size > 1 else None,
+            )
+            attn_mscale = _yarn_get_mscale(factor, float(rope_scaling.get("mscale_all_dim", 1.0)))
+        elif rope_type == "rope":
+            self.rotary = RotaryEmbedding(
+                kv_channels=qk_rope_head_dim,
+                rotary_base=rope_theta,
+                use_cpu_initialization=False,
+                cp_group=ps.cp_group if ps.cp_size > 1 else None,
+            )
+            attn_mscale = 1.0
+        else:
+            raise ValueError(f"Unsupported MLA rope type: {rope_type!r}")
+        self._softmax_scale = attn_mscale * attn_mscale / math.sqrt(self.q_head_dim)
+        self._query_scale = 1.0
+
+        cp_kwargs = {}
+        if ps.cp_size > 1:
+            if MultiLatentAttention._cp_stream is None:
+                MultiLatentAttention._cp_stream = torch.cuda.Stream()
+            cp_kwargs = dict(
+                cp_group=ps.cp_group,
+                cp_global_ranks=ps.cp_global_ranks,
+                cp_stream=MultiLatentAttention._cp_stream,
+            )
+        self._use_torch_core = ps.cp_size > 1 and v_head_dim != self.q_head_dim
+        self.core_attn = None
+        if not self._use_torch_core:
+            dpa_kwargs = dict(cp_kwargs, softmax_scale=self._softmax_scale)
+            kv_channels = (
+                (self.q_head_dim, v_head_dim) if v_head_dim != self.q_head_dim else self.q_head_dim
+            )
+            self.core_attn = te.DotProductAttention(
+                num_attention_heads=self.num_heads_local,
+                kv_channels=kv_channels,
+                attention_dropout=0.0,
+                attn_mask_type="causal",
+                qkv_format="thd" if use_thd else "sbhd",
+                **dpa_kwargs,
+            )
+
+    def forward(self, x: torch.Tensor, packed_seq_params=None) -> torch.Tensor:
+        q_compressed = self.linear_q_down_proj(x)
+        kv_combined = self.linear_kv_down_proj(x)
+        kv_compressed, k_pos_emb = kv_combined.split(
+            [self.kv_lora_rank, self.qk_rope_head_dim],
+            dim=-1,
+        )
+        if self.ps.tp_size > 1:
+            k_pos_emb = gather_from_sequence_parallel(k_pos_emb, self.ps)
+
+        q_proj = self.linear_q_up_proj(q_compressed)
+        q = q_proj.view(
+            *q_proj.shape[:-1],
+            self.num_heads_local,
+            self.q_head_dim,
+        )
+        kv_proj = self.linear_kv_up_proj(kv_compressed)
+        kv = kv_proj.view(
+            *kv_proj.shape[:-1],
+            self.num_heads_local,
+            self.qk_nope_head_dim + self.v_head_dim,
+        )
+        q_nope, q_pos = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        k_nope, value = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        k_pos = k_pos_emb.unsqueeze(-2)
+
+        is_thd = packed_seq_params is not None
+        if is_thd:
+            q_nope = q_nope.squeeze(1)
+            q_pos = q_pos.squeeze(1)
+            k_nope = k_nope.squeeze(1)
+            value = value.squeeze(1)
+            k_pos = k_pos.squeeze(1)
+
+        q_pos, k_pos = self._apply_rope(q_pos, k_pos, packed_seq_params)
+        if k_pos.dim() == q_nope.dim():
+            k_pos = k_pos.expand(*q_nope.shape[:-1], self.qk_rope_head_dim)
+        else:
+            k_pos = k_pos.expand(-1, -1, self.num_heads_local, -1)
+        query = torch.cat([q_nope, q_pos], dim=-1).contiguous()
+        key = torch.cat([k_nope, k_pos], dim=-1).contiguous()
+        value = value.contiguous()
+        if self._query_scale != 1.0:
+            query = query * self._query_scale
+
+        if is_thd:
+            if self._use_torch_core:
+                out = self._torch_core_attention_thd(
+                    query,
+                    key,
+                    value,
+                    packed_seq_params=packed_seq_params,
+                ).reshape(query.size(0), 1, -1)
+            else:
+                psp_kwargs = {
+                    k: getattr(packed_seq_params, k)
+                    for k in _KEPT_PSP_FIELDS
+                    if getattr(packed_seq_params, k, None) is not None
+                }
+                assert self.core_attn is not None
+                out = self.core_attn(
+                    query,
+                    key,
+                    value,
+                    core_attention_bias_type="no_bias",
+                    attn_mask_type="padding_causal",
+                    **psp_kwargs,
+                ).reshape(query.size(0), 1, -1)
+        else:
+            if self._use_torch_core:
+                out = self._torch_core_attention(query, key, value)
+            else:
+                assert self.core_attn is not None
+                out = self.core_attn(query, key, value, core_attention_bias_type="no_bias")
+            if out.dim() > x.dim():
+                out = out.reshape(*out.shape[:-2], self.num_heads_local * self.v_head_dim)
+        return self.linear_proj(out)
+
+    def _torch_core_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        local_seq = query.size(0)
+        if self.ps.cp_size > 1:
+            from torch.distributed.nn.functional import all_gather
+
+            query_parts = all_gather(query.contiguous(), group=self.ps.cp_group)
+            key_parts = all_gather(key.contiguous(), group=self.ps.cp_group)
+            value_parts = all_gather(value.contiguous(), group=self.ps.cp_group)
+            query = zigzag_reconstruct_from_cp_parts(query_parts, seq_dim=0)
+            key = zigzag_reconstruct_from_cp_parts(key_parts, seq_dim=0)
+            value = zigzag_reconstruct_from_cp_parts(value_parts, seq_dim=0)
+
+        q = query.permute(1, 2, 0, 3)
+        k = key.permute(1, 2, 0, 3)
+        v = value.permute(1, 2, 0, 3)
+        scale = self._softmax_scale if self._query_scale == 1.0 else None
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            dropout_p=0.0,
+            is_causal=True,
+            scale=scale,
+        )
+        out = out.permute(2, 0, 1, 3).contiguous()
+        if self.ps.cp_size > 1:
+            out = zigzag_slice_for_cp(out, self.ps.cp_rank, self.ps.cp_size, seq_dim=0)
+            if out.size(0) != local_seq:
+                raise RuntimeError("CP MLA output shard has unexpected sequence length.")
+        return out
+
+    def _torch_core_attention_thd(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        *,
+        packed_seq_params,
+    ) -> torch.Tensor:
+        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)
+        if cu_seqlens is None:
+            raise ValueError("Packed THD MLA fallback requires cu_seqlens.")
+
+        local_tokens = query.size(0)
+        if self.ps.cp_size > 1:
+            from torch.distributed.nn.functional import all_gather
+
+            query = reconstruct_packed_from_cp_parts(
+                list(all_gather(query.contiguous(), group=self.ps.cp_group)),
+                cu_seqlens_padded=cu_seqlens,
+                cp_size=self.ps.cp_size,
+                dim=0,
+            )
+            key = reconstruct_packed_from_cp_parts(
+                list(all_gather(key.contiguous(), group=self.ps.cp_group)),
+                cu_seqlens_padded=cu_seqlens,
+                cp_size=self.ps.cp_size,
+                dim=0,
+            )
+            value = reconstruct_packed_from_cp_parts(
+                list(all_gather(value.contiguous(), group=self.ps.cp_group)),
+                cu_seqlens_padded=cu_seqlens,
+                cp_size=self.ps.cp_size,
+                dim=0,
+            )
+
+        outputs = []
+        scale = self._softmax_scale if self._query_scale == 1.0 else None
+        for idx in range(int(cu_seqlens.numel()) - 1):
+            start = int(cu_seqlens[idx].item())
+            end = int(cu_seqlens[idx + 1].item())
+            if end <= start:
+                continue
+            q = query[start:end].permute(1, 0, 2).unsqueeze(0)
+            k = key[start:end].permute(1, 0, 2).unsqueeze(0)
+            v = value[start:end].permute(1, 0, 2).unsqueeze(0)
+            out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                dropout_p=0.0,
+                is_causal=True,
+                scale=scale,
+            )
+            outputs.append(out.squeeze(0).permute(1, 0, 2).contiguous())
+        full_out = torch.cat(outputs, dim=0) if outputs else value.new_empty(value.shape)
+        if self.ps.cp_size <= 1:
+            return full_out
+        local_out = split_packed_to_cp_local(
+            full_out,
+            cu_seqlens_padded=cu_seqlens,
+            cp_size=self.ps.cp_size,
+            cp_rank=self.ps.cp_rank,
+            dim=0,
+        )
+        if local_out.size(0) != local_tokens:
+            raise RuntimeError("CP THD MLA output shard has unexpected token count.")
+        return local_out
+
+    def _apply_rope(self, q_pos: torch.Tensor, k_pos: torch.Tensor, packed_seq_params):
+        is_thd = packed_seq_params is not None
+        if is_thd:
+            max_q = getattr(packed_seq_params, "max_seqlen_q", None)
+            max_kv = getattr(packed_seq_params, "max_seqlen_kv", None)
+            seq_len = (
+                int(max(max_q, max_kv))
+                if max_q is not None and max_kv is not None
+                else int(packed_seq_params.cu_seqlens_q[-1])
+            )
+            freqs = self.rotary(seq_len, packed_seq=True)
+            if isinstance(freqs, tuple):
+                freqs, mscale = freqs
+            else:
+                mscale = 1.0
+            q_pos = _apply_mla_rope_thd(
+                q_pos,
+                packed_seq_params.cu_seqlens_q,
+                freqs,
+                mscale=mscale,
+                cp_group=self.ps.cp_group,
+            )
+            k_pos = _apply_mla_rope_thd(
+                k_pos,
+                packed_seq_params.cu_seqlens_kv,
+                freqs,
+                mscale=mscale,
+                cp_group=self.ps.cp_group,
+            )
+            return q_pos, k_pos
+
+        seq_len = q_pos.size(0) * self.ps.cp_size
+        freqs = self.rotary(seq_len)
+        if isinstance(freqs, tuple):
+            freqs, mscale = freqs
+        else:
+            mscale = 1.0
+        return (
+            _apply_mla_rope_bshd(q_pos, freqs, mscale=mscale),
+            _apply_mla_rope_bshd(k_pos, freqs, mscale=mscale),
+        )
+
+
+__all__ = ["MultiLatentAttention"]