diff --git a/nemo_automodel/_transformers/capabilities.py b/nemo_automodel/_transformers/capabilities.py
index e81cbb5788..4086ba4e44 100644
--- a/nemo_automodel/_transformers/capabilities.py
+++ b/nemo_automodel/_transformers/capabilities.py
@@ -98,7 +98,7 @@ def _uses_magi_attention(model: "nn.Module") -> bool:
     """True when the model uses the MagiAttention (FFA / context-parallel) backend.
 
     MagiAttention implements context parallelism via its own load-balancing
-    dispatch (see ``components/distributed/magi_attn_utils.py``), so it supports CP.
+    dispatch (see ``components/distributed/context_parallel/magi.py``), so it supports CP.
     """
     backend = getattr(model, "backend", None)
     return getattr(backend, "attn", None) == "magi"
@@ -336,7 +336,7 @@ def supports_cp_with_sequence_packing(self) -> bool:
 
         MagiAttention dispatches the packed sequence across the CP group with its
         own load-balancing solver and a per-document varlen mask, so it supports
-        CP + packing (see ``magi_attn_utils.magi_prepare_packed_cp``). Models
+        CP + packing (see ``context_parallel.magi.magi_prepare_packed_cp``). Models
         with native THD support own their packed CP path in TileLang attention."""
         model = self._model
         if not self.supports_sequence_packing:
diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py
index 285abfe2c0..0599ee5232 100644
--- a/nemo_automodel/_transformers/infrastructure.py
+++ b/nemo_automodel/_transformers/infrastructure.py
@@ -748,7 +748,7 @@ def apply_model_infrastructure(
     # is not excluded by the _uses_te_attention check, so gate on ep_size: only
     # dense (non-MoE) models need this pass.
     if mesh.cp_size > 1 and mesh.ep_size <= 1 and not _uses_te_attention(model):
-        from nemo_automodel.components.distributed.cp_utils import (
+        from nemo_automodel.components.distributed.context_parallel.utils import (
             attach_context_parallel_hooks,
             attach_cp_sdpa_hooks,
         )
diff --git a/nemo_automodel/components/attention/utils.py b/nemo_automodel/components/attention/utils.py
index f683abc824..1bb026ffd5 100644
--- a/nemo_automodel/components/attention/utils.py
+++ b/nemo_automodel/components/attention/utils.py
@@ -108,7 +108,7 @@ def attn_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **call_kwargs:
             )
         # requires magi_attention; the guards above are exercised on CPU but the
         # kernel build is not, so exclude it from coverage.
-        from nemo_automodel.components.distributed.magi_attn_utils import (  # pragma: no cover - requires magi_attention
+        from nemo_automodel.components.distributed.context_parallel.magi import (  # pragma: no cover - requires magi_attention
             make_magi_attn_func,
         )
 
diff --git a/nemo_automodel/components/datasets/loader.py b/nemo_automodel/components/datasets/loader.py
index 5169ce8b2c..cf20ad6fc5 100644
--- a/nemo_automodel/components/datasets/loader.py
+++ b/nemo_automodel/components/datasets/loader.py
@@ -589,6 +589,13 @@ def dataset_builds_on_all_ranks(self) -> bool:
         """Whether dataset construction must bypass rank-zero-first ordering."""
         return isinstance(self.dataset_config, AllRanksDatasetConfig)
 
+    @property
+    def emits_thd(self) -> bool:
+        """Whether this configuration produces THD-formatted batches."""
+        from nemo_automodel.components.datasets.utils import packed_sequence_thd_collater
+
+        return isinstance(self.packing, ThdPackingConfig) or self.collate_fn is packed_sequence_thd_collater
+
     def _build_dataset(
         self,
         *,
diff --git a/nemo_automodel/components/distributed/context_parallel/__init__.py b/nemo_automodel/components/distributed/context_parallel/__init__.py
new file mode 100644
index 0000000000..b07012b80f
--- /dev/null
+++ b/nemo_automodel/components/distributed/context_parallel/__init__.py
@@ -0,0 +1,19 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Context-parallel batch sharding."""
+
+from nemo_automodel.components.distributed.context_parallel.sharder import ContextParallelSharder
+
+__all__ = ["ContextParallelSharder"]
diff --git a/nemo_automodel/components/distributed/magi_attn_utils.py b/nemo_automodel/components/distributed/context_parallel/magi.py
similarity index 99%
rename from nemo_automodel/components/distributed/magi_attn_utils.py
rename to nemo_automodel/components/distributed/context_parallel/magi.py
index 4e166a4e7d..b9872aca69 100644
--- a/nemo_automodel/components/distributed/magi_attn_utils.py
+++ b/nemo_automodel/components/distributed/context_parallel/magi.py
@@ -50,7 +50,7 @@
 import torch
 import torch.distributed as dist
 
-from nemo_automodel.components.distributed.cp_utils import _make_cp_batch_and_ctx, make_cp_batch_for_te
+from nemo_automodel.components.distributed.context_parallel.utils import _make_cp_batch_and_ctx, make_cp_batch_for_te
 
 logger = logging.getLogger(__name__)
 
@@ -753,7 +753,7 @@ def prepare_llm_batch(
         context). ``local_indices`` is the global stream position of every
         local token on the paths that dispatch the sequence (magi's
         ``get_position_ids``), None otherwise; the framework installs it on
-        the magi ContextParallelismSharder for the token-tensor verbs.
+        the magi ContextParallelSharder for the token-tensor verbs.
         """
         # cp=1 prefix-tree mask: the datasets layer cannot import this module (component
         # independence), so the collate attaches the tree structure and the spec is built
@@ -843,7 +843,7 @@ def make_cp_batch(
             return_local_indices: Also return the local-token global index map
                 from the dispatch that just ran (None on paths that do not
                 dispatch, e.g. cp=1 THD conversion and the VLM domain). Used
-                by the magi ContextParallelismSharder's token verbs.
+                by the magi ContextParallelSharder's token verbs.
 
         Returns:
             The dispatched (magi-sharded) batch, or ``(batch, local_indices)``
diff --git a/nemo_automodel/components/distributed/mamba_cp.py b/nemo_automodel/components/distributed/context_parallel/mamba.py
similarity index 100%
rename from nemo_automodel/components/distributed/mamba_cp.py
rename to nemo_automodel/components/distributed/context_parallel/mamba.py
diff --git a/nemo_automodel/components/distributed/cp_sharder.py b/nemo_automodel/components/distributed/context_parallel/sharder.py
similarity index 87%
rename from nemo_automodel/components/distributed/cp_sharder.py
rename to nemo_automodel/components/distributed/context_parallel/sharder.py
index a03e40d9e1..956fde0923 100644
--- a/nemo_automodel/components/distributed/cp_sharder.py
+++ b/nemo_automodel/components/distributed/context_parallel/sharder.py
@@ -14,13 +14,13 @@
 
 """Context-parallel batch-sharding contract.
 
-Every CP backend is a :class:`ContextParallelismSharder`. A model that owns its CP batch
+Every CP backend is a :class:`ContextParallelSharder`. A model that owns its CP batch
 sharding and attention transport returns one from
 ``prepare_model_inputs_for_cp`` under the ``"cp_sharder"`` batch key; the
 framework constructs its own for the remaining backends (torch
 ``context_parallel`` round-robin, TE/THD, MagiAttention) so
-the CP dispatch (``cp_utils.prepare_cp_forward``) reduces to resolving a sharder and calling
-``shard_batch``. This replaces the retired private batch keys
+constructing :class:`ContextParallelSharder` resolves and configures the backend;
+callers then invoke ``sharder.shard(batch)``. This replaces the retired private batch keys
 (``_cp_make_batch_fn``, ``_cp_metadata_seq_dims``, ``_cp_metadata_pad_values``,
 ``_cp_full_logits_grad_touch``).
 
@@ -44,7 +44,7 @@
 shared contiguous-shard batch prep used by models whose CP ranks own contiguous
 sequence slices (Gemma4, DeepSeek V4), and the torch ``context_parallel``
 round-robin load-balanced prep with its index map. The TE/THD and magi preps
-live with their dependencies (``cp_utils``, ``magi_attn_utils``); the
+live with their dependencies (``context_parallel.utils``, ``context_parallel.magi``); the
 dispatcher wraps them into sharders at resolution time.
 """
 
@@ -57,6 +57,7 @@
 
 import torch
 import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
 
 
 def _cp_rank(cp_mesh) -> int:
@@ -236,15 +237,14 @@ class ShardLayout:
     input_token_stream_positions: torch.Tensor | None = None
 
 
-@dataclass
-class ContextParallelismSharder:
+class ContextParallelSharder:
     """CP backend description: how a batch is sharded and where local tokens live.
 
     Attributes:
         shard_batch: ``(cp_mesh, tp_mesh, batch, *, loss_mask=None,
             padding_token_id=0) -> (ctx_factory, batch, ShardLayout | None)``.
             Pads and shards the batch, installs any backend-owned attention
-            transport, and reports the shard layout it computed; the dispatch
+            transport, and reports the shard layout it computed; :meth:`shard`
             stores it as ``shard_layout`` for the token verbs.
         local_token_global_indices: ``(cp_mesh, padded_seq_len, device) ->
             LongTensor`` with the global position of each local token —
@@ -252,22 +252,126 @@ class ContextParallelismSharder:
             data-dependent layouts, whose partition arrives with
             ``shard_layout`` (their token verbs raise before the first shard).
         shard_layout: The :class:`ShardLayout` of the last ``shard_batch``, set
-            by the dispatch. Sharders are built per resolution/hook call, so
+            by :meth:`shard`. Sharders are built per resolution/hook call, so
             the layout never leaks across steps.
     """
 
     shard_batch: Callable[..., tuple[Callable, dict[str, Any], "ShardLayout | None"]]
     local_token_global_indices: Callable[..., torch.Tensor] | None
-    shard_layout: "ShardLayout | None" = None
+    shard_layout: "ShardLayout | None"
+    _cp_mesh: Any
+    _tp_mesh: Any
+    _loss_mask: torch.Tensor | None
+    _padding_token_id: int
 
-    def _indices(self, cp_mesh, padded_seq_len: int, device) -> torch.Tensor:
+    def __init__(
+        self,
+        model: torch.nn.Module | None = None,
+        device_mesh: DeviceMesh | None = None,
+        batch: dict[str, Any] | None = None,
+        *,
+        shard_batch: Callable[..., tuple[Callable, dict[str, Any], "ShardLayout | None"]] | None = None,
+        local_token_global_indices: Callable[..., torch.Tensor] | None = None,
+        shard_layout: "ShardLayout | None" = None,
+        padding_token_id: int = 0,
+        num_chunks: int = 1,
+        loss_mask: torch.Tensor | None = None,
+        invoke_pre_embed: bool = True,
+        extra_seq_buffers: dict[str, int] | None = None,
+    ) -> None:
+        """Construct a strategy sharder or resolve one for a forward.
+
+        Args:
+            model: Model whose attention backend and CP preparation hook select
+                the sharding strategy, or None for the generic strategy. Omit
+                when constructing directly from ``shard_batch``.
+            device_mesh: Device mesh containing optional ``cp`` and ``tp`` axes.
+                Direct strategy construction uses it to configure the sharder;
+                omit it only when returning an unresolved model-owned strategy.
+            batch: Mutable input mapping required when resolving a strategy.
+                Token tensors normally have shape
+                [batch, sequence, ...]; THD source batches declare
+                ``qkv_format="thd"`` and are flattened during :meth:`shard`.
+                Model hook metadata is merged into this mapping in place.
+            shard_batch: Optional backend callback for direct strategy
+                construction. It accepts token tensors with backend-defined
+                layouts and returns the sharded batch plus its layout.
+            local_token_global_indices: Optional callback returning a tensor of
+                shape [local_tokens] with each local token's global position.
+            shard_layout: Optional captured layout for direct strategy
+                construction. Tensor fields use the layouts documented by
+                :class:`ShardLayout`.
+            padding_token_id: Value used to pad ``input_ids`` on the sequence axis.
+            num_chunks: Number of THD chunks created during sharding.
+            loss_mask: Optional tensor of shape [batch, sequence] sharded with
+                the batch.
+            invoke_pre_embed: Whether to invoke a model-owned CP preparation hook.
+            extra_seq_buffers: Additional batch keys mapped to their sequence axes.
+        """
+        if shard_batch is not None:
+            has_resolution_args = (
+                model is not None
+                or batch is not None
+                or num_chunks != 1
+                or not invoke_pre_embed
+                or extra_seq_buffers is not None
+            )
+            if has_resolution_args:
+                raise TypeError("shard_batch is mutually exclusive with model, batch, and strategy-resolution options")
+
+            self.shard_batch = shard_batch
+            self.local_token_global_indices = local_token_global_indices
+            self.shard_layout = shard_layout
+            mesh_dim_names = getattr(device_mesh, "mesh_dim_names", ())
+            self._cp_mesh = device_mesh["cp"] if "cp" in mesh_dim_names else None
+            self._tp_mesh = device_mesh["tp"] if "tp" in mesh_dim_names else None
+            self._loss_mask = loss_mask
+            self._padding_token_id = padding_token_id
+            return
+        if local_token_global_indices is not None or shard_layout is not None:
+            raise TypeError("local_token_global_indices and shard_layout require shard_batch")
+        if batch is None:
+            raise TypeError("batch is required when shard_batch is not provided")
+
+        from nemo_automodel.components.distributed.context_parallel.utils import _prepare_cp_sharder
+
+        resolved = _prepare_cp_sharder(
+            model,
+            device_mesh,
+            batch,
+            padding_token_id=padding_token_id,
+            num_chunks=num_chunks,
+            loss_mask=loss_mask,
+            invoke_pre_embed=invoke_pre_embed,
+            extra_seq_buffers=extra_seq_buffers,
+        )
+        self.shard_batch = resolved.shard_batch
+        self.local_token_global_indices = resolved.local_token_global_indices
+        self.shard_layout = resolved.shard_layout
+        self._cp_mesh = resolved._cp_mesh
+        self._tp_mesh = resolved._tp_mesh
+        self._loss_mask = resolved._loss_mask
+        self._padding_token_id = resolved._padding_token_id
+
+    def shard(self, batch: dict[str, Any]) -> tuple[Callable, dict[str, Any]]:
+        """Shard a batch and retain its layout for token-aligned tensors."""
+        ctx, batch, self.shard_layout = self.shard_batch(
+            self._cp_mesh,
+            self._tp_mesh,
+            batch,
+            loss_mask=self._loss_mask,
+            padding_token_id=self._padding_token_id,
+        )
+        return ctx, batch
+
+    def _indices(self, padded_seq_len: int, device) -> torch.Tensor:
         layout = self.shard_layout or _NO_SHARD_LAYOUT
         captured = layout.local_token_global_indices
         if captured is not None:
             # Data-dependent layout: use the partition the shard reported, and
             # validate the requested length against it so a mismatched tensor
             # cannot be silently mis-sharded.
-            cp_size = cp_mesh.size() if cp_mesh is not None else 1
+            cp_size = self._cp_mesh.size() if self._cp_mesh is not None else 1
             expected = captured.numel() * cp_size
             if padded_seq_len != expected:
                 raise ValueError(
@@ -278,14 +382,14 @@ def _indices(self, cp_mesh, padded_seq_len: int, device) -> torch.Tensor:
             return captured.reshape(-1).to(device=device, dtype=torch.long)
         if self.local_token_global_indices is None:
             raise NotImplementedError(
-                "This ContextParallelismSharder has a data-dependent token layout; its index map "
+                "This ContextParallelSharder has a data-dependent token layout; its index map "
                 "arrives with the shard layout — token-tensor shard/gather are unavailable before "
                 "the first shard."
             )
-        return self.local_token_global_indices(cp_mesh, padded_seq_len, device)
+        return self.local_token_global_indices(self._cp_mesh, padded_seq_len, device)
 
     def shard_token_tensor(
-        self, cp_mesh, tensor: torch.Tensor, seq_dim: int = 1, fill: float | int | None = None
+        self, tensor: torch.Tensor, seq_dim: int = 1, fill: float | int | None = None
     ) -> torch.Tensor:
         """Shard a full-length token-aligned tensor exactly like the model inputs.
 
@@ -330,16 +434,15 @@ def shard_token_tensor(
                 tensor = _pad_tensor_seq_dim_(tensor, seq_dim, layout.padded_seq_len - length, fill)
             else:
                 raise ValueError(
-                    f"This ContextParallelismSharder sharded a batch of padded_seq_len={layout.padded_seq_len} "
+                    f"This ContextParallelSharder sharded a batch of padded_seq_len={layout.padded_seq_len} "
                     f"(original_seq_len={layout.original_seq_len}), got a tensor of length {length} on dim {seq_dim}. "
                     "Pass the original-length tensor with an explicit `fill`, or pre-pad it yourself."
                 )
-        indices = self._indices(cp_mesh, tensor.shape[seq_dim], tensor.device)
+        indices = self._indices(tensor.shape[seq_dim], tensor.device)
         return shard_token_tensor_by_indices(tensor, indices, seq_dim=seq_dim)
 
     def gather_token_tensor(
         self,
-        cp_mesh,
         tensor: torch.Tensor,
         seq_dim: int = 1,
         trim: bool = False,
@@ -355,9 +458,9 @@ def gather_token_tensor(
         (nothing to trim to).
         """
         layout = self.shard_layout or _NO_SHARD_LAYOUT
-        padded_seq_len = tensor.shape[seq_dim] * (cp_mesh.size() if cp_mesh is not None else 1)
-        indices = self._indices(cp_mesh, padded_seq_len, tensor.device)
-        full = gather_token_tensor_by_indices(cp_mesh, tensor, indices, seq_dim=seq_dim)
+        padded_seq_len = tensor.shape[seq_dim] * (self._cp_mesh.size() if self._cp_mesh is not None else 1)
+        indices = self._indices(padded_seq_len, tensor.device)
+        full = gather_token_tensor_by_indices(self._cp_mesh, tensor, indices, seq_dim=seq_dim)
         if not trim:
             return full
         if layout.padded_seq_len is not None and full.shape[seq_dim] != layout.padded_seq_len:
@@ -377,7 +480,7 @@ def gather_token_tensor(
         if layout.original_seq_len is not None:
             return full.narrow(seq_dim, 0, layout.original_seq_len)
         raise NotImplementedError(
-            "This ContextParallelismSharder has no shard layout to trim to; "
+            "This ContextParallelSharder has no shard layout to trim to; "
             "gather with trim=False and restore the layout with the batch metadata "
             "(padding_mask / cu_seqlens)."
         )
@@ -718,7 +821,7 @@ def shard_batch_load_balanced(
 ):
     """Shard a batch with torch ``context_parallel`` round-robin load balancing.
 
-    ``ContextParallelismSharder.shard_batch`` implementation for the default framework-owned CP
+    ``ContextParallelSharder.shard_batch`` implementation for the default framework-owned CP
     path (layout ``"round_robin"``, indices from
     :func:`round_robin_local_indices`). Assumes an active CP mesh (size > 1).
     ``padding_token_id`` is accepted per the contract but unused: CP-pad slots
@@ -729,9 +832,8 @@ def shard_batch_load_balanced(
         ``(ctx_factory, batch, ShardLayout)`` where entering ``ctx_factory()``
         installs the SDPA-kernel + ``context_parallel`` context for the forward.
     """
-    # Call-time import: the torch-CP transport machinery stays in cp_utils
-    # (NeMo-RL imports it from there), and cp_utils imports this module.
-    from nemo_automodel.components.distributed.cp_utils import (  # noqa: PLC0415
+    # Call-time import avoids a cycle: utils imports this module's strategy helpers.
+    from nemo_automodel.components.distributed.context_parallel.utils import (  # noqa: PLC0415
         _shard_grad_buffer_for_cp,
         create_context_parallel_ctx,
         get_train_context,
@@ -865,7 +967,7 @@ def shard_batch_aux_only(
         The layout's ``padded_seq_len`` is what the model must pad its primary
         stream to before sharding.
     """
-    from nemo_automodel.components.distributed.cp_utils import (  # noqa: PLC0415
+    from nemo_automodel.components.distributed.context_parallel.utils import (  # noqa: PLC0415
         create_context_parallel_ctx,
         get_train_context,
     )
diff --git a/nemo_automodel/components/distributed/cp_utils.py b/nemo_automodel/components/distributed/context_parallel/utils.py
similarity index 85%
rename from nemo_automodel/components/distributed/cp_utils.py
rename to nemo_automodel/components/distributed/context_parallel/utils.py
index 488769f48c..57f58eb783 100644
--- a/nemo_automodel/components/distributed/cp_utils.py
+++ b/nemo_automodel/components/distributed/context_parallel/utils.py
@@ -14,13 +14,13 @@
 
 import contextlib
 from functools import partial
-from typing import List, Optional, Set
+from typing import Any, List, Optional, Set
 
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     ShardLayout,
     identity_local_indices,
     round_robin_local_indices,
@@ -269,25 +269,68 @@ def _mesh_dim_size(device_mesh, dim: str) -> int:
     return submesh.size() if submesh is not None else 0
 
 
-def prepare_cp_forward(
-    model,
-    device_mesh,
-    batch,
+def _attention_backend(model) -> str | None:
+    """Read the configured attention backend from a live model."""
+    backend = getattr(getattr(model, "backend", None), "attn", None)
+    if backend is not None:
+        return str(backend)
+    config = getattr(model, "config", None)
+    for candidate in (config, getattr(config, "text_config", None)):
+        implementation = getattr(candidate, "_attn_implementation", None) or getattr(
+            candidate, "_attn_implementation_internal", None
+        )
+        if implementation is not None:
+            return str(implementation)
+    return None
+
+
+def _uses_te_attention(model) -> bool:
+    """Whether a live model uses native or injected TE attention."""
+    return _attention_backend(model) == "te" or bool(getattr(model, "_te_attention_injected", False))
+
+
+def _is_multimodal_model(model) -> bool:
+    """Whether a live model owns a vision or audio tower."""
+    config = getattr(model, "config", None)
+    return any(getattr(config, name, None) is not None for name in ("vision_config", "audio_config")) or any(
+        getattr(model, name, None) is not None
+        for name in ("visual", "vision_model", "vision_tower", "audio_model", "audio_tower")
+    )
+
+
+def _magi_state_from_model(model, device_mesh):
+    """Recreate the per-forward Magi handle from the model and device mesh."""
+    if model is None or _attention_backend(model) != "magi":
+        return None
+    from nemo_automodel.components.distributed.context_parallel.magi import MagiState, get_cp_group
+
+    cp_group = get_cp_group(device_mesh)
+    return MagiState(
+        enabled=True,
+        custom=getattr(getattr(model, "backend", None), "attn", None) == "magi",
+        cp_group=cp_group,
+        cp_size=cp_group.size() if cp_group is not None else 1,
+        domain="vlm" if _is_multimodal_model(model) else "llm",
+        device_mesh=device_mesh,
+    )
+
+
+def _prepare_cp_sharder(
+    model: Any,
+    device_mesh: DeviceMesh | None,
+    batch: dict[str, Any],
     *,
-    magi=None,
-    use_te: bool = False,
     padding_token_id: int = 0,
     num_chunks: int = 1,
-    loss_mask=None,
+    loss_mask: torch.Tensor | None = None,
     invoke_pre_embed: bool = True,
     extra_seq_buffers: Optional[dict[str, int]] = None,
-):
-    """Single CP dispatch for a training/eval forward: hook -> sharder -> (ctx, batch).
+) -> ContextParallelSharder:
+    """Resolve and configure a CP sharder for its public constructor.
 
-    Collapses the per-recipe CP branching into one call: the model hook may
-    return a ContextParallelismSharder, ``_make_cp_batch_and_ctx`` resolves it against the
-    framework-owned sharders (magi / TE / generic torch ``context_parallel``)
-    and calls ``shard_batch``. When CP is active and the model exposes
+    The model hook may return a ContextParallelSharder; otherwise this
+    function resolves a framework-owned sharder from the live model's attention
+    backend and the batch's token layout. When CP is active and the model exposes
     ``prepare_model_inputs_for_cp``, that sharder-only hook is invoked directly as
     a plain method (it constructs a sharder and touches no weights; embed / vision
     splice / sequence shard run in the model's own forward per microbatch).
@@ -295,15 +338,8 @@ def prepare_cp_forward(
     Args:
         model: The (first) model part, or None (e.g. no-model contexts).
         device_mesh: The full device mesh (``cp``/``tp`` submeshes are read).
-        batch: The full-sequence batch; mutated and sharded in place.
-        magi: Optional recipe MagiState, threaded to ``_make_cp_batch_and_ctx``
-            where it occupies the same dispatch rung as the TE path. Its
-            recipe domain is bound at ``setup_magi``; for llm-domain magi the
-            model hook is skipped (mirrors the recipes' historical branching),
-            while vlm-domain magi still runs the pre-embed first (vision stays
-            on SDPA under magi).
-        use_te: THD-packed collator is active (TE/THD sharding; also magi's
-            ``is_thd``).
+        batch: The full-sequence batch. Model hook updates are merged in place;
+            :meth:`ContextParallelSharder.shard` performs the actual sharding.
         padding_token_id: Pad sentinel for ``input_ids``.
         num_chunks: THD chunk count, forwarded to the hook and TE sharding.
         loss_mask: Optional per-token mask forwarded to the batch sharding.
@@ -315,22 +351,24 @@ def prepare_cp_forward(
             batch on the generic torch path (rejected on the TE THD path;
             ignored by backends that own their transport).
     Returns:
-        ``(ctx_factory, batch, sharder)`` — the resolved :class:`ContextParallelismSharder`
-        (the identity sharder when no CP prep applies), whose token
-        verbs keep per-token tensors aligned with the sharded inputs.
+        The resolved and mesh-configured :class:`ContextParallelSharder` (the
+        identity sharder when no CP prep applies).
     """
 
-    magi_enabled = magi is not None and getattr(magi, "enabled", False)
+    batch_is_thd = batch.get("qkv_format") == "thd"
+    magi_state = _magi_state_from_model(model, device_mesh)
+    magi_enabled = magi_state is not None and getattr(magi_state, "enabled", False)
+    backend_uses_thd = batch_is_thd and (magi_enabled or _uses_te_attention(model))
     cp_sharder = None
     has_hook = model is not None and hasattr(model, "prepare_model_inputs_for_cp")
     effective_cp_size = _mesh_dim_size(device_mesh, "cp")
 
     # llm-domain magi replaces the whole batch prep (no model has both a CP
     # hook and magi); vlm-domain magi composes with the vision pre-embed.
-    magi_replaces_hook = magi_enabled and getattr(magi, "domain", "llm") == "llm"
-    model_owns_thd = use_te and bool(getattr(model, "supports_thd", False))
+    magi_replaces_hook = magi_enabled and getattr(magi_state, "domain", "llm") == "llm"
+    model_owns_thd = batch_is_thd and bool(getattr(model, "supports_thd", False))
     if (effective_cp_size > 1 or model_owns_thd) and has_hook and not magi_replaces_hook and invoke_pre_embed:
-        # Every CP hook is sharder-only: it constructs a ContextParallelismSharder
+        # Every CP hook is sharder-only: it constructs a ContextParallelSharder
         # and consumes nothing (embed / vision splice / sequence shard happen in the
         # model's own forward). It touches no weights — a plain method call, no
         # ``__call__`` routing or FSDP2 unshard — and leaves the batch intact.
@@ -338,32 +376,41 @@ def prepare_cp_forward(
         cp_sharder = prepared.get("cp_sharder")
         batch.update({key: value for key, value in prepared.items() if key != "cp_sharder"})
 
-    return _make_cp_batch_and_ctx(
-        device_mesh,
-        batch,
-        loss_mask,
-        use_te=use_te,
-        padding_token_id=padding_token_id,
+    cp_mesh = _get_submesh(device_mesh, "cp")
+    if backend_uses_thd and extra_seq_buffers:
+        raise ValueError("extra_seq_buffers are not supported by the TE THD context-parallel path")
+    strategy = _resolve_cp_sharder(
+        cp_mesh,
+        cp_sharder,
+        magi=magi_state,
+        is_thd=backend_uses_thd,
         num_chunks=num_chunks,
-        magi=magi,
+        seq_lens_padding_value=-1000,
         model=model,
-        cp_sharder=cp_sharder,
         extra_seq_buffers=extra_seq_buffers,
     )
+    return ContextParallelSharder(
+        device_mesh=device_mesh,
+        shard_batch=strategy.shard_batch,
+        local_token_global_indices=strategy.local_token_global_indices,
+        shard_layout=strategy.shard_layout,
+        loss_mask=loss_mask,
+        padding_token_id=padding_token_id,
+    )
 
 
 def _resolve_cp_sharder(
     cp_mesh,
-    model_sharder: Optional[ContextParallelismSharder],
+    model_sharder: Optional[ContextParallelSharder],
     *,
     magi,
-    use_te: bool,
+    is_thd: bool,
     num_chunks: int,
     seq_lens_padding_value: int,
     model,
     extra_seq_buffers: Optional[dict[str, int]] = None,
-) -> ContextParallelismSharder:
-    """Resolve the ContextParallelismSharder for this forward: model-owned > magi > TE > generic > none.
+) -> ContextParallelSharder:
+    """Resolve the ContextParallelSharder for this forward: model-owned > magi > TE > generic > none.
 
     Always returns a sharder: when no CP prep applies, an identity sharder,
     so callers hold working token verbs at every cp_size and
@@ -378,7 +425,7 @@ def _resolve_cp_sharder(
     """
     cp_active = cp_mesh is not None and cp_mesh.size() > 1
 
-    # A model that owns its CP attention returns a ContextParallelismSharder from its CP
+    # A model that owns its CP attention returns a ContextParallelSharder from its CP
     # input-prep hook. Honor it instead of any framework-owned path so the
     # implementation stays with the model.
     if model_sharder is not None:
@@ -388,7 +435,7 @@ def _resolve_cp_sharder(
         # Backend-owned prep (MagiAttention): magi manages its own CP transport,
         # so like the TE path shard_batch returns (nullcontext, prepped_batch).
         # All magi internals (HF-vs-custom, recipe domain, cp group) stay in
-        # magi_attn_utils. The dispatch-solver partition is data-dependent, so
+        # context_parallel.magi. The dispatch-solver partition is data-dependent, so
         # shard_batch installs the index map it just computed (magi's
         # get_position_ids) on the sharder for the token verbs.
         def _shard_batch_magi(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_id=0):
@@ -399,7 +446,7 @@ def _shard_batch_magi(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_
                 batch,
                 padding_token_id=padding_token_id,
                 num_chunks=num_chunks,
-                is_thd=use_te,
+                is_thd=is_thd,
                 model=model,
                 return_local_indices=True,
             )
@@ -424,9 +471,9 @@ def _shard_batch_magi(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_
                 )
             return contextlib.nullcontext, prepped, layout
 
-        return ContextParallelismSharder(shard_batch=_shard_batch_magi, local_token_global_indices=None)
+        return ContextParallelSharder(shard_batch=_shard_batch_magi)
 
-    if use_te:
+    if is_thd:
         # The THD partition is data-dependent (cu_seqlens), so shard_batch
         # installs the index map it just computed on the sharder for the token
         # verbs (chunked streams carry none). The BSHD->THD flatten is a pure
@@ -453,17 +500,17 @@ def _shard_batch_te(cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_id
                 )
             return contextlib.nullcontext, prepped, layout
 
-        return ContextParallelismSharder(shard_batch=_shard_batch_te, local_token_global_indices=None)
+        return ContextParallelSharder(shard_batch=_shard_batch_te)
 
     if cp_active:
-        return ContextParallelismSharder(
+        return ContextParallelSharder(
             shard_batch=partial(shard_batch_load_balanced, extra_seq_buffers=extra_seq_buffers),
             local_token_global_indices=round_robin_local_indices,
         )
 
     # No CP prep applies: the identity sharder, so callers hold working token
     # verbs at every cp_size.
-    return ContextParallelismSharder(
+    return ContextParallelSharder(
         shard_batch=shard_batch_identity,
         local_token_global_indices=identity_local_indices,
     )
@@ -506,13 +553,13 @@ def _make_cp_batch_and_ctx(
     seq_lens_padding_value: int = -1000,
     magi=None,
     model=None,
-    cp_sharder: Optional[ContextParallelismSharder] = None,
+    cp_sharder: Optional[ContextParallelSharder] = None,
     extra_seq_buffers: Optional[dict[str, int]] = None,
 ):
     """
-    Resolve a ContextParallelismSharder and shard the batch; a no-op when no CP prep applies.
+    Resolve a ContextParallelSharder and shard the batch; a no-op when no CP prep applies.
 
-    Every CP backend is a :class:`ContextParallelismSharder`. A model that owns its CP
+    Every CP backend is a :class:`ContextParallelSharder`. A model that owns its CP
     attention returns one from its ``prepare_model_inputs_for_cp`` hook
     (threaded here as ``cp_sharder`` — an explicit parameter, never a batch
     key, so the batch stays pure tensors); the framework constructs one for
@@ -527,7 +574,7 @@ def _make_cp_batch_and_ctx(
         batch (Dict[str, torch.Tensor]): The input batch containing (string, torch.Tensor)
 
     Returns:
-        tuple (contextmanager, dict[str, torch.Tensor], ContextParallelismSharder): The forward
+        tuple (contextmanager, dict[str, torch.Tensor], ContextParallelSharder): The forward
         context factory (nullcontext when the backend owns its transport or CP
         is inactive), the prepared/sharded batch, and the resolved sharder —
         callers use its token verbs (``shard_token_tensor`` /
@@ -536,25 +583,29 @@ def _make_cp_batch_and_ctx(
     """
 
     cp_mesh = _get_submesh(device_mesh, "cp")
-    tp_mesh = _get_submesh(device_mesh, "tp")
 
     if use_te and extra_seq_buffers:
         raise ValueError("extra_seq_buffers are not supported by the TE THD context-parallel path")
 
-    sharder = _resolve_cp_sharder(
+    strategy = _resolve_cp_sharder(
         cp_mesh,
         cp_sharder,
         magi=magi,
-        use_te=use_te,
+        is_thd=use_te,
         num_chunks=num_chunks,
         seq_lens_padding_value=seq_lens_padding_value,
         model=model,
         extra_seq_buffers=extra_seq_buffers,
     )
-    ctx, batch, layout = sharder.shard_batch(
-        cp_mesh, tp_mesh, batch, loss_mask=loss_mask, padding_token_id=padding_token_id
+    sharder = ContextParallelSharder(
+        device_mesh=device_mesh,
+        shard_batch=strategy.shard_batch,
+        local_token_global_indices=strategy.local_token_global_indices,
+        shard_layout=strategy.shard_layout,
+        loss_mask=loss_mask,
+        padding_token_id=padding_token_id,
     )
-    sharder.shard_layout = layout
+    ctx, batch = sharder.shard(batch)
     return ctx, batch, sharder
 
 
@@ -598,7 +649,7 @@ def make_cp_batch_for_te(
         return_local_indices (bool): Also return this rank's local-token global
             index map (the ``thd_get_partitioned_indices`` partition; an
             identity arange when CP is inactive; None in chunked mode, where
-            each chunk is its own token space). Used by the THD ContextParallelismSharder's
+            each chunk is its own token space). Used by the THD ContextParallelSharder's
             token verbs.
 
     Returns:
@@ -715,7 +766,7 @@ def _shard_thd_chunk_for_te(
 
     # The partition is the same for every token-aligned key; it is also this
     # rank's local-token global index map, returned so the caller can install
-    # it on the THD sharder (ContextParallelismSharder token verbs).
+    # it on the THD sharder (ContextParallelSharder token verbs).
     local_indices = tex.thd_get_partitioned_indices(
         filtered_cu_seqlens_padded, batch["input_ids"].size(0), cp_size, cp_rank
     )
diff --git a/nemo_automodel/components/distributed/parallelizer.py b/nemo_automodel/components/distributed/parallelizer.py
index eda8127f99..85686cfd9b 100644
--- a/nemo_automodel/components/distributed/parallelizer.py
+++ b/nemo_automodel/components/distributed/parallelizer.py
@@ -510,7 +510,7 @@ def parallelize(
 
             for layer in layers:
                 if hasattr(layer, "block_type") and layer.block_type == "mamba":
-                    from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+                    from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
                     mixer = layer.mixer
                     mixer.cp = MambaContextParallel(
diff --git a/nemo_automodel/components/models/deepseek_v4/cp.py b/nemo_automodel/components/models/deepseek_v4/cp.py
index f946070274..44baf63fc4 100644
--- a/nemo_automodel/components/models/deepseek_v4/cp.py
+++ b/nemo_automodel/components/models/deepseek_v4/cp.py
@@ -222,10 +222,10 @@ def build_dsv4_cp_causal_padding_mask(
 # ---------------------------------------------------------------------------
 # Model-owned CP batch sharding (Miles-style contiguous query shard).
 #
-# The CP dispatch (``cp_utils.prepare_cp_forward``) delegates manual all-gather CP to the model
-# via the ``ContextParallelismSharder`` returned by ``prepare_model_inputs_for_cp``. DSV4's
+# ``ContextParallelSharder`` construction delegates manual all-gather CP to the model
+# via the ``ContextParallelSharder`` returned by ``prepare_model_inputs_for_cp``. DSV4's
 # sharder pads + contiguously shards the sequence per CP rank (via the shared
-# contiguous implementation in ``components/distributed/cp_sharder.py``) and
+# contiguous implementation in ``components/distributed/context_parallel/sharder.py``) and
 # hands the CP process group to the forward (``_dsv4_cp_group``) so DSV4
 # attention can all-gather K/V.
 # ---------------------------------------------------------------------------
@@ -288,7 +288,7 @@ def _repad_dsv4_packed_batch(
 
     batch_size = primary.shape[0]
     # Per-row map from input position to rebuilt-row column (-1 = input pad
-    # slot whose token was dropped); the ContextParallelismSharder token verbs restore the
+    # slot whose token was dropped); the ContextParallelSharder token verbs restore the
     # caller's coordinates through it.
     input_positions = torch.full((batch_size, primary.shape[1]), -1, dtype=torch.long, device=primary.device)
     rebuilt_primary = []
@@ -433,7 +433,7 @@ def make_dsv4_contiguous_shard_cp_batch_and_ctx(
 ):
     """Contiguously shard a batch for DeepSeek V4 Miles-style context parallelism.
 
-    Exposed as ``ContextParallelismSharder.shard_batch`` (via ``functools.partial`` to bind
+    Exposed as ``ContextParallelSharder.shard_batch`` (via ``functools.partial`` to bind
     ``pad_multiple``) and invoked by the CP dispatch. HybridEP can
     first max-reduce packed lengths so every rank contributes a uniform token count.
     Each CP rank then keeps one ``seq_start:seq_end`` slice; DSV4 attention all-gathers
@@ -447,7 +447,7 @@ def make_dsv4_contiguous_shard_cp_batch_and_ctx(
     """
     import contextlib
 
-    from nemo_automodel.components.distributed.cp_sharder import (  # noqa: PLC0415
+    from nemo_automodel.components.distributed.context_parallel.sharder import (  # noqa: PLC0415
         ShardLayout,
         convert_attention_mask_to_padding_mask,
         shard_batch_contiguous,
diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py
index ff4a9ad6c6..a37e105f65 100644
--- a/nemo_automodel/components/models/deepseek_v4/model.py
+++ b/nemo_automodel/components/models/deepseek_v4/model.py
@@ -842,7 +842,7 @@ def prepare_model_inputs_for_cp(
     ) -> dict[str, Any]:
         """Model-owned context-parallel batch prep (Miles-style contiguous shard).
 
-        Returns a ``ContextParallelismSharder`` (under the ``"cp_sharder"`` batch key) so
+        Returns a ``ContextParallelSharder`` (under the ``"cp_sharder"`` batch key) so
         the CP dispatch delegates CP sharding back to this
         model, with the config-derived per-rank shard multiple bound. DSV4
         embeds internally, so (unlike VLM models) this does not pre-embed --
@@ -850,12 +850,12 @@ def prepare_model_inputs_for_cp(
         """
         from functools import partial  # noqa: PLC0415
 
-        from nemo_automodel.components.distributed.cp_sharder import (  # noqa: PLC0415
-            ContextParallelismSharder,
+        from nemo_automodel.components.distributed.context_parallel.sharder import (  # noqa: PLC0415
+            ContextParallelSharder,
             contiguous_local_indices,
         )
 
-        cp_sharder = ContextParallelismSharder(
+        cp_sharder = ContextParallelSharder(
             shard_batch=partial(
                 make_dsv4_contiguous_shard_cp_batch_and_ctx,
                 pad_multiple=dsv4_cp_local_seq_multiple(self.config),
diff --git a/nemo_automodel/components/models/gemma4_moe/cp_attention.py b/nemo_automodel/components/models/gemma4_moe/cp_attention.py
index 62aa078827..7209b16f03 100644
--- a/nemo_automodel/components/models/gemma4_moe/cp_attention.py
+++ b/nemo_automodel/components/models/gemma4_moe/cp_attention.py
@@ -1145,7 +1145,7 @@ def _gemma4_cp_manual_attention(
 ) -> torch.Tensor:
     """Gemma4-owned manual ring CP attention entry.
 
-    Plugs into cp_utils' generic ``run_cp_manual_attention`` seam: receives the
+    Plugs into context_parallel.utils' generic ``run_cp_manual_attention`` seam: receives the
     raw local (un-gathered) Q/K/V plus ``cp_mesh``, builds the ring context, and
     runs the p2p ring FlexAttention. K/V are rotated across CP ranks inside the
     ring autograd function -- they are never all-gathered.
@@ -1188,7 +1188,7 @@ def _gemma4_cp_manual_attention(
 def _install_gemma4_cp_ring_sdpa(attention_module: torch.nn.Module, cp_mesh) -> None:
     """Swap ``F.scaled_dot_product_attention`` -> Gemma4 ring CP attention on this module.
 
-    Gemma4 owns its CP attention end-to-end (it does not use cp_utils' generic CP
+    Gemma4 owns its CP attention end-to-end (it does not use context_parallel.utils' generic CP
     SDPA hooks). It installs its own ``@torch._dynamo.disable`` SDPA wrapper -- on
     the inner attention module so it also fires during gradient-checkpointing
     recompute -- that runs the p2p ring FlexAttention. The per-forward attention
@@ -1256,7 +1256,7 @@ def attach_gemma4_cp_ring_attention(attention_module: torch.nn.Module, *, use_ff
 
     Declares the metadata keys the ring needs and exposes ``setup_cp_attention(cp_mesh)``
     -- the model-owned CP-attention seam the parallelizer calls (with the CP mesh)
-    instead of cp_utils' generic SDPA hooks. ``run_cp_manual_attention`` is also bound
+    instead of context_parallel.utils' generic SDPA hooks. ``run_cp_manual_attention`` is also bound
     as the ring entry point.
 
     ``use_ffpa`` opts the (full-attention, head_dim=512) ring chunks into the FFPA
diff --git a/nemo_automodel/components/models/gemma4_moe/cp_batch.py b/nemo_automodel/components/models/gemma4_moe/cp_batch.py
index 3896ab173f..0f790897ab 100644
--- a/nemo_automodel/components/models/gemma4_moe/cp_batch.py
+++ b/nemo_automodel/components/models/gemma4_moe/cp_batch.py
@@ -21,12 +21,12 @@
 metadata inside its forward and contiguously slices them there; the dispatch-time
 sharder therefore only touches the no-grad auxiliary streams.
 
-The generic slicing lives in ``components/distributed/cp_sharder.py``; this
+The generic slicing lives in ``components/distributed/context_parallel/sharder.py``; this
 module owns the one Gemma4-specific piece the aux-only shard still needs: the
 ``_packed_seq_ids`` document-boundary synthesis its manual CP attention mask
 builder requires (its pad-region zeros depend on the global pad tail, which the
 forward -- holding only this rank's slice -- cannot reconstruct). Gemma4's
-``prepare_model_inputs_for_cp`` exposes it through the ``ContextParallelismSharder``
+``prepare_model_inputs_for_cp`` exposes it through the ``ContextParallelSharder``
 it returns under the ``"cp_sharder"`` batch key, which the CP dispatch invokes in
 place of the default load-balanced ``context_parallel`` path.
 """
@@ -35,7 +35,7 @@
 
 import torch
 
-from nemo_automodel.components.distributed.cp_sharder import (
+from nemo_automodel.components.distributed.context_parallel.sharder import (
     convert_attention_mask_to_padding_mask,
     shard_batch_contiguous,
 )
@@ -77,7 +77,7 @@ def make_contiguous_aux_only_shard_cp_batch_and_ctx(
 ):
     """Aux-only contiguous CP shard for Gemma4's sunk (in-forward) pre-embed.
 
-    Exposed as ``ContextParallelismSharder.shard_batch`` by Gemma4's sharder-only
+    Exposed as ``ContextParallelSharder.shard_batch`` by Gemma4's sharder-only
     ``prepare_model_inputs_for_cp``. It shards only the no-grad auxiliary streams
     (``labels`` / ``position_ids`` / ``loss_mask`` / ``padding_mask`` plus the
     synthesized ``_packed_seq_ids`` document map) and leaves ``input_ids`` /
diff --git a/nemo_automodel/components/models/gemma4_moe/model.py b/nemo_automodel/components/models/gemma4_moe/model.py
index 08251959eb..cfff41cb7d 100644
--- a/nemo_automodel/components/models/gemma4_moe/model.py
+++ b/nemo_automodel/components/models/gemma4_moe/model.py
@@ -78,8 +78,8 @@ def _make_missing(name: str):
     CausalLMOutputWithPast = _make_missing("CausalLMOutputWithPast")
 
 from nemo_automodel._transformers.model_capabilities import ModelCapabilities
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     contiguous_local_indices,
     shard_sequence_for_cp_contiguous,
 )
@@ -1004,7 +1004,7 @@ def setup_cp_attention(self, cp_mesh) -> None:
     def _cp_shard_batch_aux_only(self, cp_mesh, tp_mesh, batch, *, loss_mask=None, padding_token_id=0):
         """Gemma4-owned aux-only CP batch sharder that also self-installs the ring.
 
-        Exposed as ``ContextParallelismSharder.shard_batch`` by the sharder-only
+        Exposed as ``ContextParallelSharder.shard_batch`` by the sharder-only
         ``prepare_model_inputs_for_cp``. The CP dispatch calls it with the CP
         submesh, which is the one place Gemma4 reliably receives ``cp_mesh`` on a
         model-owned path (dense variants are not guaranteed to run the MoE
@@ -1436,7 +1436,7 @@ def prepare_model_inputs_for_cp(
         """Return a sharder-only CP backend; embed + splice + slice happen in forward.
 
         Sunk (Megatron-style per-microbatch) CP: the returned
-        :class:`ContextParallelismSharder` contiguously shards only the no-grad
+        :class:`ContextParallelSharder` contiguously shards only the no-grad
         aux streams (labels/position_ids/loss_mask/padding_mask + the synthesized
         ``_packed_seq_ids`` document map) via
         :func:`make_contiguous_aux_only_shard_cp_batch_and_ctx` and leaves
@@ -1457,7 +1457,7 @@ def prepare_model_inputs_for_cp(
         if batch.get("input_ids") is None:
             raise ValueError("prepare_model_inputs_for_cp requires input_ids.")
         return {
-            "cp_sharder": ContextParallelismSharder(
+            "cp_sharder": ContextParallelSharder(
                 shard_batch=self._cp_shard_batch_aux_only,
                 local_token_global_indices=contiguous_local_indices,
             )
diff --git a/nemo_automodel/components/models/glm_moe_dsa/cp.py b/nemo_automodel/components/models/glm_moe_dsa/cp.py
index cec654da0a..e27b02b6b2 100644
--- a/nemo_automodel/components/models/glm_moe_dsa/cp.py
+++ b/nemo_automodel/components/models/glm_moe_dsa/cp.py
@@ -21,7 +21,7 @@
 import torch
 import torch.distributed as dist
 
-from nemo_automodel.components.distributed.cp_sharder import ShardLayout
+from nemo_automodel.components.distributed.context_parallel.sharder import ShardLayout
 from nemo_automodel.components.distributed.thd_utils import split_batch_into_thd_chunks
 
 
@@ -176,7 +176,7 @@ def shard_glm_dsa_packed_cp_batch(
     num_chunks: int = 1,
     seq_lens_padding_value: int = -1000,
 ):
-    """``ContextParallelismSharder.shard_batch`` wrapper for GLM DSA packed CP."""
+    """``ContextParallelSharder.shard_batch`` wrapper for GLM DSA packed CP."""
     layout = _packed_cp_layout(batch, num_chunks=num_chunks)
     ctx_factory, sharded_batch = make_glm_dsa_packed_cp_batch_and_ctx(
         cp_mesh,
diff --git a/nemo_automodel/components/models/glm_moe_dsa/model.py b/nemo_automodel/components/models/glm_moe_dsa/model.py
index 08f266d26a..4b18259657 100644
--- a/nemo_automodel/components/models/glm_moe_dsa/model.py
+++ b/nemo_automodel/components/models/glm_moe_dsa/model.py
@@ -349,15 +349,15 @@ def prepare_model_inputs_for_cp(
         """
         from functools import partial  # noqa: PLC0415
 
-        from nemo_automodel.components.distributed.cp_sharder import (  # noqa: PLC0415
-            ContextParallelismSharder,
+        from nemo_automodel.components.distributed.context_parallel.sharder import (  # noqa: PLC0415
+            ContextParallelSharder,
             contiguous_local_indices,
         )
 
         if getattr(self.backend, "attn", None) != "tilelang":
             raise NotImplementedError("GLM DSA context parallelism is implemented only for backend.attn='tilelang'.")
 
-        cp_sharder = ContextParallelismSharder(
+        cp_sharder = ContextParallelSharder(
             shard_batch=partial(
                 shard_glm_dsa_packed_cp_batch,
                 num_chunks=int(num_chunks),
diff --git a/nemo_automodel/components/models/minimax_m3_vl/model.py b/nemo_automodel/components/models/minimax_m3_vl/model.py
index 6cbc5b2842..107f6a61b2 100644
--- a/nemo_automodel/components/models/minimax_m3_vl/model.py
+++ b/nemo_automodel/components/models/minimax_m3_vl/model.py
@@ -26,8 +26,8 @@
 import torch
 import torch.nn as nn
 
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     round_robin_local_indices,
     shard_batch_aux_only,
     shard_sequence_for_cp_round_robin,
@@ -598,7 +598,9 @@ def _splice_multimodal(
         # this embed+splice runs in-forward under an active CP ring context it must
         # suspend the ring dispatcher, or torch's load-balanced ring SDPA rejects
         # the non-causal attention. No-op when CP is inactive.
-        from nemo_automodel.components.distributed.cp_utils import cp_dispatcher_suspended  # noqa: PLC0415
+        from nemo_automodel.components.distributed.context_parallel.utils import (
+            cp_dispatcher_suspended,  # noqa: PLC0415
+        )
 
         with cp_dispatcher_suspended(self.cp_mesh):
             features = self.vision_tower(pixel_values, self._to_grid_list(grid_thw))
@@ -647,7 +649,7 @@ def prepare_model_inputs_for_cp(
     ) -> dict[str, Any]:
         """Return a sharder-only CP backend; embed + splice + shard happen in forward.
 
-        The returned :class:`ContextParallelismSharder` round-robin-shards only the
+        The returned :class:`ContextParallelSharder` round-robin-shards only the
         no-grad aux streams (labels/position_ids/loss_mask/padding_mask) via
         :func:`shard_batch_aux_only`, leaving ``input_ids`` and the multimodal inputs
         full-length; the forward then embeds + splices and calls
@@ -662,7 +664,7 @@ def prepare_model_inputs_for_cp(
         """
         del batch, num_chunks
         return {
-            "cp_sharder": ContextParallelismSharder(
+            "cp_sharder": ContextParallelSharder(
                 shard_batch=shard_batch_aux_only,
                 local_token_global_indices=round_robin_local_indices,
             )
diff --git a/nemo_automodel/components/models/nemotron_omni/model.py b/nemo_automodel/components/models/nemotron_omni/model.py
index 019c95ce79..d7461abfb1 100644
--- a/nemo_automodel/components/models/nemotron_omni/model.py
+++ b/nemo_automodel/components/models/nemotron_omni/model.py
@@ -34,13 +34,13 @@
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     round_robin_local_indices,
     shard_batch_aux_only,
     shard_sequence_for_cp_round_robin,
 )
-from nemo_automodel.components.distributed.cp_utils import cp_dispatcher_suspended
+from nemo_automodel.components.distributed.context_parallel.utils import cp_dispatcher_suspended
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.models.common.tie_word_embeddings import (
@@ -786,7 +786,7 @@ def prepare_model_inputs_for_cp(
         ``forward`` per microbatch (the existing ``inputs_embeds is None`` block),
         which then round-robin shards the result with
         :func:`shard_sequence_for_cp_round_robin`. The returned
-        :class:`ContextParallelismSharder` round-robin-shards only the no-grad aux
+        :class:`ContextParallelSharder` round-robin-shards only the no-grad aux
         streams (labels/position_ids/loss_mask/padding_mask) and leaves
         ``input_ids`` and the media inputs full-length for the forward. NemotronOmni
         uses plain 1-D positions, so no ``position_ids`` are computed here.
@@ -798,7 +798,7 @@ def prepare_model_inputs_for_cp(
         """
         del batch, num_chunks
         return {
-            "cp_sharder": ContextParallelismSharder(
+            "cp_sharder": ContextParallelSharder(
                 shard_batch=shard_batch_aux_only,
                 local_token_global_indices=round_robin_local_indices,
             )
diff --git a/nemo_automodel/components/models/qwen3_5/model.py b/nemo_automodel/components/models/qwen3_5/model.py
index eb6349417f..6c369d8b4f 100644
--- a/nemo_automodel/components/models/qwen3_5/model.py
+++ b/nemo_automodel/components/models/qwen3_5/model.py
@@ -38,8 +38,8 @@
     Qwen3_5Model as HFQwen3_5Model,
 )
 
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     round_robin_local_indices,
     shard_batch_aux_only,
     shard_sequence_for_cp_round_robin,
@@ -991,7 +991,7 @@ def prepare_model_inputs_for_cp(
         (a) computes the mRoPE ``position_ids`` on the *full* (unsharded) sequence
         via ``get_rope_index`` and returns them for :func:`shard_batch_aux_only`
         to round-robin-shard on the mRoPE axis, and (b) returns the
-        :class:`ContextParallelismSharder`. ``input_ids`` and the media inputs are
+        :class:`ContextParallelSharder`. ``input_ids`` and the media inputs are
         left in the batch for the forward; ``mm_token_type_ids`` is consumed here
         (only ``get_rope_index`` needs it) so the sharded forward never sees a
         full-length copy.
@@ -1047,7 +1047,7 @@ def prepare_model_inputs_for_cp(
             self.model.rope_deltas = rope_deltas
 
         return {
-            "cp_sharder": ContextParallelismSharder(
+            "cp_sharder": ContextParallelSharder(
                 shard_batch=shard_batch_aux_only,
                 local_token_global_indices=round_robin_local_indices,
             ),
@@ -1089,7 +1089,9 @@ def _embed_and_splice_for_cp(
         # splice runs in-forward under an active CP ring context it must suspend the
         # ring dispatcher, or torch's load-balanced ring SDPA all-gathers the vision
         # Q/K/V and rejects the non-causal attention. No-op when CP is inactive.
-        from nemo_automodel.components.distributed.cp_utils import cp_dispatcher_suspended  # noqa: PLC0415
+        from nemo_automodel.components.distributed.context_parallel.utils import (
+            cp_dispatcher_suspended,  # noqa: PLC0415
+        )
 
         with cp_dispatcher_suspended(self.cp_mesh):
             if pixel_values is not None:
diff --git a/nemo_automodel/components/models/qwen3_5_moe/model.py b/nemo_automodel/components/models/qwen3_5_moe/model.py
index 3212d871cc..94d15c3882 100644
--- a/nemo_automodel/components/models/qwen3_5_moe/model.py
+++ b/nemo_automodel/components/models/qwen3_5_moe/model.py
@@ -60,8 +60,8 @@ def _make_missing(name: str):
     Qwen3_5MoeVisionRotaryEmbedding = _make_missing("Qwen3_5MoeVisionRotaryEmbedding")
     HFQwen3_5MoeModel = _make_missing("Qwen3_5MoeModel")
 
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     round_robin_local_indices,
     shard_batch_aux_only,
     shard_sequence_for_cp_round_robin,
@@ -867,7 +867,7 @@ def prepare_model_inputs_for_cp(
         computes the mRoPE ``position_ids`` on the full (unsharded) sequence via
         ``get_rope_index`` and returns them for :func:`shard_batch_aux_only` to
         round-robin-shard on the mRoPE axis, plus the
-        :class:`ContextParallelismSharder`. ``input_ids`` and the media inputs are
+        :class:`ContextParallelSharder`. ``input_ids`` and the media inputs are
         left in the batch for the forward; ``mm_token_type_ids`` is consumed here
         (only ``get_rope_index`` needs it).
 
@@ -922,7 +922,7 @@ def prepare_model_inputs_for_cp(
             self.model.rope_deltas = rope_deltas
 
         return {
-            "cp_sharder": ContextParallelismSharder(
+            "cp_sharder": ContextParallelSharder(
                 shard_batch=shard_batch_aux_only,
                 local_token_global_indices=round_robin_local_indices,
             ),
@@ -963,7 +963,9 @@ def _embed_and_splice_for_cp(
         # splice runs in-forward under an active CP ring context it must suspend the
         # ring dispatcher, or torch's load-balanced ring SDPA all-gathers the vision
         # Q/K/V and rejects the non-causal attention. No-op when CP is inactive.
-        from nemo_automodel.components.distributed.cp_utils import cp_dispatcher_suspended  # noqa: PLC0415
+        from nemo_automodel.components.distributed.context_parallel.utils import (
+            cp_dispatcher_suspended,  # noqa: PLC0415
+        )
 
         with cp_dispatcher_suspended(self.cp_mesh):
             if pixel_values is not None:
diff --git a/nemo_automodel/components/models/step3p7/model.py b/nemo_automodel/components/models/step3p7/model.py
index aad1869a44..c2df2d9f8d 100644
--- a/nemo_automodel/components/models/step3p7/model.py
+++ b/nemo_automodel/components/models/step3p7/model.py
@@ -23,8 +23,8 @@
 import torch.nn as nn
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     round_robin_local_indices,
     shard_batch_aux_only,
     shard_sequence_for_cp_round_robin,
@@ -522,7 +522,7 @@ def prepare_model_inputs_for_cp(
         Embedding and the vision multimodal scatter now run inside ``forward``
         per microbatch (see the CP branch that calls ``get_multimodal_embeddings``
         + ``prepare_inputs_embeds`` + :func:`shard_sequence_for_cp_round_robin`). The returned
-        :class:`ContextParallelismSharder` round-robin-shards only the no-grad aux
+        :class:`ContextParallelSharder` round-robin-shards only the no-grad aux
         streams (labels/position_ids/loss_mask/padding_mask) and leaves
         ``input_ids`` and the media inputs full-length for the forward. Step3.7
         uses plain 1-D positions, so no ``position_ids`` are computed here (the
@@ -537,7 +537,7 @@ def prepare_model_inputs_for_cp(
             raise ValueError("Step3p7 CP pre-embedding requires input_ids.")
         del num_chunks
         return {
-            "cp_sharder": ContextParallelismSharder(
+            "cp_sharder": ContextParallelSharder(
                 shard_batch=shard_batch_aux_only,
                 local_token_global_indices=round_robin_local_indices,
             )
diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py
index 9b68951e54..bd06bf1d27 100644
--- a/nemo_automodel/components/moe/parallelizer.py
+++ b/nemo_automodel/components/moe/parallelizer.py
@@ -798,7 +798,7 @@ def apply_cp(model: torch.nn.Module, cp_mesh: DeviceMesh, cp_comm_type: str = "p
                     type(attn_module).__name__ if attn_module is not None else type(self_attn).__name__,
                 )
         elif layer_type == "mamba":
-            from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+            from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
             mixer = block.self_attn  # NemotronV3Block.self_attn aliases mixer
             mixer.cp = MambaContextParallel(
diff --git a/nemo_automodel/components/speculative/dflash/target.py b/nemo_automodel/components/speculative/dflash/target.py
index b6dfd6203c..9d53259a72 100644
--- a/nemo_automodel/components/speculative/dflash/target.py
+++ b/nemo_automodel/components/speculative/dflash/target.py
@@ -81,7 +81,7 @@ def __init__(
         self.cp_mesh = cp_mesh
         self._cp_size = cp_mesh.size() if cp_mesh is not None else 1
         if self._cp_size > 1:
-            from nemo_automodel.components.distributed.cp_utils import attach_context_parallel_hooks
+            from nemo_automodel.components.distributed.context_parallel.utils import attach_context_parallel_hooks
             from nemo_automodel.components.speculative.target_cp import attach_cp_kv_gather_hooks
 
             # Strip the 4D mask, and all-gather K/V so each rank attends its local Q
diff --git a/nemo_automodel/components/speculative/dspark/target.py b/nemo_automodel/components/speculative/dspark/target.py
index d0ae712f67..3ef97bccab 100644
--- a/nemo_automodel/components/speculative/dspark/target.py
+++ b/nemo_automodel/components/speculative/dspark/target.py
@@ -74,7 +74,7 @@ def __init__(self, model: nn.Module, target_layer_ids: Sequence[int], cp_mesh=No
         self.cp_mesh = cp_mesh
         self._cp_size = cp_mesh.size() if cp_mesh is not None else 1
         if self._cp_size > 1:
-            from nemo_automodel.components.distributed.cp_utils import attach_context_parallel_hooks
+            from nemo_automodel.components.distributed.context_parallel.utils import attach_context_parallel_hooks
             from nemo_automodel.components.speculative.target_cp import attach_cp_kv_gather_hooks
 
             # Strip the 4D mask (self_attn then calls SDPA on the local shard), and
diff --git a/nemo_automodel/components/speculative/eagle/target.py b/nemo_automodel/components/speculative/eagle/target.py
index 165e0e4ed2..f9db7b9a45 100644
--- a/nemo_automodel/components/speculative/eagle/target.py
+++ b/nemo_automodel/components/speculative/eagle/target.py
@@ -168,7 +168,7 @@ def __init__(self, model: nn.Module, aux_layer_ids: Sequence[int] | None = None,
         self.cp_mesh = cp_mesh
         self._cp_size = cp_mesh.size() if cp_mesh is not None else 1
         if self._cp_size > 1:
-            from nemo_automodel.components.distributed.cp_utils import attach_context_parallel_hooks
+            from nemo_automodel.components.distributed.context_parallel.utils import attach_context_parallel_hooks
             from nemo_automodel.components.speculative.target_cp import attach_cp_kv_gather_hooks
 
             attach_context_parallel_hooks(self.model)
diff --git a/nemo_automodel/components/speculative/target_cp.py b/nemo_automodel/components/speculative/target_cp.py
index ba3ea80882..cf9369a165 100644
--- a/nemo_automodel/components/speculative/target_cp.py
+++ b/nemo_automodel/components/speculative/target_cp.py
@@ -19,7 +19,7 @@
 full sequence before handing them to the draft. That flow is specific to
 speculative decoding -- it needs no gradients through the target and no
 model-owned CP sharder -- so it lives here rather than in the shared
-``components/distributed/cp_utils`` surface.
+``components/distributed/context_parallel/utils`` surface.
 """
 
 from typing import Callable, List, Optional
diff --git a/nemo_automodel/recipes/dllm/strategy.py b/nemo_automodel/recipes/dllm/strategy.py
index 708697be70..6730ea7c0f 100644
--- a/nemo_automodel/recipes/dllm/strategy.py
+++ b/nemo_automodel/recipes/dllm/strategy.py
@@ -42,7 +42,7 @@
     corrupt_uniform,
     corrupt_uniform_random,
 )
-from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
 from nemo_automodel.components.distributed.utils import get_sync_ctx
 from nemo_automodel.components.loss.dllm_loss import (
     BlockDiffusionCrossEntropyLoss,
@@ -701,7 +701,8 @@ def forward_backward(
             torch.autocast(device_type="cuda", dtype=autocast_dtype) if autocast_dtype is not None else nullcontext()
         )
         fp8_ctx = recipe.te_fp8.maybe_te_autocast() if recipe.te_fp8 is not None else nullcontext()
-        train_ctx, _, _ = prepare_cp_forward(None, recipe.device_mesh, {})
+        cp_sharder = ContextParallelSharder(None, recipe.device_mesh, {})
+        train_ctx, _ = cp_sharder.shard({})
 
         with train_ctx(), sync_ctx, fp8_ctx, autocast_ctx:
             draft_kwargs = dict(
diff --git a/nemo_automodel/recipes/dllm/train_ft.py b/nemo_automodel/recipes/dllm/train_ft.py
index e20084daa6..f854865b93 100644
--- a/nemo_automodel/recipes/dllm/train_ft.py
+++ b/nemo_automodel/recipes/dllm/train_ft.py
@@ -46,7 +46,7 @@
 
 from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
 from nemo_automodel.components.datasets.dllm.collate import DLLMCollator
-from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
 from nemo_automodel.components.distributed.utils import get_sync_ctx
 from nemo_automodel.components.loggers.metric_logger import MetricsSample
 from nemo_automodel.components.loggers.mlflow_utils import to_float_metrics
@@ -309,7 +309,8 @@ def _forward_backward_step(
         model = self.model_parts[0]
 
         # Context parallel setup (no labels to pass for dLLM)
-        train_ctx, batch, _ = prepare_cp_forward(None, self.device_mesh, batch)
+        cp_sharder = ContextParallelSharder(None, self.device_mesh, batch)
+        train_ctx, batch = cp_sharder.shard(batch)
         fp8_ctx = self.te_fp8.maybe_te_autocast() if self.te_fp8 is not None else nullcontext()
         sync_ctx = (
             get_sync_ctx(
@@ -1083,7 +1084,8 @@ def _forward_backward_step(
         p_mask = window["p_mask"]
 
         model = self.model_parts[0]
-        train_ctx, batch, _ = prepare_cp_forward(None, self.device_mesh, batch)
+        cp_sharder = ContextParallelSharder(None, self.device_mesh, batch)
+        train_ctx, batch = cp_sharder.shard(batch)
         fp8_ctx = self.te_fp8.maybe_te_autocast() if self.te_fp8 is not None else nullcontext()
         sync_ctx = (
             get_sync_ctx(
diff --git a/nemo_automodel/recipes/kd_utils.py b/nemo_automodel/recipes/kd_utils.py
index d93bff0dfe..3f065e13ce 100644
--- a/nemo_automodel/recipes/kd_utils.py
+++ b/nemo_automodel/recipes/kd_utils.py
@@ -25,7 +25,7 @@
 from torch.distributed.tensor import DTensor, Shard
 
 from nemo_automodel.components.distributed.config import DDPConfig, DistributedSetup
-from nemo_automodel.components.distributed.cp_utils import unshard_context_parallel_tensor
+from nemo_automodel.components.distributed.context_parallel.utils import unshard_context_parallel_tensor
 from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config, parse_distributed_section
 
 if TYPE_CHECKING:
diff --git a/nemo_automodel/recipes/llm/kd.py b/nemo_automodel/recipes/llm/kd.py
index cb57fa47a2..1b51403078 100644
--- a/nemo_automodel/recipes/llm/kd.py
+++ b/nemo_automodel/recipes/llm/kd.py
@@ -51,7 +51,7 @@
 from nemo_automodel._transformers.auto_tokenizer import NeMoAutoTokenizer
 from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
 from nemo_automodel.components.distributed.config import DistributedSetup
-from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
 from nemo_automodel.components.distributed.pipelining.config import PipelineConfig
 from nemo_automodel.components.distributed.utils import get_sync_ctx
 from nemo_automodel.components.loggers.metric_logger import MetricsSample
@@ -77,9 +77,6 @@
 )
 from nemo_automodel.recipes.llm.train_ft import (
     TrainFinetuneRecipeForNextTokenPrediction,
-    _get_num_thd_chunks,
-    _uses_te_dot_product_attention,
-    _uses_thd_collater,
     build_model,
 )
 
@@ -493,12 +490,12 @@ def _teacher_forward_separate(self, batch: dict[str, Any]) -> torch.Tensor | Non
         """
         batch = self.kd_mesh_bridge.move_to_device(batch)
         sequence_length = batch["labels"].shape[1]
-        train_ctx, batch, _ = prepare_cp_forward(
+        cp_sharder = ContextParallelSharder(
             self.teacher_model,
             self.device_mesh,
             batch,
-            use_te=False,
         )
+        train_ctx, batch = cp_sharder.shard(batch)
         labels = batch.pop("labels")
         with train_ctx(), torch.no_grad():
             if self.pp_enabled:
@@ -597,7 +594,7 @@ def _forward_backward_step(
         labels = batch.pop("labels")
         # KD has not wired model-owned CP; skip the pre-embed hook explicitly.
         # Separate-mesh teacher logits ride the batch through CP sharding.
-        train_ctx, batch, _ = prepare_cp_forward(
+        cp_sharder = ContextParallelSharder(
             self.model_parts[0],
             self.device_mesh,
             batch,
@@ -605,6 +602,7 @@ def _forward_backward_step(
             invoke_pre_embed=False,
             extra_seq_buffers={"teacher_logits": 1} if separate_teacher_logits is not None else None,
         )
+        train_ctx, batch = cp_sharder.shard(batch)
         separate_teacher_logits = batch.pop("teacher_logits", None)
 
         model = self.model_parts[0]
@@ -712,16 +710,16 @@ def _forward_backward_step_pp(
         if separate_teacher_logits is not None:
             batch["teacher_logits"] = separate_teacher_logits
         # KD has not wired model-owned CP; skip the pre-embed hook explicitly.
-        train_ctx, batch, _ = prepare_cp_forward(
+        cp_sharder = ContextParallelSharder(
             self.model_parts[0],
             self.device_mesh,
             batch,
-            use_te=_uses_te_dot_product_attention(self.cfg.model) and _uses_thd_collater(self.cfg.dataloader),
             padding_token_id=self.tokenizer.pad_token_id if self.tokenizer else 0,
-            num_chunks=_get_num_thd_chunks(True, self.cfg),
+            num_chunks=self.pp.pp_batch_size // self.pp.pp_microbatch_size,
             invoke_pre_embed=False,
             extra_seq_buffers={"teacher_logits": 1} if separate_teacher_logits is not None else None,
         )
+        train_ctx, batch = cp_sharder.shard(batch)
         separate_teacher_logits = batch.pop("teacher_logits", None)
         labels = batch.pop("labels")
         input_ids = batch.pop("input_ids")
diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py
index 9d5dd6993a..77140f772c 100644
--- a/nemo_automodel/recipes/llm/train_ft.py
+++ b/nemo_automodel/recipes/llm/train_ft.py
@@ -51,10 +51,11 @@
 from nemo_automodel._transformers.mfu import AutoMFU
 from nemo_automodel._transformers.utils import apply_cache_compatibility_patches
 from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
+from nemo_automodel.components.datasets.loader import DataloaderConfig
 from nemo_automodel.components.distributed.config import DistributedSetup, FSDP2Config, MegatronFSDPConfig
-from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
+from nemo_automodel.components.distributed.context_parallel.magi import MagiState, setup_magi
 from nemo_automodel.components.distributed.init_utils import initialize_distributed
-from nemo_automodel.components.distributed.magi_attn_utils import MagiState, setup_magi
 from nemo_automodel.components.distributed.mesh import MeshContext
 from nemo_automodel.components.distributed.pipelining import AutoPipeline
 from nemo_automodel.components.distributed.utils import FirstRankPerNode, dp_eval_sample_shard, get_sync_ctx
@@ -117,53 +118,31 @@ def _get_model_name(cfg_model):
         return None
 
 
-def _uses_te_dot_product_attention(model_or_cfg):
-    """Check whether the model uses TE DotProductAttention.
-
-    Accepts either an instantiated nn.Module (preferred — inspects actual modules)
-    or a config object (fallback — checks backend.attn string).
-    """
-    if isinstance(model_or_cfg, torch.nn.Module):
-        try:
-            from transformer_engine.pytorch.attention import DotProductAttention
-        except ImportError:
-            return False
-        return any(isinstance(m, DotProductAttention) for m in model_or_cfg.modules())
-    # Config fallback for call sites before model is built
-    return (
-        hasattr(model_or_cfg, "backend") and hasattr(model_or_cfg.backend, "attn") and model_or_cfg.backend.attn == "te"
-    )
-
-
-def _uses_thd_collater(cfg_dataloader):
-    """Return True if the dataloader's collate_fn is ``packed_sequence_thd_collater``.
-
-    ``collate_fn`` ends in ``_fn``, so ConfigNode resolves the YAML dotted-path string to
-    the actual callable at load time — the value here is always the function, never a string.
-    """
-    from nemo_automodel.components.datasets.utils import packed_sequence_thd_collater
-
-    return getattr(cfg_dataloader, "collate_fn", None) is packed_sequence_thd_collater
-
-
-def _should_pack_validation(cfg: RecipeConfig, model: nn.Module) -> bool:
+def _should_pack_validation(
+    training_dataloader: DataloaderConfig | None,
+    validation_dataloader: DataloaderConfig,
+    model: nn.Module,
+) -> bool:
     """Return whether validation must use the configured training packer."""
-    if cfg.get("packed_sequence.packed_sequence_size", 0) <= 0:
+    if validation_dataloader.packing is None:
         return False
-
-    validation_uses_thd = _uses_thd_collater(cfg.get("validation_dataloader", None))
-    if validation_uses_thd:
+    if replace(validation_dataloader, packing=None).emits_thd:
         return True
-
+    if training_dataloader is None or not training_dataloader.emits_thd:
+        return False
     model_requires_packing = bool(
         callable(getattr(model, "should_pack_validation_with_training", None))
         and model.should_pack_validation_with_training()
     )
-    magi_backend = (
-        str(cfg.get("model.backend.attn", "")) == "magi" or str(cfg.get("model.attn_implementation", "")) == "magi"
+    model_config = getattr(model, "config", None)
+    attention_backend = getattr(getattr(model, "backend", None), "attn", None) or getattr(
+        model_config, "_attn_implementation", None
+    )
+    return (
+        attention_backend in ("te", "magi")
+        or bool(getattr(model, "_te_attention_injected", False))
+        or model_requires_packing
     )
-    backend_requires_packing = _uses_te_dot_product_attention(cfg.model) or magi_backend or model_requires_packing
-    return backend_requires_packing and _uses_thd_collater(cfg.get("dataloader", None))
 
 
 def _should_precompute_pp_causal_masks(model_config: Any) -> bool:
@@ -171,12 +150,6 @@ def _should_precompute_pp_causal_masks(model_config: Any) -> bool:
     return getattr(model_config, "model_type", None) != "deepseek_v4"
 
 
-def _get_num_thd_chunks(pp_enabled, cfg):
-    if pp_enabled:
-        return cfg.get("step_scheduler.local_batch_size", 1) // cfg.get("distributed.pipeline.pp_microbatch_size", 1)
-    return 1
-
-
 def _maybe_downgrade_loss_fn(loss_fn: nn.Module, probe_module: nn.Module, pp_enabled: bool) -> nn.Module:
     """Downgrade to MaskedCrossEntropy when the requested loss cannot run."""
     if not _supports_logits_to_keep(probe_module) and not isinstance(loss_fn, MaskedCrossEntropy):
@@ -537,8 +510,9 @@ def setup(self):
             # THD override logic
             if (
                 self.mesh_context.cp_size > 1
-                and _uses_te_dot_product_attention(self.cfg.model)
-                and _uses_thd_collater(self.cfg.get("dataloader", None))
+                and self.cfg.get("model.backend.attn", self.cfg.get("model.attn_implementation", None)) == "te"
+                and self.cfg.dataloader is not None
+                and self.cfg.dataloader.emits_thd
             ):
                 pp_microbatch_size = 1
                 pp_batch_size = pp_batch_size // self.cfg.get("distributed.pipeline.pp_microbatch_size", 1)
@@ -700,9 +674,12 @@ def materialize_loader(config):
                 )
 
         self.dataloader = materialize_loader(self.cfg.dataloader)
-        pack_validation = _should_pack_validation(self.cfg, self.model_parts[0])
         self.val_dataloaders = {
-            name: materialize_loader(dl_config if pack_validation else replace(dl_config, packing=None))
+            name: materialize_loader(
+                dl_config
+                if _should_pack_validation(self.cfg.dataloader, dl_config, self.model_parts[0])
+                else replace(dl_config, packing=None)
+            )
             for name, dl_config in self.cfg.validation_dataloaders.items()
         }
         # Optional tool-call accuracy evaluator for agent SFT runs.
@@ -972,24 +949,14 @@ def _forward_backward_step(
             )
             for k, v in batch.items()
         }
-        _thd_collater = _uses_thd_collater(self.cfg.get("dataloader", None))
-        # Gate THD/cu_seqlens processing on the dataset being THD-packed, not on TE
-        # attention being present on this rank: both TE attention and mamba need
-        # cu_seqlens, and gating on attention would drop PP stages with no attention
-        # layers (mamba+moe only) and leave cu_seqlens unbuilt downstream.
-        _use_te_value = _thd_collater
-        _num_chunks_value = _get_num_thd_chunks(self.pp_enabled, self.cfg)
-        # Single CP dispatch: magi / model-owned (ContextParallelismSharder) / TE-THD / generic
-        # torch context_parallel.
-        train_ctx, batch, _ = prepare_cp_forward(
+        cp_sharder = ContextParallelSharder(
             self.model_parts[0] if hasattr(self, "model_parts") else None,
             self.device_mesh,
             batch,
-            magi=self.magi,
-            use_te=_use_te_value,
             padding_token_id=self.tokenizer.pad_token_id if self.tokenizer else 0,
-            num_chunks=_num_chunks_value,
+            num_chunks=self.pp.pp_batch_size // self.pp.pp_microbatch_size if self.pp_enabled else 1,
         )
+        train_ctx, batch = cp_sharder.shard(batch)
         labels = batch.pop("labels")
         fp8_ctx = self.te_fp8.maybe_te_autocast() if self.te_fp8 is not None else nullcontext()
 
diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py
index fb03f5a6fc..6e7b46da5f 100644
--- a/nemo_automodel/recipes/vlm/finetune.py
+++ b/nemo_automodel/recipes/vlm/finetune.py
@@ -48,9 +48,9 @@
 from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
 from nemo_automodel.components.datasets.vlm.pp_media import stage_vlm_media_for_pp
 from nemo_automodel.components.distributed.config import DistributedSetup, MegatronFSDPConfig
-from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
+from nemo_automodel.components.distributed.context_parallel.magi import MagiState, setup_magi
 from nemo_automodel.components.distributed.init_utils import initialize_distributed
-from nemo_automodel.components.distributed.magi_attn_utils import MagiState, setup_magi
 from nemo_automodel.components.distributed.pipelining import AutoPipeline
 from nemo_automodel.components.distributed.utils import FirstRankPerNode, get_sync_ctx
 from nemo_automodel.components.loggers.log_utils import setup_logging
@@ -816,15 +816,14 @@ def _forward_backward_step(
                 "context-parallel THD for mRoPE VLMs is not yet implemented."
             )
         _padding_id = getattr(getattr(getattr(self, "processor", None), "tokenizer", None), "pad_token_id", 0) or 0
-        train_ctx, batch, _ = prepare_cp_forward(
+        cp_sharder = ContextParallelSharder(
             self.model_parts[0],
             self.device_mesh,
             batch,
-            magi=self.magi,
-            use_te=_use_te_vlm,
             padding_token_id=_padding_id,
             invoke_pre_embed=True,
         )
+        train_ctx, batch = cp_sharder.shard(batch)
         labels = batch.pop("labels")
 
         if self.pp_enabled:
@@ -1111,12 +1110,13 @@ def _run_validation_epoch(self, val_dataloader):
                 }
                 num_label_tokens = (batch["labels"] != -100).sum().item()
 
-                train_ctx, batch, _ = prepare_cp_forward(
+                cp_sharder = ContextParallelSharder(
                     self.model_parts[0],
                     self.device_mesh,
                     batch,
                     invoke_pre_embed=not self.pp_enabled,
                 )
+                train_ctx, batch = cp_sharder.shard(batch)
                 labels = batch.pop("labels")
                 with train_ctx():
                     batch = filter_forward_kwargs(self.model_parts[0], batch)
diff --git a/nemo_automodel/recipes/vlm/kd.py b/nemo_automodel/recipes/vlm/kd.py
index 5c0e33e51f..b173899baa 100644
--- a/nemo_automodel/recipes/vlm/kd.py
+++ b/nemo_automodel/recipes/vlm/kd.py
@@ -52,7 +52,7 @@
 from nemo_automodel._transformers.auto_tokenizer import NeMoAutoTokenizer
 from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
 from nemo_automodel.components.distributed.config import DistributedSetup
-from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
 from nemo_automodel.components.distributed.utils import get_sync_ctx
 from nemo_automodel.components.loggers.metric_logger import MetricsSample
 from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy
@@ -264,7 +264,8 @@ def _teacher_forward_separate(self, batch: dict[str, Any]) -> torch.Tensor:
         model = self.teacher_model
         # Single CP dispatch: invokes the teacher's pre-embed hook (when CP is
         # active and the model has one) and shards the batch.
-        train_ctx, batch, _ = prepare_cp_forward(model, self.device_mesh, batch)
+        cp_sharder = ContextParallelSharder(model, self.device_mesh, batch)
+        train_ctx, batch = cp_sharder.shard(batch)
         batch.pop("labels")
         with train_ctx(), torch.no_grad():
             teacher_batch = filter_forward_kwargs(model, batch)
@@ -328,13 +329,14 @@ def _forward_backward_step(
             batch["teacher_logits"] = separate_teacher_logits
 
         # Separate-mesh teacher logits ride the batch through CP sharding.
-        train_ctx, batch, _ = prepare_cp_forward(
+        cp_sharder = ContextParallelSharder(
             self.model_parts[0],
             self.device_mesh,
             batch,
             invoke_pre_embed=not self.pp_enabled,
             extra_seq_buffers={"teacher_logits": 1} if separate_teacher_logits is not None else None,
         )
+        train_ctx, batch = cp_sharder.shard(batch)
         separate_teacher_logits = batch.pop("teacher_logits", None)
         labels = batch.pop("labels")
 
diff --git a/skills/nemo-automodel-distributed-training/SKILL.md b/skills/nemo-automodel-distributed-training/SKILL.md
index 7b95db3f69..29861167b1 100644
--- a/skills/nemo-automodel-distributed-training/SKILL.md
+++ b/skills/nemo-automodel-distributed-training/SKILL.md
@@ -528,7 +528,7 @@ components/distributed/pipelining/hf_utils.py      -- HF model validation for PP
 Context parallelism:
 
 ```
-components/distributed/cp_utils.py
+components/distributed/context_parallel/utils.py
     make_cp_batch_and_ctx()            -- creates CP context manager + shards batch
     create_context_parallel_ctx()      -- wraps torch.distributed.tensor.experimental.context_parallel
     attach_context_parallel_hooks()    -- strips attention_mask, sets is_causal=True
diff --git a/tests/functional_tests/attention/prefix_tree_flex_parity.py b/tests/functional_tests/attention/prefix_tree_flex_parity.py
index e5f43f0cfe..f0da17caea 100644
--- a/tests/functional_tests/attention/prefix_tree_flex_parity.py
+++ b/tests/functional_tests/attention/prefix_tree_flex_parity.py
@@ -40,7 +40,7 @@
 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 
 from nemo_automodel.components.datasets.llm.prefix_tree import fold_shared_prefix_rollouts
-from nemo_automodel.components.distributed.magi_attn_utils import AttnMaskSpec
+from nemo_automodel.components.distributed.context_parallel.magi import AttnMaskSpec
 
 # fp32 both sides: a tight tolerance that flags any mask error rather than noise.
 MAX_DIFF_TOL = 1e-3
diff --git a/tests/functional_tests/attention/prefix_tree_magi_parity.py b/tests/functional_tests/attention/prefix_tree_magi_parity.py
index 1572ee2d57..b69a7e1263 100644
--- a/tests/functional_tests/attention/prefix_tree_magi_parity.py
+++ b/tests/functional_tests/attention/prefix_tree_magi_parity.py
@@ -38,7 +38,7 @@
 from _prefix_tree_reference import build_reference_mask
 
 from nemo_automodel.components.datasets.llm.prefix_tree import fold_shared_prefix_rollouts
-from nemo_automodel.components.distributed.magi_attn_utils import (
+from nemo_automodel.components.distributed.context_parallel.magi import (
     AttnMaskSpec,
     is_magi_available,
     make_magi_attn_func,
diff --git a/tests/functional_tests/context_parallel/run_attention_cp.py b/tests/functional_tests/context_parallel/run_attention_cp.py
index c2daf92b58..c901c0ed77 100644
--- a/tests/functional_tests/context_parallel/run_attention_cp.py
+++ b/tests/functional_tests/context_parallel/run_attention_cp.py
@@ -513,7 +513,7 @@ def _run_thd_te_qwen_deepseek(model_type, config, rank, world_size, device, attn
         dist.broadcast(param_with_cp.data, src=0)
 
     # Create packed sequence batch
-    from nemo_automodel.components.distributed.cp_utils import make_cp_batch_for_te
+    from nemo_automodel.components.distributed.context_parallel.utils import make_cp_batch_for_te
 
     batch_size = 4
     seq_lens_per_batch = [[32], [40], [36], [44]]
diff --git a/tests/functional_tests/context_parallel/run_cp_dispatcher_suspend.py b/tests/functional_tests/context_parallel/run_cp_dispatcher_suspend.py
index 4f8fe1ca7f..f517ff46af 100644
--- a/tests/functional_tests/context_parallel/run_cp_dispatcher_suspend.py
+++ b/tests/functional_tests/context_parallel/run_cp_dispatcher_suspend.py
@@ -41,8 +41,8 @@ def main():
 
     from torch.distributed.device_mesh import init_device_mesh
 
-    from nemo_automodel.components.distributed.cp_sharder import shard_batch_aux_only
-    from nemo_automodel.components.distributed.cp_utils import cp_dispatcher_suspended
+    from nemo_automodel.components.distributed.context_parallel.sharder import shard_batch_aux_only
+    from nemo_automodel.components.distributed.context_parallel.utils import cp_dispatcher_suspended
 
     mesh = init_device_mesh("cuda", (dist.get_world_size(),), mesh_dim_names=("cp",))
     cp_mesh = mesh["cp"]
diff --git a/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py b/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py
index a1265f6564..b63bc805a0 100644
--- a/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py
+++ b/tests/functional_tests/context_parallel/run_cp_pp_image_sink.py
@@ -135,7 +135,7 @@ def main():
     from torch.distributed.device_mesh import init_device_mesh
 
     from nemo_automodel.components.datasets.vlm.pp_media import prepare_vlm_media_for_pp, stage_vlm_media_for_pp
-    from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+    from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
     from nemo_automodel.components.distributed.pipelining import AutoPipeline
     from nemo_automodel.components.moe.parallelizer import apply_cp
 
@@ -191,7 +191,8 @@ def cp_only(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name=None, **kw):
             batch = prepare_vlm_media_for_pp(batch, batch_size=2, n_microbatches=2)
         else:
             batch.update({"pixel_values": pv, "image_grid_thw": grid})
-        train_ctx, batch, _ = prepare_cp_forward(model_part0, mesh, batch)
+        cp_sharder = ContextParallelSharder(model_part0, mesh, batch)
+        train_ctx, batch = cp_sharder.shard(batch)
         labels = batch.pop("labels")
         if pp_size > 1:
             with train_ctx(), stage_vlm_media_for_pp(pp, pp.parts, batch):
diff --git a/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py b/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py
index 4eca05c01a..063bfdca01 100644
--- a/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py
+++ b/tests/functional_tests/context_parallel/run_cp_pp_layer2_sink.py
@@ -105,7 +105,7 @@ def main():
 
     from torch.distributed.device_mesh import init_device_mesh
 
-    from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+    from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
     from nemo_automodel.components.distributed.pipelining import AutoPipeline
     from nemo_automodel.components.moe.parallelizer import apply_cp
 
@@ -162,7 +162,8 @@ def cp_only_parallelize(m, world_mesh, moe_mesh, *, dp_axis_names, cp_axis_name=
         dist.broadcast(input_ids, src=0)
         pos = torch.arange(seqlen, device=device).unsqueeze(0).expand(2, -1).contiguous()
         batch = {"input_ids": input_ids.clone(), "labels": input_ids.clone(), "position_ids": pos.clone()}
-        train_ctx, batch, _ = prepare_cp_forward(model_part0, mesh, batch)
+        cp_sharder = ContextParallelSharder(model_part0, mesh, batch)
+        train_ctx, batch = cp_sharder.shard(batch)
         labels = batch.pop("labels")
         if pp_size > 1:
             with train_ctx():
diff --git a/tests/functional_tests/context_parallel/run_cp_sharder_token_verbs.py b/tests/functional_tests/context_parallel/run_cp_sharder_token_verbs.py
index 7f4ddc92fd..0856419d6e 100644
--- a/tests/functional_tests/context_parallel/run_cp_sharder_token_verbs.py
+++ b/tests/functional_tests/context_parallel/run_cp_sharder_token_verbs.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Multi-rank functional test for the ContextParallelismSharder token verbs (L1, 2+ GPUs).
+"""Multi-rank functional test for the ContextParallelSharder token verbs (L1, 2+ GPUs).
 
 The single-process unit suite can only exercise the identity early-returns of
 ``gather_token_tensor``; this driver runs the real collectives:
@@ -34,8 +34,10 @@
 import torch.distributed as dist
 from torch.distributed.device_mesh import init_device_mesh
 
-from nemo_automodel.components.distributed.cp_sharder import round_robin_local_indices
-from nemo_automodel.components.distributed.cp_utils import _resolve_cp_sharder
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
+    round_robin_local_indices,
+)
 
 
 def main() -> None:
@@ -46,10 +48,6 @@ def main() -> None:
     mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("cp",))
     cp_mesh = mesh["cp"]
 
-    sharder = _resolve_cp_sharder(
-        cp_mesh, None, magi=None, use_te=False, num_chunks=1, seq_lens_padding_value=-1000, model=None
-    )
-
     # Shard a batch whose length needs CP padding (6 -> 8 at cp=2) so the
     # captured layout (original=6, padded=8) drive fill/trim below.
     seq_len = 3 * world
@@ -57,22 +55,22 @@ def main() -> None:
         "input_ids": torch.arange(seq_len, device=device).unsqueeze(0),
         "labels": torch.arange(seq_len, device=device).unsqueeze(0),
     }
-    _, _, layout = sharder.shard_batch(cp_mesh, None, batch)
-    sharder.shard_layout = layout
+    sharder = ContextParallelSharder(None, mesh, batch)
+    sharder.shard(batch)
     padded = sharder.shard_layout.padded_seq_len
     assert sharder.shard_layout.original_seq_len == seq_len, sharder.shard_layout.original_seq_len
     assert padded == seq_len + (-seq_len) % (2 * world), padded
 
     # --- down: caller-coordinate tensor rides the same layout -------------
     full = torch.arange(float(seq_len), device=device).unsqueeze(0)
-    local = sharder.shard_token_tensor(cp_mesh, full, fill=-1.0)
+    local = sharder.shard_token_tensor(full, fill=-1.0)
     indices = round_robin_local_indices(cp_mesh, padded, device=device)
     expected_local = torch.where(indices < seq_len, indices.float(), torch.tensor(-1.0, device=device)).unsqueeze(0)
     assert torch.equal(local, expected_local), (rank, local.tolist(), expected_local.tolist())
 
     # --- up: differentiable gather + trim back to caller coordinates ------
     local_leaf = local.detach().clone().requires_grad_(True)
-    gathered = sharder.gather_token_tensor(cp_mesh, local_leaf, trim=True)
+    gathered = sharder.gather_token_tensor(local_leaf, trim=True)
     assert gathered.shape == (1, seq_len), gathered.shape
     # global order restored: position i holds the token with global index i
     assert torch.equal(gathered, full), (rank, gathered.tolist())
diff --git a/tests/functional_tests/context_parallel/run_gemma4_vl_cp_sink.py b/tests/functional_tests/context_parallel/run_gemma4_vl_cp_sink.py
index 4a39b173b7..8eaff3e815 100644
--- a/tests/functional_tests/context_parallel/run_gemma4_vl_cp_sink.py
+++ b/tests/functional_tests/context_parallel/run_gemma4_vl_cp_sink.py
@@ -15,7 +15,7 @@
 """GPU CP forward-equivalence for the Gemma4 in-forward pre-embed shard (L1, 2 GPUs).
 
 Exercises the sunk contiguous-CP path end to end for the E-series-shaped dense
-Gemma4: ``prepare_cp_forward`` invokes the sharder-only
+Gemma4: ``ContextParallelSharder`` invokes the sharder-only
 ``prepare_model_inputs_for_cp`` hook (``shard_batch_contiguous(shard_primary=False)``
 shards labels/position_ids and the synthesized ``_packed_seq_ids``; the model records
 ``cp_mesh`` and installs its p2p flex ring), and
@@ -24,7 +24,7 @@
 slices this rank's shard. The unsharded cp2 logits must match the cp1 eager
 forward.
 
-  cp1 eager (cp_mesh unset)  ==  cp2 (apply_cp + prepare_cp_forward + in-forward slice)
+  cp1 eager (cp_mesh unset)  ==  cp2 (apply_cp + ContextParallelSharder + in-forward slice)
 
 vision-bidirectional mask is driven by ``mm_token_type_ids`` (a vision block) so
 the ``_gemma4_vision_group_ids`` cumsum-then-slice path is exercised without
@@ -58,7 +58,7 @@ def main():
 
     from torch.distributed.device_mesh import init_device_mesh
 
-    from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+    from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
     from nemo_automodel.components.models.common import BackendConfig
     from nemo_automodel.components.models.gemma4_moe.model import (
         Gemma4Config,
@@ -146,7 +146,7 @@ def main():
     if rank == 0:
         print("[progress] cp1 eager forward done", flush=True)
 
-    # cp2: the recipe path -- apply_cp installs the ring, prepare_cp_forward runs the
+    # cp2: the recipe path -- apply_cp installs the ring, ContextParallelSharder runs the
     # sharder-only hook (aux-only contiguous shard), forward embeds + contiguously
     # slices this rank's shard.
     device_mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("cp",))
@@ -158,7 +158,8 @@ def main():
     }
     if mm_arg is not None:
         batch["mm_token_type_ids"] = mm_arg.clone()
-    train_ctx, batch, _ = prepare_cp_forward(model, device_mesh, batch)
+    cp_sharder = ContextParallelSharder(model, device_mesh, batch)
+    train_ctx, batch = cp_sharder.shard(batch)
     batch.pop("labels", None)
     # Drop the synthesized single-document _packed_seq_ids: it reaches the ring by
     # the identical contiguous slice in old and new (unit-proven slice-equivalence),
diff --git a/tests/functional_tests/context_parallel/run_hybrid_nemotron_v3_cp.py b/tests/functional_tests/context_parallel/run_hybrid_nemotron_v3_cp.py
index 7afab519c2..a052468c46 100644
--- a/tests/functional_tests/context_parallel/run_hybrid_nemotron_v3_cp.py
+++ b/tests/functional_tests/context_parallel/run_hybrid_nemotron_v3_cp.py
@@ -140,7 +140,7 @@ def _wire_te_cp(model, cp_group, config):
     """Wire TE-based CP on each hybrid layer (p2p for attention, hidden-parallel for mamba)."""
     from transformer_engine.pytorch.attention import DotProductAttention
 
-    from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+    from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
     for layer in model.layers.values():
         if layer.block_type == "mamba":
@@ -171,7 +171,7 @@ def _wire_sdpa_cp(model, cp_group):
     kernel, matching the reordering applied by both TE CP and PyTorch's
     context_parallel(allgather).
     """
-    from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+    from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
     for layer in model.layers.values():
         if layer.block_type == "mamba":
diff --git a/tests/functional_tests/context_parallel/run_mamba_cp.py b/tests/functional_tests/context_parallel/run_mamba_cp.py
index f420a73da0..dca10a7e94 100644
--- a/tests/functional_tests/context_parallel/run_mamba_cp.py
+++ b/tests/functional_tests/context_parallel/run_mamba_cp.py
@@ -179,7 +179,7 @@ def run_bshd_te(rank, world_size, device, config):
     """Config 1: 3D BSHD input with TE p2p CP and DualChunkSwap."""
     from torch.distributed.device_mesh import init_device_mesh
 
-    from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+    from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
     mixer_baseline, mixer_cp = _create_mixer_pair(config, device)
 
@@ -263,7 +263,7 @@ def run_thd_te(rank, world_size, device, config):
     """Config 2: 2D THD input with TE p2p CP and DualChunkSwap."""
     from torch.distributed.device_mesh import init_device_mesh
 
-    from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+    from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
     mixer_baseline, mixer_cp = _create_mixer_pair(config, device)
 
@@ -348,7 +348,7 @@ def run_thd_te_packed(rank, world_size, device, config):
     """Config 3: 2D THD input with TE p2p CP, multi-sequence packing, and seq_idx."""
     from torch.distributed.device_mesh import init_device_mesh
 
-    from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+    from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
     mixer_baseline, mixer_cp = _create_mixer_pair(config, device)
 
@@ -445,7 +445,7 @@ def run_bshd_sdpa(rank, world_size, device, config):
     """Config 4: 3D BSHD input with DualChunkSwap split (matches context_parallel allgather)."""
     from torch.distributed.device_mesh import init_device_mesh
 
-    from nemo_automodel.components.distributed.mamba_cp import (
+    from nemo_automodel.components.distributed.context_parallel.mamba import (
         MambaContextParallel,
         _redo_attention_load_balancing,
         _undo_attention_load_balancing,
diff --git a/tests/functional_tests/context_parallel/run_minimax_m3_full_model_cp.py b/tests/functional_tests/context_parallel/run_minimax_m3_full_model_cp.py
index 6d1b59dc11..dcf813f3ee 100644
--- a/tests/functional_tests/context_parallel/run_minimax_m3_full_model_cp.py
+++ b/tests/functional_tests/context_parallel/run_minimax_m3_full_model_cp.py
@@ -20,7 +20,8 @@
 drives the WHOLE model exactly as recipes/vlm/finetune.py does under CP:
 
   apply_cp(model, cp_mesh)                      # sets _cp_mesh on the sparse layers
-  train_ctx, batch, _ = prepare_cp_forward(None, ...) # torch context_parallel shards the seq
+  cp_sharder = ContextParallelSharder(None, ...)
+  train_ctx, batch = cp_sharder.shard(batch)          # torch context_parallel shards the seq
   with train_ctx(): logits_local = model(**batch)
   logits_full = context_parallel_unshard(...)   # undo the load-balanced layout
 
@@ -104,7 +105,7 @@ def main():
     from torch.distributed.device_mesh import init_device_mesh
     from torch.distributed.tensor.experimental._attention import context_parallel_unshard
 
-    from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+    from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
     from nemo_automodel.components.models.common import BackendConfig
     from nemo_automodel.components.models.minimax_m3_vl.config import MiniMaxM3VLTextConfig
     from nemo_automodel.components.models.minimax_m3_vl.model import MiniMaxM3SparseForCausalLM
@@ -173,7 +174,8 @@ def main():
         "labels": input_ids.clone(),  # required by the CP dispatch
         "position_ids": position_ids.clone(),
     }
-    train_ctx, batch, _ = prepare_cp_forward(None, device_mesh, batch)
+    cp_sharder = ContextParallelSharder(None, device_mesh, batch)
+    train_ctx, batch = cp_sharder.shard(batch)
     with torch.no_grad(), train_ctx():
         logits_local = _logits(model(input_ids=batch["input_ids"], position_ids=batch["position_ids"])).float()
 
diff --git a/tests/functional_tests/context_parallel/run_minimax_m3_vl_cp_sink.py b/tests/functional_tests/context_parallel/run_minimax_m3_vl_cp_sink.py
index 864a56a416..ea3a231d5b 100644
--- a/tests/functional_tests/context_parallel/run_minimax_m3_vl_cp_sink.py
+++ b/tests/functional_tests/context_parallel/run_minimax_m3_vl_cp_sink.py
@@ -14,14 +14,14 @@
 
 """GPU CP forward-equivalence for the MiniMax M3 VL in-forward pre-embed shard (L1, 2 GPUs).
 
-Exercises the sunk CP path end to end: ``prepare_cp_forward`` invokes the
+Exercises the sunk CP path end to end: ``ContextParallelSharder`` invokes the
 sharder-only ``prepare_model_inputs_for_cp`` hook (``shard_batch_aux_only``
 round-robin shards labels/position_ids and installs the ring-SDPA context), and
 ``MiniMaxM3SparseForConditionalGeneration.forward`` embeds the full sequence then
 ``shard_sequence_for_cp_round_robin`` shards ``inputs_embeds`` per rank. The unsharded cp2
 logits must match the cp1 eager forward.
 
-  cp1 eager (cp_mesh unset)  ==  cp2 (apply_cp + prepare_cp_forward + in-forward shard)
+  cp1 eager (cp_mesh unset)  ==  cp2 (apply_cp + ContextParallelSharder + in-forward shard)
 
 Text-only batch, dense layers (sparse disabled) so torch ``context_parallel`` ring
 SDPA is the transport. bf16 (the dense CP SDPA kernel has no fp32 path); per-token
@@ -50,7 +50,7 @@ def main():
     from torch.distributed.device_mesh import init_device_mesh
     from torch.distributed.tensor.experimental._attention import context_parallel_unshard
 
-    from nemo_automodel.components.distributed.cp_utils import prepare_cp_forward
+    from nemo_automodel.components.distributed.context_parallel import ContextParallelSharder
     from nemo_automodel.components.models.common import BackendConfig
     from nemo_automodel.components.models.minimax_m3_vl.config import MiniMaxM3VLConfig
     from nemo_automodel.components.models.minimax_m3_vl.model import MiniMaxM3SparseForConditionalGeneration
@@ -93,12 +93,13 @@ def main():
     with torch.no_grad():
         logits_eager = model(input_ids=input_ids, position_ids=position_ids.clone()).float()
 
-    # cp2: the recipe path -- apply_cp installs model.cp_mesh, prepare_cp_forward runs
+    # cp2: the recipe path -- apply_cp installs model.cp_mesh, ContextParallelSharder runs
     # the sharder-only hook, forward embeds then in-forward shards inputs_embeds.
     device_mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("cp",))
     apply_cp(model, device_mesh["cp"])
     batch = {"input_ids": input_ids.clone(), "labels": input_ids.clone(), "position_ids": position_ids.clone()}
-    train_ctx, batch, _ = prepare_cp_forward(model, device_mesh, batch)
+    cp_sharder = ContextParallelSharder(model, device_mesh, batch)
+    train_ctx, batch = cp_sharder.shard(batch)
     batch.pop("labels", None)
     with torch.no_grad(), train_ctx():
         logits_local = model(input_ids=batch["input_ids"], position_ids=batch["position_ids"]).float()
diff --git a/tests/unit_tests/_transformers/test_infrastructure.py b/tests/unit_tests/_transformers/test_infrastructure.py
index 99a7b54cc4..bff4096ad7 100644
--- a/tests/unit_tests/_transformers/test_infrastructure.py
+++ b/tests/unit_tests/_transformers/test_infrastructure.py
@@ -698,11 +698,11 @@ def test_apply_model_infrastructure_attaches_cp_hooks_for_non_te(monkeypatch):
         patch(f"{_INFRA_MODULE}._uses_te_attention", return_value=False),
         patch(f"{_INFRA_MODULE}.Checkpointer") as MockCheckpointer,
         patch(
-            "nemo_automodel.components.distributed.cp_utils.attach_context_parallel_hooks",
+            "nemo_automodel.components.distributed.context_parallel.utils.attach_context_parallel_hooks",
             side_effect=lambda mp: attached.__setitem__("ctx", attached["ctx"] + 1),
         ),
         patch(
-            "nemo_automodel.components.distributed.cp_utils.attach_cp_sdpa_hooks",
+            "nemo_automodel.components.distributed.context_parallel.utils.attach_cp_sdpa_hooks",
             side_effect=lambda mp, cp_mesh: attached.__setitem__("attn", attached["attn"] + 1),
         ),
     ):
diff --git a/tests/unit_tests/datasets/vlm/test_collate_fns.py b/tests/unit_tests/datasets/vlm/test_collate_fns.py
index 821725a750..5e0a72c52f 100644
--- a/tests/unit_tests/datasets/vlm/test_collate_fns.py
+++ b/tests/unit_tests/datasets/vlm/test_collate_fns.py
@@ -3091,7 +3091,7 @@ def test_sdpa_returns_4d_mask(self):
     def test_single_sequence_omits_packed_seq_ids(self):
         """A single (unpacked) sequence carries no ``_packed_seq_ids``; the all-gather
         CP path synthesizes the trivial one-document map downstream (see
-        ``cp_utils._synthesize_single_document_seq_ids``)."""
+        ``context_parallel.utils._synthesize_single_document_seq_ids``)."""
         from nemo_automodel.components.datasets.vlm.collate_fns import neat_packed_vlm_collater
 
         batch = [self._make_packed_sample(4, 0)]
diff --git a/tests/unit_tests/distributed/test_cp_sharder.py b/tests/unit_tests/distributed/test_cp_sharder.py
index bbe4d4c2ef..a6c15d6761 100644
--- a/tests/unit_tests/distributed/test_cp_sharder.py
+++ b/tests/unit_tests/distributed/test_cp_sharder.py
@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for the ContextParallelismSharder contract in components/distributed/cp_sharder.py.
+"""Unit tests for the ContextParallelSharder contract in components/distributed/context_parallel/sharder.py.
 
 Collectives are not exercised here (CPU CI): the tests cover the pure layout
 math — local index generation, index-based token-tensor shard, gathered-shard
-reordering — and the ContextParallelismSharder default/override resolution.
+reordering — and the ContextParallelSharder default/override resolution.
 """
 
 from __future__ import annotations
@@ -26,7 +26,7 @@
 import pytest
 import torch
 
-from nemo_automodel.components.distributed import cp_sharder as cs
+from nemo_automodel.components.distributed.context_parallel import sharder as cs
 
 
 class _FakeMesh:
@@ -46,6 +46,17 @@ def get_group(self):
         return None
 
 
+class _FakeDeviceMesh(dict):
+    """Minimal device mesh exposing CP and TP submeshes."""
+
+    def __init__(self, cp_mesh: _FakeMesh | None):
+        super().__init__()
+        self.mesh_dim_names = []
+        if cp_mesh is not None:
+            self["cp"] = cp_mesh
+            self.mesh_dim_names.append("cp")
+
+
 @pytest.fixture(autouse=True)
 def _force_no_dist(monkeypatch):
     """Pin rank resolution to the fake mesh's local rank."""
@@ -129,43 +140,45 @@ def test_gather_token_tensor_identity_without_cp():
 
 
 # ---------------------------------------------------------------------------
-# ContextParallelismSharder default/override resolution
+# ContextParallelSharder default/override resolution
 # ---------------------------------------------------------------------------
 def test_sharder_default_shard_token_tensor_uses_indices():
-    sharder = cs.ContextParallelismSharder(
+    sharder = cs.ContextParallelSharder(
+        device_mesh=_FakeDeviceMesh(_FakeMesh(2, 1)),
         shard_batch=lambda *a, **k: (contextlib.nullcontext, {}, None),
         local_token_global_indices=cs.contiguous_local_indices,
     )
     full = torch.randn(1, 8)
-    local = sharder.shard_token_tensor(_FakeMesh(2, 1), full, seq_dim=1)
+    local = sharder.shard_token_tensor(full, seq_dim=1)
     torch.testing.assert_close(local, full[:, 4:])
 
 
 def test_shard_batch_contiguous_records_shard_layout():
     """shard_batch reports original/padded lengths as ShardLayout; once installed,
     the token verbs accept caller-coordinate tensors and reject mismatched ones."""
-    sharder = cs.ContextParallelismSharder(
-        shard_batch=None,
+    mesh = _FakeMesh(2, 0)
+    sharder = cs.ContextParallelSharder(
+        device_mesh=_FakeDeviceMesh(mesh),
+        shard_batch=cs.shard_batch_identity,
         local_token_global_indices=cs.contiguous_local_indices,
     )
-    mesh = _FakeMesh(2, 0)
     batch = {"input_ids": torch.arange(6).unsqueeze(0), "labels": torch.arange(6).unsqueeze(0)}
     _, _, layout = cs.shard_batch_contiguous(mesh, None, batch)  # pads 6 -> 8
     sharder.shard_layout = layout
 
     assert (sharder.shard_layout.original_seq_len, sharder.shard_layout.padded_seq_len) == (6, 8)
     # down: unpadded tensor auto-pads with the explicit fill, rank 0 owns [0:4]
-    local = sharder.shard_token_tensor(mesh, torch.arange(6.0).unsqueeze(0), fill=-1.0)
+    local = sharder.shard_token_tensor(torch.arange(6.0).unsqueeze(0), fill=-1.0)
     assert torch.equal(local, torch.tensor([[0.0, 1.0, 2.0, 3.0]]))
     # the pad_multiple silent-misalignment window is closed: a plausible but
     # wrong length raises even though it divides cp_size
     with pytest.raises(ValueError, match="padded_seq_len=8"):
-        sharder.shard_token_tensor(mesh, torch.zeros(1, 4))
+        sharder.shard_token_tensor(torch.zeros(1, 4))
     # up: trim validates the gathered length against the captured layout (no
     # collective runs in this single-process test, so the gather stays local
     # and the guard must fire rather than mis-trim)
     with pytest.raises(ValueError, match="reported padded_seq_len 8"):
-        sharder.gather_token_tensor(mesh, torch.zeros(1, 4), trim=True)
+        sharder.gather_token_tensor(torch.zeros(1, 4), trim=True)
 
 
 def test_sharder_repositioned_layout_round_trips_input_coordinates():
@@ -174,60 +187,80 @@ def test_sharder_repositioned_layout_round_trips_input_coordinates():
     columns and dropped input pad slots come back as fill."""
     # input row: [a, b, PAD] -> rebuilt row: [a, b, X, X] (doc re-padded to 4)
     positions = torch.tensor([[0, 1, -1]])
-    sharder = cs.ContextParallelismSharder(
-        shard_batch=None,
+    sharder = cs.ContextParallelSharder(
+        device_mesh=_FakeDeviceMesh(_FakeMesh(2, 0)),
+        shard_batch=cs.shard_batch_identity,
         local_token_global_indices=cs.contiguous_local_indices,
         shard_layout=cs.ShardLayout(padded_seq_len=4, input_token_stream_positions=positions),
     )
-    mesh = _FakeMesh(2, 0)  # rank 0 owns columns [0:2]
-    local = sharder.shard_token_tensor(mesh, torch.tensor([[10.0, 20.0, 99.0]]), fill=0.0)
+    local = sharder.shard_token_tensor(torch.tensor([[10.0, 20.0, 99.0]]), fill=0.0)
     assert torch.equal(local, torch.tensor([[10.0, 20.0]]))
     with pytest.raises(ValueError, match="fill"):
-        sharder.shard_token_tensor(mesh, torch.tensor([[10.0, 20.0, 99.0]]))
+        sharder.shard_token_tensor(torch.tensor([[10.0, 20.0, 99.0]]))
 
     # up: gather (identity at cp<=1 here) then map back to input coordinates
     full_rows = torch.tensor([[10.0, 20.0, 7.0, 7.0]])
-    out = sharder.gather_token_tensor(_FakeMesh(1), full_rows, trim=True, fill=-5.0)
+    gather_sharder = cs.ContextParallelSharder(
+        device_mesh=_FakeDeviceMesh(_FakeMesh(1)),
+        shard_batch=cs.shard_batch_identity,
+        local_token_global_indices=cs.contiguous_local_indices,
+        shard_layout=sharder.shard_layout,
+    )
+    out = gather_sharder.gather_token_tensor(full_rows, trim=True, fill=-5.0)
     assert torch.equal(out, torch.tensor([[10.0, 20.0, -5.0]]))
     with pytest.raises(ValueError, match="fill"):
-        sharder.gather_token_tensor(_FakeMesh(1), full_rows, trim=True)
+        gather_sharder.gather_token_tensor(full_rows, trim=True)
 
 
 def test_gather_trim_raises_without_captured_facts():
-    sharder = cs.ContextParallelismSharder(
-        shard_batch=None,
+    sharder = cs.ContextParallelSharder(
+        device_mesh=_FakeDeviceMesh(_FakeMesh(1)),
+        shard_batch=cs.shard_batch_identity,
         local_token_global_indices=cs.contiguous_local_indices,
     )
     with pytest.raises(NotImplementedError, match="no shard layout to trim to"):
-        sharder.gather_token_tensor(_FakeMesh(1), torch.zeros(1, 4), trim=True)
+        sharder.gather_token_tensor(torch.zeros(1, 4), trim=True)
 
 
 def test_reported_indices_validate_stream_length():
     # Reported index maps flatten + cast to long, and reject a padded_seq_len
     # that does not match the partition the shard reported.
-    sharder = cs.ContextParallelismSharder(
-        shard_batch=lambda *a, **k: (contextlib.nullcontext, {}, None),
-        local_token_global_indices=None,
-        shard_layout=cs.ShardLayout(local_token_global_indices=torch.tensor([[1, 0]], dtype=torch.int32)),
-    )
-    assert torch.equal(sharder._indices(_FakeMesh(2, 0), 4, None), torch.tensor([1, 0]))
-    assert torch.equal(sharder._indices(None, 2, None), torch.tensor([1, 0]))  # no mesh -> cp_size 1
+    kwargs = {
+        "shard_batch": lambda *a, **k: (contextlib.nullcontext, {}, None),
+        "local_token_global_indices": None,
+        "shard_layout": cs.ShardLayout(local_token_global_indices=torch.tensor([[1, 0]], dtype=torch.int32)),
+    }
+    sharder = cs.ContextParallelSharder(device_mesh=_FakeDeviceMesh(_FakeMesh(2, 0)), **kwargs)
+    assert torch.equal(sharder._indices(4, None), torch.tensor([1, 0]))
+    no_mesh_sharder = cs.ContextParallelSharder(**kwargs)
+    assert torch.equal(no_mesh_sharder._indices(2, None), torch.tensor([1, 0]))  # no mesh -> cp_size 1
     with pytest.raises(ValueError, match="does not match"):
-        sharder._indices(_FakeMesh(2, 0), 6, None)
+        sharder._indices(6, None)
 
 
 def test_sharder_token_verbs_unavailable_for_data_dependent_layouts():
     # THD/magi layouts depend on batch content (cu_seqlens / dispatch solver),
     # so their framework sharders carry no index map and the token-tensor verbs
     # must fail loudly rather than shard the wrong slice.
-    sharder = cs.ContextParallelismSharder(
+    sharder = cs.ContextParallelSharder(
+        device_mesh=_FakeDeviceMesh(_FakeMesh(2, 0)),
         shard_batch=lambda *a, **k: (contextlib.nullcontext, {}, None),
         local_token_global_indices=None,
     )
     with pytest.raises(NotImplementedError, match="data-dependent"):
-        sharder.shard_token_tensor(_FakeMesh(2, 0), torch.randn(1, 8))
+        sharder.shard_token_tensor(torch.randn(1, 8))
     with pytest.raises(NotImplementedError, match="data-dependent"):
-        sharder.gather_token_tensor(_FakeMesh(2, 0), torch.randn(1, 4))
+        sharder.gather_token_tensor(torch.randn(1, 4))
+
+
+def test_constructor_rejects_mixed_resolution_and_strategy_arguments():
+    with pytest.raises(TypeError, match="mutually exclusive"):
+        cs.ContextParallelSharder(
+            None,
+            _FakeMesh(1),
+            {},
+            shard_batch=cs.shard_batch_identity,
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -309,7 +342,7 @@ def test_shard_batch_aux_only_matches_load_balanced(monkeypatch):
     """The aux-only shard pads labels/position_ids/loss_mask identically to the
     load-balanced shard but leaves the primary stream full-length and out of the
     CP buffer list."""
-    from nemo_automodel.components.distributed import cp_utils
+    from nemo_automodel.components.distributed.context_parallel import utils as cp_utils
 
     captured: dict = {}
 
@@ -356,7 +389,7 @@ def make_batch():
 
 def test_shard_batch_aux_only_reports_padded_layout(monkeypatch):
     """The returned ShardLayout carries the primary stream's target padded length."""
-    from nemo_automodel.components.distributed import cp_utils
+    from nemo_automodel.components.distributed.context_parallel import utils as cp_utils
 
     monkeypatch.setattr(cp_utils, "create_context_parallel_ctx", lambda *a, **k: contextlib.nullcontext())
     monkeypatch.setattr(cp_utils, "get_train_context", lambda *a, **k: contextlib.nullcontext)
diff --git a/tests/unit_tests/distributed/test_cp_utils.py b/tests/unit_tests/distributed/test_cp_utils.py
index 4e47054d63..b1649326c1 100644
--- a/tests/unit_tests/distributed/test_cp_utils.py
+++ b/tests/unit_tests/distributed/test_cp_utils.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for :pyfile:`nemo_automodel/components/distributed/cp_utils.py`.
+"""Unit tests for :pyfile:`nemo_automodel/components/distributed/context_parallel/utils.py`.
 
 The real implementation relies heavily on ``torch.distributed`` and GPU-specific
 behavior.  These unit-tests therefore *mock* the heavyweight distributed pieces
@@ -29,9 +29,9 @@
 import torch
 
 # Import module under test
-from nemo_automodel.components.distributed import cp_utils as _cu
-from nemo_automodel.components.distributed.cp_sharder import (
-    ContextParallelismSharder,
+from nemo_automodel.components.distributed.context_parallel import utils as _cu
+from nemo_automodel.components.distributed.context_parallel.sharder import (
+    ContextParallelSharder,
     contiguous_local_indices,
     round_robin_local_indices,
     shard_batch_aux_only,
@@ -40,11 +40,11 @@
 from nemo_automodel.components.models.gemma4_moe import cp_batch as _cm
 
 
-# ContextParallelismSharder used by the model-owned dispatch tests below (passed as an explicit
+# ContextParallelSharder used by the model-owned dispatch tests below (passed as an explicit
 # _make_cp_batch_and_ctx parameter; the batch itself stays pure tensors). Exercises the public
 # contiguous shard (the production entry DSV4/Gemma4 wrap) on the model-provided per-token keys.
 def _contiguous_sharder():
-    return ContextParallelismSharder(
+    return ContextParallelSharder(
         shard_batch=partial(
             shard_batch_contiguous,
             extra_seq_keys={"per_layer_inputs": 1, "_packed_seq_ids": 1, "mm_token_type_ids": 1},
@@ -95,6 +95,16 @@ def __init__(self, cp_size: int, tp_size: int, cp_rank: int = 0):
         self.mesh_dim_names = ["cp", "tp"]
 
 
+def _construct_strategy_sharder(strategy, device_mesh):
+    """Construct a mesh-configured sharder from a resolved strategy."""
+    return ContextParallelSharder(
+        device_mesh=device_mesh,
+        shard_batch=strategy.shard_batch,
+        local_token_global_indices=strategy.local_token_global_indices,
+        shard_layout=strategy.shard_layout,
+    )
+
+
 def test_make_cp_batch_and_ctx_no_mesh():
     """When *no* device mesh is provided the call should be a no-op."""
     input_ids = torch.tensor([[1, 2, 3]])
@@ -136,7 +146,7 @@ def make_native_batch(cp_mesh, tp_mesh, batch, **kwargs):
         "input_ids": torch.tensor([[1, 2, 3, 4]]),
         "labels": torch.tensor([[1, 2, 3, 4]]),
     }
-    sharder = ContextParallelismSharder(
+    sharder = ContextParallelSharder(
         shard_batch=make_native_batch,
         local_token_global_indices=contiguous_local_indices,
     )
@@ -610,7 +620,7 @@ def test_synthesize_single_document_seq_ids_noop_when_present():
     assert torch.equal(batch["_packed_seq_ids"], existing)
 
 
-def test_magi_dispatches_at_the_te_rung():
+def test_sharder_constructor_derives_magi_and_thd_without_sharding(monkeypatch):
     """An enabled magi occupies the same _make_cp_batch_and_ctx rung as the TE
     path: (nullcontext, prepped batch), never the torch-native CP context."""
     import contextlib as _ctxlib
@@ -628,16 +638,19 @@ def make_cp_batch(
             seen.update(cp_mesh=cp_mesh, model=model, is_thd=is_thd, pad=padding_token_id, chunks=num_chunks)
             return ({"prepared": True}, None) if return_local_indices else {"prepared": True}
 
-    model = SimpleNamespace()  # no prepare_model_inputs_for_cp -> hook path skipped
-    ctx, batch, _ = _cu.prepare_cp_forward(
+    magi = _FakeMagi()
+    model = SimpleNamespace(backend=SimpleNamespace(attn="magi"))
+    monkeypatch.setattr(_cu, "_magi_state_from_model", lambda actual, mesh: magi if actual is model else None)
+    batch = {"input_ids": torch.tensor([[1, 2]]), "qkv_format": "thd"}
+    sharder = ContextParallelSharder(
         model,
         _DummyDeviceMesh(cp_size=2, tp_size=1),
-        {"input_ids": torch.tensor([[1, 2]])},
-        magi=_FakeMagi(),
-        use_te=True,
+        batch,
         padding_token_id=7,
         num_chunks=3,
     )
+    assert not seen
+    ctx, batch = sharder.shard(batch)
     assert ctx is _ctxlib.nullcontext
     assert batch == {"prepared": True}
     assert seen["model"] is model
@@ -648,9 +661,22 @@ def make_cp_batch(
     assert batch2 == {"prepared": True} and seen["cp_mesh"] is None
 
 
-def test_te_dispatches_through_a_framework_sharder(monkeypatch):
-    """use_te resolves to a framework-built THD ContextParallelismSharder whose shard_batch
-    wraps make_cp_batch_for_te, threading the recipe-static args through."""
+def test_magi_state_is_derived_from_live_model():
+    """Magi backend kind and domain come from the model, not recipe arguments."""
+    from types import SimpleNamespace
+
+    model = SimpleNamespace(
+        backend=SimpleNamespace(attn="magi"),
+        config=SimpleNamespace(vision_config=SimpleNamespace()),
+    )
+    state = _cu._magi_state_from_model(model, _DummyDeviceMesh(cp_size=1, tp_size=1))
+    assert state.enabled and state.custom
+    assert state.domain == "vlm"
+    assert state.cp_size == 1
+
+
+def test_sharder_constructor_derives_te_from_model_and_thd_from_batch(monkeypatch):
+    """A TE model and THD batch resolve a sharder without recipe-owned flags."""
     seen = {}
 
     def fake_make_cp_batch_for_te(
@@ -664,26 +690,39 @@ def fake_make_cp_batch_for_te(
     monkeypatch.setattr(_cu, "make_cp_batch_for_te", fake_make_cp_batch_for_te)
 
     device_mesh = _DummyDeviceMesh(cp_size=2, tp_size=1)
-    ctx, batch, _ = _cu._make_cp_batch_and_ctx(
+    model = type("_Model", (), {"backend": type("_Backend", (), {"attn": "te"})()})()
+    batch = {"input_ids": torch.tensor([[1, 2]]), "qkv_format": "thd"}
+    sharder = ContextParallelSharder(
+        model,
         device_mesh,
-        {"input_ids": torch.tensor([[1, 2]])},
-        use_te=True,
+        batch,
         padding_token_id=7,
         num_chunks=3,
-        seq_lens_padding_value=-5,
     )
+    assert not seen
+    ctx, batch = sharder.shard(batch)
     assert ctx is contextlib.nullcontext
     assert batch == {"thd": True}
     assert seen["cp_mesh"] is device_mesh["cp"]
-    assert (seen["pad"], seen["fmt"], seen["chunks"], seen["sent"]) == (7, "thd", 3, -5)
+    assert (seen["pad"], seen["fmt"], seen["chunks"], seen["sent"]) == (7, "thd", 3, -1000)
 
-    # THD conversion also runs at cp<=1 (packing at cp=1), like before.
-    seen.clear()
-    _, batch2, _ = _cu._make_cp_batch_and_ctx(None, {"input_ids": torch.tensor([[1, 2]])}, use_te=True)
-    assert batch2 == {"thd": True} and seen["cp_mesh"] is None
+
+def test_sharder_constructor_does_not_infer_te_from_batch_alone(monkeypatch):
+    """A THD-origin batch does not force TE preparation on a non-TE model."""
+    monkeypatch.setattr(
+        _cu,
+        "make_cp_batch_for_te",
+        lambda *args, **kwargs: pytest.fail("TE batch preparation should not run"),
+    )
+    model = type("_Model", (), {"backend": type("_Backend", (), {"attn": "sdpa"})()})()
+    batch = {"input_ids": torch.tensor([[1, 2]]), "qkv_format": "thd"}
+    sharder = ContextParallelSharder(model, _DummyDeviceMesh(cp_size=1, tp_size=1), batch)
+    ctx, out = sharder.shard(batch)
+    assert ctx is contextlib.nullcontext
+    assert out is batch
 
 
-def test_prepare_cp_forward_merges_model_hook_batch_updates(monkeypatch):
+def test_sharder_constructor_merges_model_hook_batch_updates(monkeypatch):
     """Model-owned hooks may return batch metadata in addition to the sharder."""
 
     cp_context_kwargs = {}
@@ -703,7 +742,7 @@ def prepare_model_inputs_for_cp(self, batch, *, num_chunks):
             assert num_chunks == 3
             assert batch["mm_token_type_ids"].shape == (1, 4)
             return {
-                "cp_sharder": ContextParallelismSharder(
+                "cp_sharder": ContextParallelSharder(
                     shard_batch=shard_batch_aux_only,
                     local_token_global_indices=round_robin_local_indices,
                 ),
@@ -720,17 +759,19 @@ def prepare_model_inputs_for_cp(self, batch, *, num_chunks):
         "image_grid_hws": torch.tensor([[2, 2]]),
     }
 
-    ctx, out, sharder = _cu.prepare_cp_forward(_Model(), _DummyDeviceMesh(cp_size=2, tp_size=1), batch, num_chunks=3)
-
+    sharder = ContextParallelSharder(_Model(), _DummyDeviceMesh(cp_size=2, tp_size=1), batch, num_chunks=3)
+
+    assert "cp_sharder" not in batch
+    assert torch.equal(batch["input_ids"], torch.tensor([[1, 2, 3, 4]]))
+    assert torch.equal(batch["labels"], torch.tensor([[10, 20, 30, 40]]))
+    assert batch["position_ids"] is position_ids
+    assert batch["mm_token_type_ids"] is None
+    assert batch["image_grid_thw"] is image_grid_thw
+    assert batch["image_grid_hws"] is None
+    assert cp_context_kwargs == {}
+    ctx, out = sharder.shard(batch)
     assert ctx is contextlib.nullcontext
     assert out is batch
-    assert "cp_sharder" not in out
-    assert torch.equal(out["input_ids"], torch.tensor([[1, 2, 3, 4]]))
-    assert torch.equal(out["labels"], torch.tensor([[10, 20, 30, 40]]))
-    assert out["position_ids"] is position_ids
-    assert out["mm_token_type_ids"] is None
-    assert out["image_grid_thw"] is image_grid_thw
-    assert out["image_grid_hws"] is None
     assert cp_context_kwargs["cp_buffers"][1] is position_ids
     assert cp_context_kwargs["cp_seq_dims"] == [1, 2]
     assert sharder.shard_layout.original_seq_len == 4
@@ -749,19 +790,19 @@ def fake_make_cp_batch_for_te(cp_mesh, batch, *, return_local_indices=False, **k
     monkeypatch.setattr(_cu, "make_cp_batch_for_te", fake_make_cp_batch_for_te)
 
     cp2 = _DummySubMesh(2)
-    sharder = _cu._resolve_cp_sharder(
-        cp2, None, magi=None, use_te=True, num_chunks=1, seq_lens_padding_value=-1000, model=None
+    strategy = _cu._resolve_cp_sharder(
+        cp2, None, magi=None, is_thd=True, num_chunks=1, seq_lens_padding_value=-1000, model=None
     )
+    sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=2, tp_size=1))
     full = torch.arange(4.0)  # [T] token-aligned tensor, THD seq_dim=0
 
     with pytest.raises(NotImplementedError, match="before the first shard"):
-        sharder.shard_token_tensor(cp2, full, seq_dim=0)
+        sharder.shard_token_tensor(full, seq_dim=0)
 
-    _, _, layout = sharder.shard_batch(cp2, None, {"input_ids": torch.tensor([1, 2, 3, 4])})
-    sharder.shard_layout = layout
-    assert torch.equal(sharder.shard_token_tensor(cp2, full, seq_dim=0), torch.tensor([0.0, 3.0]))
+    sharder.shard({"input_ids": torch.tensor([1, 2, 3, 4])})
+    assert torch.equal(sharder.shard_token_tensor(full, seq_dim=0), torch.tensor([0.0, 3.0]))
     with pytest.raises(ValueError, match="does not match"):
-        sharder.shard_token_tensor(cp2, torch.arange(6.0), seq_dim=0)
+        sharder.shard_token_tensor(torch.arange(6.0), seq_dim=0)
 
 
 class _FakeMagiState:
@@ -784,20 +825,20 @@ def test_magi_sharder_captures_hf_dispatch_facts():
     the original length and the verbs work in the caller's [1, S] coordinates."""
     cp2 = _DummySubMesh(2)
     # global padded length 4 = 2 local x cp 2; input was [1, 3] -> tail pad of 1
-    sharder = _cu._resolve_cp_sharder(
+    strategy = _cu._resolve_cp_sharder(
         cp2,
         None,
         magi=_FakeMagiState(torch.tensor([[0, 2]])),
-        use_te=False,
+        is_thd=False,
         num_chunks=1,
         seq_lens_padding_value=-1000,
         model=None,
     )
-    _, _, layout = sharder.shard_batch(cp2, None, {"input_ids": torch.tensor([[1, 2, 3]])})
-    sharder.shard_layout = layout
+    sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=2, tp_size=1))
+    sharder.shard({"input_ids": torch.tensor([[1, 2, 3]])})
     assert (sharder.shard_layout.original_seq_len, sharder.shard_layout.padded_seq_len) == (3, 4)
     # down: original-length tensor auto-pads then follows the dispatch permutation
-    local = sharder.shard_token_tensor(cp2, torch.tensor([[10.0, 20.0, 30.0]]), fill=0.0)
+    local = sharder.shard_token_tensor(torch.tensor([[10.0, 20.0, 30.0]]), fill=0.0)
     assert torch.equal(local, torch.tensor([[10.0, 30.0]]))
 
 
@@ -805,21 +846,21 @@ def test_magi_sharder_captures_packed_row_shape():
     """Packed magi over a THD flatten with no extra dispatch pad: the sharder
     captures the pre-flatten row shape (padded == rows x cols)."""
     cp2 = _DummySubMesh(2)
-    sharder = _cu._resolve_cp_sharder(
+    strategy = _cu._resolve_cp_sharder(
         cp2,
         None,
         magi=_FakeMagiState(torch.tensor([[0, 3]])),
-        use_te=True,
+        is_thd=True,
         num_chunks=1,
         seq_lens_padding_value=-1000,
         model=None,
     )
-    _, _, layout = sharder.shard_batch(cp2, None, {"input_ids": torch.tensor([[1, 2], [3, 4]])})
-    sharder.shard_layout = layout
+    sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=2, tp_size=1))
+    sharder.shard({"input_ids": torch.tensor([[1, 2], [3, 4]])})
     assert sharder.shard_layout.input_row_shape == (2, 2)
     assert sharder.shard_layout.padded_seq_len == 4
     rows = torch.tensor([[10.0, 20.0], [30.0, 40.0]])
-    assert torch.equal(sharder.shard_token_tensor(cp2, rows), torch.tensor([10.0, 40.0]))
+    assert torch.equal(sharder.shard_token_tensor(rows), torch.tensor([10.0, 40.0]))
 
 
 def test_make_cp_batch_for_te_identity_indices_without_cp():
@@ -848,15 +889,15 @@ def test_round_robin_sharder_captures_lengths_and_pads_token_tensors(monkeypatch
 
     assert (sharder.shard_layout.original_seq_len, sharder.shard_layout.padded_seq_len) == (6, 8)
     # down: unpadded [1, 6] advantages ride with an explicit fill
-    local = sharder.shard_token_tensor(device_mesh["cp"], torch.arange(6.0).unsqueeze(0), fill=0.0)
+    local = sharder.shard_token_tensor(torch.arange(6.0).unsqueeze(0), fill=0.0)
     # rank 0 under 2*cp=4 chunks of len 2: chunks 0 and 3 -> positions [0,1,6,7]
     assert torch.equal(local, torch.tensor([[0.0, 1.0, 0.0, 0.0]]))
     # mismatched length is loud, not silently mis-sharded
     with pytest.raises(ValueError, match="padded_seq_len=8"):
-        sharder.shard_token_tensor(device_mesh["cp"], torch.zeros(1, 7), fill=0.0)
+        sharder.shard_token_tensor(torch.zeros(1, 7), fill=0.0)
     # unpadded without fill is loud too
     with pytest.raises(ValueError, match="fill"):
-        sharder.shard_token_tensor(device_mesh["cp"], torch.zeros(1, 6))
+        sharder.shard_token_tensor(torch.zeros(1, 6))
 
 
 def test_none_sharder_captures_lengths_for_trim():
@@ -867,7 +908,7 @@ def test_none_sharder_captures_lengths_for_trim():
     _, _, sharder = _cu._make_cp_batch_and_ctx(device_mesh, batch)
     assert (sharder.shard_layout.original_seq_len, sharder.shard_layout.padded_seq_len) == (6, 6)
     t = torch.randn(1, 6)
-    assert torch.equal(sharder.gather_token_tensor(device_mesh["cp"], t, trim=True), t)
+    assert torch.equal(sharder.gather_token_tensor(t, trim=True), t)
 
 
 def test_te_sharder_captures_row_shape(monkeypatch):
@@ -882,34 +923,34 @@ def fake_make_cp_batch_for_te(cp_mesh, batch, *, return_local_indices=False, **k
     monkeypatch.setattr(_cu, "make_cp_batch_for_te", fake_make_cp_batch_for_te)
 
     cp1 = _DummySubMesh(1)
-    sharder = _cu._resolve_cp_sharder(
-        cp1, None, magi=None, use_te=True, num_chunks=1, seq_lens_padding_value=-1000, model=None
+    strategy = _cu._resolve_cp_sharder(
+        cp1, None, magi=None, is_thd=True, num_chunks=1, seq_lens_padding_value=-1000, model=None
     )
-    _, _, layout = sharder.shard_batch(cp1, None, {"input_ids": torch.arange(4).view(2, 2)})
-    sharder.shard_layout = layout
+    sharder = _construct_strategy_sharder(strategy, _DummyDeviceMesh(cp_size=1, tp_size=1))
+    sharder.shard({"input_ids": torch.arange(4).view(2, 2)})
     assert sharder.shard_layout.input_row_shape == (2, 2)
     assert sharder.shard_layout.padded_seq_len == 4
 
     # down: row-coordinate [2, 2] flattens to the stream before sharding
     rows = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
-    assert torch.equal(sharder.shard_token_tensor(cp1, rows), torch.tensor([1.0, 2.0, 3.0, 4.0]))
+    assert torch.equal(sharder.shard_token_tensor(rows), torch.tensor([1.0, 2.0, 3.0, 4.0]))
     # up: gather restores the row coordinate
-    assert torch.equal(sharder.gather_token_tensor(cp1, torch.tensor([1.0, 2.0, 3.0, 4.0]), seq_dim=0, trim=True), rows)
+    assert torch.equal(sharder.gather_token_tensor(torch.tensor([1.0, 2.0, 3.0, 4.0]), seq_dim=0, trim=True), rows)
 
 
 def test_resolve_cp_sharder_layers():
     """Resolution order: model-owned > magi > TE > generic round-robin > none."""
-    from nemo_automodel.components.distributed.cp_sharder import round_robin_local_indices
+    from nemo_automodel.components.distributed.context_parallel.sharder import round_robin_local_indices
 
     cp2 = _DummySubMesh(2)
     model_sharder = _contiguous_sharder()
-    common = dict(magi=None, use_te=False, num_chunks=1, seq_lens_padding_value=-1000, model=None)
+    common = dict(magi=None, is_thd=False, num_chunks=1, seq_lens_padding_value=-1000, model=None)
 
     # model-owned wins over everything, including native THD prep at cp<=1
-    assert _cu._resolve_cp_sharder(cp2, model_sharder, **{**common, "use_te": True}) is model_sharder
-    assert _cu._resolve_cp_sharder(None, model_sharder, **{**common, "use_te": True}) is model_sharder
+    assert _cu._resolve_cp_sharder(cp2, model_sharder, **{**common, "is_thd": True}) is model_sharder
+    assert _cu._resolve_cp_sharder(None, model_sharder, **{**common, "is_thd": True}) is model_sharder
     # TE resolves at cp<=1 when no model-owned sharder is present
-    assert _cu._resolve_cp_sharder(None, None, **{**common, "use_te": True}).local_token_global_indices is None
+    assert _cu._resolve_cp_sharder(None, None, **{**common, "is_thd": True}).local_token_global_indices is None
     # generic torch context_parallel is the framework default at cp>1
     generic = _cu._resolve_cp_sharder(cp2, None, **common)
     assert generic.local_token_global_indices is round_robin_local_indices
@@ -917,9 +958,13 @@ def test_resolve_cp_sharder_layers():
     for mesh in (None, _DummySubMesh(1)):
         none_sharder = _cu._resolve_cp_sharder(mesh, None, **common)
         batch = {"input_ids": torch.tensor([[1, 2, 3]])}
-        ctx, out, _ = none_sharder.shard_batch(mesh, None, batch)
+        none_sharder = _construct_strategy_sharder(
+            none_sharder,
+            _DummyDeviceMesh(cp_size=mesh.size() if mesh is not None else 1, tp_size=1) if mesh is not None else None,
+        )
+        ctx, out = none_sharder.shard(batch)
         assert ctx is contextlib.nullcontext and out is batch
         # token verbs are identities at cp<=1 (lengths were captured: 3 == 3)
         t = torch.randn(1, 3)
-        assert torch.equal(none_sharder.shard_token_tensor(mesh, t), t)
-        assert none_sharder.gather_token_tensor(mesh, t) is t
+        assert torch.equal(none_sharder.shard_token_tensor(t), t)
+        assert none_sharder.gather_token_tensor(t) is t
diff --git a/tests/unit_tests/distributed/test_cp_utils_diffcov.py b/tests/unit_tests/distributed/test_cp_utils_diffcov.py
index 690cc4b790..356c9026c1 100644
--- a/tests/unit_tests/distributed/test_cp_utils_diffcov.py
+++ b/tests/unit_tests/distributed/test_cp_utils_diffcov.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Targeted unit tests for context-parallel helper paths in cp_utils.
+"""Targeted unit tests for context-parallel helper paths in context_parallel.utils.
 
 These exercise the CP attention SDPA-swap hooks, the linear-attn position hook,
 the classic DTensor _cp_sdpa path (DTensor mocked), padding/prepare helpers, and
@@ -26,8 +26,8 @@
 import torch
 import torch.nn.functional as F
 
-from nemo_automodel.components.distributed import cp_sharder as cm
-from nemo_automodel.components.distributed import cp_utils as cu
+from nemo_automodel.components.distributed.context_parallel import sharder as cm
+from nemo_automodel.components.distributed.context_parallel import utils as cu
 
 
 @pytest.fixture(autouse=True)
diff --git a/tests/unit_tests/distributed/test_cp_utils_inputs_embeds.py b/tests/unit_tests/distributed/test_cp_utils_inputs_embeds.py
index db655e4e49..2991dcf356 100644
--- a/tests/unit_tests/distributed/test_cp_utils_inputs_embeds.py
+++ b/tests/unit_tests/distributed/test_cp_utils_inputs_embeds.py
@@ -29,7 +29,7 @@
 import pytest
 import torch
 
-from nemo_automodel.components.distributed import cp_utils as _cu
+from nemo_automodel.components.distributed.context_parallel import utils as _cu
 
 
 class _DummySubMesh:
@@ -555,7 +555,7 @@ def test_padding_attention_mask_pad_value_is_zero(monkeypatch):
     the function so this case is moot, but the PAD_FILL table is the right
     place to encode the semantic in case the strip is ever revisited.
     """
-    from nemo_automodel.components.distributed import cp_sharder as _cs
+    from nemo_automodel.components.distributed.context_parallel import sharder as _cs
 
     # Just verify the PAD_FILL table itself maps attention_mask -> False
     # (the runtime code path is currently unreachable because attention_mask
diff --git a/tests/unit_tests/distributed/test_magi_attn_utils.py b/tests/unit_tests/distributed/test_magi_attn_utils.py
index 2fbbdd108d..c7206377ff 100644
--- a/tests/unit_tests/distributed/test_magi_attn_utils.py
+++ b/tests/unit_tests/distributed/test_magi_attn_utils.py
@@ -28,8 +28,8 @@
 import torch
 import torch.nn as nn
 
-import nemo_automodel.components.distributed.magi_attn_utils as mu
-from nemo_automodel.components.distributed.magi_attn_utils import AttnMaskSpec, MagiState, setup_magi
+import nemo_automodel.components.distributed.context_parallel.magi as mu
+from nemo_automodel.components.distributed.context_parallel.magi import AttnMaskSpec, MagiState, setup_magi
 
 
 class _FakeCfg:
diff --git a/tests/unit_tests/distributed/test_mamba_cp.py b/tests/unit_tests/distributed/test_mamba_cp.py
index 30f268c4be..913656998d 100644
--- a/tests/unit_tests/distributed/test_mamba_cp.py
+++ b/tests/unit_tests/distributed/test_mamba_cp.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for :pyfile:`nemo_automodel/components/distributed/mamba_cp.py`.
+"""Unit tests for :pyfile:`nemo_automodel/components/distributed/context_parallel/mamba.py`.
 
 Tests mock the distributed process group so they can run on CPU-only CI
 systems while still verifying dimension calculations, parameter slicing,
@@ -27,7 +27,7 @@
 import torch
 import torch.nn as nn
 
-from nemo_automodel.components.distributed.mamba_cp import MambaContextParallel
+from nemo_automodel.components.distributed.context_parallel.mamba import MambaContextParallel
 
 # ---------------------------------------------------------------------------
 # Lightweight stubs for torch.distributed.ProcessGroup
@@ -391,7 +391,9 @@ def fake_cp2hp(tensor, cp_group, batch_size):
             H_local = H // cp_size
             return torch.randn(batch_size, L_local * cp_size, H_local)
 
-        with patch("nemo_automodel.components.distributed.mamba_cp._all_to_all_cp2hp", side_effect=fake_cp2hp):
+        with patch(
+            "nemo_automodel.components.distributed.context_parallel.mamba._all_to_all_cp2hp", side_effect=fake_cp2hp
+        ):
             mcp.pre_conv_ssm(projected)
 
         assert len(captured_calls) == 5, f"Expected 5 all-to-all calls, got {len(captured_calls)}"
@@ -438,7 +440,9 @@ def fake_cp2hp(tensor, cp_group, batch_size):
             H_local = H // cp_size
             return torch.randn(batch_size, L_local * cp_size, H_local)
 
-        with patch("nemo_automodel.components.distributed.mamba_cp._all_to_all_cp2hp", side_effect=fake_cp2hp):
+        with patch(
+            "nemo_automodel.components.distributed.context_parallel.mamba._all_to_all_cp2hp", side_effect=fake_cp2hp
+        ):
             mcp.pre_conv_ssm(projected)
 
         b_state_input = captured_calls[2]
@@ -483,7 +487,9 @@ def fake_cp2hp(tensor, cp_group, batch_size):
             H_local = H_t // cp_size
             return torch.randn(B_t, L_t * cp_size, H_local)
 
-        with patch("nemo_automodel.components.distributed.mamba_cp._all_to_all_cp2hp", side_effect=fake_cp2hp):
+        with patch(
+            "nemo_automodel.components.distributed.context_parallel.mamba._all_to_all_cp2hp", side_effect=fake_cp2hp
+        ):
             output = mcp.pre_conv_ssm(projected)
 
         d_inner_local = d_inner // cp_size
@@ -526,7 +532,9 @@ def fake_hp2cp(tensor, cp_group, batch_size):
             H_out = H_t * cp_size
             return torch.randn(B_t, L_out, H_out)
 
-        with patch("nemo_automodel.components.distributed.mamba_cp._all_to_all_hp2cp", side_effect=fake_hp2cp):
+        with patch(
+            "nemo_automodel.components.distributed.context_parallel.mamba._all_to_all_hp2cp", side_effect=fake_hp2cp
+        ):
             output = mcp.post_conv_ssm(ssm_output)
 
         assert output.shape == (B, L_local, d_inner), f"Expected ({B}, {L_local}, {d_inner}), got {output.shape}"
@@ -555,7 +563,7 @@ class TestAllToAllLayoutTransforms:
 
     def test_cp2hp_shape(self):
         """Verify _all_to_all_cp2hp output shape with identity all-to-all."""
-        from nemo_automodel.components.distributed.mamba_cp import _all_to_all_cp2hp
+        from nemo_automodel.components.distributed.context_parallel.mamba import _all_to_all_cp2hp
 
         cp_size = 2
         B, L_local, H = 2, 4, 8
@@ -563,14 +571,16 @@ def test_cp2hp_shape(self):
 
         inp = torch.randn(B, L_local, H)
 
-        with patch("nemo_automodel.components.distributed.mamba_cp._all_to_all", side_effect=lambda t, g: t):
+        with patch(
+            "nemo_automodel.components.distributed.context_parallel.mamba._all_to_all", side_effect=lambda t, g: t
+        ):
             out = _all_to_all_cp2hp(inp, pg, B)
 
         assert out.shape == (B, L_local * cp_size, H // cp_size)
 
     def test_hp2cp_shape(self):
         """Verify _all_to_all_hp2cp output shape with identity all-to-all."""
-        from nemo_automodel.components.distributed.mamba_cp import _all_to_all_hp2cp
+        from nemo_automodel.components.distributed.context_parallel.mamba import _all_to_all_hp2cp
 
         cp_size = 2
         B, L_global, H_local = 2, 8, 4
@@ -578,7 +588,9 @@ def test_hp2cp_shape(self):
 
         inp = torch.randn(B, L_global, H_local)
 
-        with patch("nemo_automodel.components.distributed.mamba_cp._all_to_all", side_effect=lambda t, g: t):
+        with patch(
+            "nemo_automodel.components.distributed.context_parallel.mamba._all_to_all", side_effect=lambda t, g: t
+        ):
             out = _all_to_all_hp2cp(inp, pg, B)
 
         assert out.shape == (B, L_global // cp_size, H_local * cp_size)
diff --git a/tests/unit_tests/models/common/test_cp_pre_embed_protocol.py b/tests/unit_tests/models/common/test_cp_pre_embed_protocol.py
index 7a79afb85f..453bd9b8bd 100644
--- a/tests/unit_tests/models/common/test_cp_pre_embed_protocol.py
+++ b/tests/unit_tests/models/common/test_cp_pre_embed_protocol.py
@@ -14,7 +14,7 @@
 
 """Signature contract for the sharder-only ``prepare_model_inputs_for_cp`` CP hook.
 
-``prepare_cp_forward`` invokes each hook-aware model directly as
+``ContextParallelSharder`` construction invokes each hook-aware model directly as
 ``model.prepare_model_inputs_for_cp(batch, num_chunks=n)`` (a plain method: the
 sharder-only hook touches no weights, so no ``__call__`` / FSDP2 unshard routing).
 Every model that answers the CP protocol must therefore expose that method with a
@@ -50,5 +50,5 @@ def test_prepare_model_inputs_for_cp_binds_dispatch_call(module_path, class_name
         pytest.fail(
             f"{class_name}.prepare_model_inputs_for_cp cannot be called as "
             f"model.prepare_model_inputs_for_cp(batch, num_chunks=n): {err}. "
-            "The CP dispatch in prepare_cp_forward calls it exactly that way."
+            "ContextParallelSharder construction calls it exactly that way."
         )
diff --git a/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py b/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py
index 0ddfd83a8c..6d8ae0aa1c 100644
--- a/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py
+++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_cp_batch.py
@@ -15,7 +15,7 @@
 """CPU unit tests for DeepSeek V4 model-owned context-parallel batch prep.
 
 Covers the model-owned CP path that runs without a real process group:
-``make_dsv4_contiguous_shard_cp_batch_and_ctx`` (the ``ContextParallelismSharder.shard_batch``
+``make_dsv4_contiguous_shard_cp_batch_and_ctx`` (the ``ContextParallelSharder.shard_batch``
 callable), the scalar group helpers, ``dsv4_cp_local_seq_multiple``, and the
 sharder-only ``DeepseekV4ForCausalLM`` CP-prep hook (``prepare_model_inputs_for_cp``).
 """
@@ -505,7 +505,7 @@ def test_prepare_model_inputs_for_cp_returns_sharder():
     prepared = DeepseekV4ForCausalLM.prepare_model_inputs_for_cp(fake_self, {"input_ids": torch.arange(8).view(1, 8)})
 
     sharder = prepared["cp_sharder"]
-    from nemo_automodel.components.distributed.cp_sharder import contiguous_local_indices
+    from nemo_automodel.components.distributed.context_parallel.sharder import contiguous_local_indices
 
     assert sharder.local_token_global_indices is contiguous_local_indices
     fn = sharder.shard_batch
@@ -539,7 +539,7 @@ def test_setup_cp_attention_stores_group():
 
 
 def test_module_exposes_pad_helper_noops():
-    from nemo_automodel.components.distributed import cp_sharder
+    from nemo_automodel.components.distributed.context_parallel import sharder as cp_sharder
 
     # pad_len <= 0 is a no-op (returns the same tensor object) for both pad helpers.
     t = torch.arange(6).view(1, 6)
diff --git a/tests/unit_tests/models/gemma4/test_gemma4_2b4b_cp.py b/tests/unit_tests/models/gemma4/test_gemma4_2b4b_cp.py
index 3c5cad1f52..8e93d29a6a 100644
--- a/tests/unit_tests/models/gemma4/test_gemma4_2b4b_cp.py
+++ b/tests/unit_tests/models/gemma4/test_gemma4_2b4b_cp.py
@@ -29,7 +29,7 @@
 import pytest
 import torch
 
-from nemo_automodel.components.distributed.cp_sharder import shard_batch_contiguous
+from nemo_automodel.components.distributed.context_parallel.sharder import shard_batch_contiguous
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.gemma4_moe.model import (
     Gemma4Config,
diff --git a/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py b/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py
index 36370dae05..0e3bc661ab 100644
--- a/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py
+++ b/tests/unit_tests/models/glm_moe_dsa/test_glm_moe_dsa_tilelang.py
@@ -1144,7 +1144,7 @@ def test_glm_dsa_prepare_model_inputs_for_cp_binds_batch_sharder():
     prepared = model.prepare_model_inputs_for_cp({"input_ids": torch.arange(8).view(1, 8)}, num_chunks=3)
 
     sharder = prepared["cp_sharder"]
-    from nemo_automodel.components.distributed.cp_sharder import contiguous_local_indices
+    from nemo_automodel.components.distributed.context_parallel.sharder import contiguous_local_indices
 
     assert sharder.local_token_global_indices is contiguous_local_indices
     fn = sharder.shard_batch
diff --git a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py
index 782ca85990..fd5d925134 100644
--- a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py
+++ b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py
@@ -176,11 +176,11 @@ def get_group(self):
 
 def test_prepare_model_inputs_for_cp_is_sharder_only(vlm_model):
     """The CP hook is sharder-only: it consumes nothing and returns a
-    ContextParallelismSharder whose shard_batch is the aux-only round-robin shard.
+    ContextParallelSharder whose shard_batch is the aux-only round-robin shard.
     Embedding + vision splice + sequence shard now run inside forward, so the raw
     batch (input_ids, pixel_values) is left intact for the forward to consume."""
-    from nemo_automodel.components.distributed.cp_sharder import (
-        ContextParallelismSharder,
+    from nemo_automodel.components.distributed.context_parallel.sharder import (
+        ContextParallelSharder,
         round_robin_local_indices,
         shard_batch_aux_only,
     )
@@ -194,7 +194,7 @@ def test_prepare_model_inputs_for_cp_is_sharder_only(vlm_model):
 
     assert set(prepared) == {"cp_sharder"}
     sharder = prepared["cp_sharder"]
-    assert isinstance(sharder, ContextParallelismSharder)
+    assert isinstance(sharder, ContextParallelSharder)
     assert sharder.shard_batch is shard_batch_aux_only
     assert sharder.local_token_global_indices is round_robin_local_indices
     # nothing consumed: the raw streams stay for the forward
diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_cp.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_cp.py
index 252fe04ace..cf3a1b267d 100644
--- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_cp.py
+++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_cp.py
@@ -129,9 +129,9 @@ def _forward_embeds(model, **forward_kwargs):
 
 def test_prepare_model_inputs_for_cp_is_sharder_only():
     """The CP hook is sharder-only: it consumes nothing and returns a
-    ContextParallelismSharder; embed + splice + shard run inside forward."""
-    from nemo_automodel.components.distributed.cp_sharder import (
-        ContextParallelismSharder,
+    ContextParallelSharder; embed + splice + shard run inside forward."""
+    from nemo_automodel.components.distributed.context_parallel.sharder import (
+        ContextParallelSharder,
         round_robin_local_indices,
         shard_batch_aux_only,
     )
@@ -143,7 +143,7 @@ def test_prepare_model_inputs_for_cp_is_sharder_only():
 
     assert set(out) == {"cp_sharder"}
     sharder = out["cp_sharder"]
-    assert isinstance(sharder, ContextParallelismSharder)
+    assert isinstance(sharder, ContextParallelSharder)
     assert sharder.shard_batch is shard_batch_aux_only
     assert sharder.local_token_global_indices is round_robin_local_indices
     # nothing consumed: input_ids / media stay for the forward
diff --git a/tests/unit_tests/models/qwen3_5/test_qwen3_5_cp_preembed.py b/tests/unit_tests/models/qwen3_5/test_qwen3_5_cp_preembed.py
index bc4efdadbb..18927ce928 100644
--- a/tests/unit_tests/models/qwen3_5/test_qwen3_5_cp_preembed.py
+++ b/tests/unit_tests/models/qwen3_5/test_qwen3_5_cp_preembed.py
@@ -74,8 +74,8 @@ def test_requires_input_ids(self):
     def test_returns_sharder_and_positions_only(self):
         """Sharder-only hook: no inputs_embeds (the forward embeds), full mRoPE
         positions returned for the aux shard, mm_token_type_ids consumed."""
-        from nemo_automodel.components.distributed.cp_sharder import (
-            ContextParallelismSharder,
+        from nemo_automodel.components.distributed.context_parallel.sharder import (
+            ContextParallelSharder,
             round_robin_local_indices,
             shard_batch_aux_only,
         )
@@ -85,7 +85,7 @@ def test_returns_sharder_and_positions_only(self):
 
         assert "inputs_embeds" not in out  # embedding happens in forward now
         sharder = out["cp_sharder"]
-        assert isinstance(sharder, ContextParallelismSharder)
+        assert isinstance(sharder, ContextParallelSharder)
         assert sharder.shard_batch is shard_batch_aux_only
         assert sharder.local_token_global_indices is round_robin_local_indices
         # position_ids came from get_rope_index (mRoPE [3, B, S]); aux shard slices it.
diff --git a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py
index bc8125f48b..295865be13 100644
--- a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py
+++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_cp_preembed.py
@@ -69,8 +69,8 @@ def test_requires_input_ids(self):
     def test_returns_sharder_and_positions_only(self):
         """Sharder-only hook: no inputs_embeds (the forward embeds), full mRoPE
         positions returned for the aux shard, mm_token_type_ids consumed."""
-        from nemo_automodel.components.distributed.cp_sharder import (
-            ContextParallelismSharder,
+        from nemo_automodel.components.distributed.context_parallel.sharder import (
+            ContextParallelSharder,
             round_robin_local_indices,
             shard_batch_aux_only,
         )
@@ -80,7 +80,7 @@ def test_returns_sharder_and_positions_only(self):
 
         assert "inputs_embeds" not in out  # embedding happens in forward now
         sharder = out["cp_sharder"]
-        assert isinstance(sharder, ContextParallelismSharder)
+        assert isinstance(sharder, ContextParallelSharder)
         assert sharder.shard_batch is shard_batch_aux_only
         assert sharder.local_token_global_indices is round_robin_local_indices
         assert out["position_ids"].shape == (3, 1, 4)  # mRoPE [3, B, S]
diff --git a/tests/unit_tests/models/step3p7/test_model.py b/tests/unit_tests/models/step3p7/test_model.py
index 1b2e117498..6b88b7f9f1 100644
--- a/tests/unit_tests/models/step3p7/test_model.py
+++ b/tests/unit_tests/models/step3p7/test_model.py
@@ -289,8 +289,8 @@ def test_from_pretrained_uses_step3p7_config(monkeypatch):
 
 
 def test_prepare_model_inputs_for_cp_is_sharder_only():
-    from nemo_automodel.components.distributed.cp_sharder import (
-        ContextParallelismSharder,
+    from nemo_automodel.components.distributed.context_parallel.sharder import (
+        ContextParallelSharder,
         round_robin_local_indices,
         shard_batch_aux_only,
     )
@@ -303,7 +303,7 @@ def test_prepare_model_inputs_for_cp_is_sharder_only():
     assert "inputs_embeds" not in result
     assert set(result) == {"cp_sharder"}
     sharder = result["cp_sharder"]
-    assert isinstance(sharder, ContextParallelismSharder)
+    assert isinstance(sharder, ContextParallelSharder)
     assert sharder.shard_batch is shard_batch_aux_only
     assert sharder.local_token_global_indices is round_robin_local_indices
 
diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py
index fd3cda8099..0f5aaef8b1 100644
--- a/tests/unit_tests/moe/test_parallelizer.py
+++ b/tests/unit_tests/moe/test_parallelizer.py
@@ -2636,10 +2636,10 @@ def __init__(self, attn_module, moe=None, layer_type=None, attention_type=None):
 
 
 def _stub_dense_cp_hooks(monkeypatch):
-    cp_utils_stub = types.ModuleType("nemo_automodel.components.distributed.cp_utils")
+    cp_utils_stub = types.ModuleType("nemo_automodel.components.distributed.context_parallel.utils")
     cp_utils_stub.attach_context_parallel_hooks = MagicMock()
     cp_utils_stub.attach_cp_sdpa_hooks = MagicMock()
-    monkeypatch.setitem(sys.modules, "nemo_automodel.components.distributed.cp_utils", cp_utils_stub)
+    monkeypatch.setitem(sys.modules, "nemo_automodel.components.distributed.context_parallel.utils", cp_utils_stub)
     return cp_utils_stub
 
 
diff --git a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py
index b988104fce..f3623746a4 100644
--- a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py
+++ b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py
@@ -19,12 +19,12 @@
 full recipe — exercising the code shape that gets shipped:
 
   - Invoke the sharder-only ``prepare_model_inputs_for_cp`` directly through
-    ``prepare_cp_forward`` (a plain method call; nothing consumed, so input_ids
+    ``ContextParallelSharder`` construction (a plain method call; nothing consumed, so input_ids
     and multimodal inputs stay in the batch for the model's own forward)
   - PP gating: the sharder-only hook is invoked on every stage (all PP-capable
     VLMs are sunk — they embed + shard in their own forward); media is dropped on
     non-first stages so those stage forwards see only text inputs
-  - Validation: count labels after _make_cp_batch_and_ctx and inside train_ctx
+  - Validation: count labels after ``ContextParallelSharder.shard`` and inside train_ctx
   - Validation: position_ids ``.to(self.dist_env.device)`` (not model.device)
 """
 
@@ -38,10 +38,24 @@
 
 import nemo_automodel.recipes.vlm.finetune as vlm_finetune
 from nemo_automodel.components.config.loader import ConfigNode
-from nemo_automodel.components.distributed import cp_utils as cp_utils_mod
 from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM
 
 
+def _identity_cp_shard(sharder, batch):
+    """Bypass CP transport while preserving constructor-side strategy resolution.
+
+    Args:
+        sharder: Sharder whose resolved strategy is not exercised by this test.
+        batch: Mutable model-input mapping whose tensor values retain their
+            existing shapes.
+
+    Returns:
+        The null context factory and the same input mapping.
+    """
+    del sharder
+    return nullcontext, batch
+
+
 def _make_recipe_with_pp_stages(*, pp_enabled=True, has_first_stage=True, pp_microbatch_size=2):
     first_stage = SimpleNamespace(is_first=True, inputs_meta=("old-first",))
     later_stage = SimpleNamespace(is_first=False, inputs_meta=("old-later",))
@@ -138,11 +152,22 @@ def test_forward_backward_step_pp_cp_first_stage_sunk_keeps_input_ids_full(monke
     }
     seen_cp_batch = {}
 
-    def _make_cp_batch_and_ctx(device_mesh, cp_batch, *args, **kwargs):
+    def _shard(sharder, cp_batch):
+        """Capture the global model-input mapping before CP transport.
+
+        Args:
+            sharder: Sharder configured by the VLM recipe.
+            cp_batch: Mutable model-input mapping whose tensor values have
+                global batch and sequence extents.
+
+        Returns:
+            The null context factory and the same input mapping.
+        """
+        del sharder
         seen_cp_batch.update(cp_batch)
-        return nullcontext, cp_batch, None
+        return nullcontext, cp_batch
 
-    monkeypatch.setattr(cp_utils_mod, "_make_cp_batch_and_ctx", _make_cp_batch_and_ctx)
+    monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _shard)
     monkeypatch.setattr(vlm_finetune, "stage_vlm_media_for_pp", lambda *args, **kwargs: nullcontext())
     monkeypatch.setattr(FinetuneRecipeForVLM, "_maybe_set_pp_first_stage_embed_input_meta", lambda self, mi: None)
 
@@ -211,11 +236,22 @@ def _run_nonfirst_stage_fbstep(monkeypatch, model):
     }
     seen_cp_batch = {}
 
-    def _make_cp_batch_and_ctx(device_mesh, cp_batch, *args, **kwargs):
+    def _shard(sharder, cp_batch):
+        """Capture the global model-input mapping before CP transport.
+
+        Args:
+            sharder: Sharder configured by the VLM recipe.
+            cp_batch: Mutable model-input mapping whose tensor values have
+                global batch and sequence extents.
+
+        Returns:
+            The null context factory and the same input mapping.
+        """
+        del sharder
         seen_cp_batch.update(cp_batch)
-        return nullcontext, cp_batch, None
+        return nullcontext, cp_batch
 
-    monkeypatch.setattr(cp_utils_mod, "_make_cp_batch_and_ctx", _make_cp_batch_and_ctx)
+    monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _shard)
     monkeypatch.setattr(vlm_finetune, "stage_vlm_media_for_pp", lambda *args, **kwargs: nullcontext())
     monkeypatch.setattr(FinetuneRecipeForVLM, "_maybe_set_pp_first_stage_embed_input_meta", lambda self, mi: None)
 
@@ -481,7 +517,7 @@ def test_run_validation_epoch_does_not_sum_tokens_over_cp(monkeypatch):
 
     # No-op replacements for the heavy collaborators.
     monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda *a, **k: nullcontext())
-    monkeypatch.setattr(cp_utils_mod, "_make_cp_batch_and_ctx", lambda mesh, batch, *a, **k: (nullcontext, batch, None))
+    monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _identity_cp_shard)
     monkeypatch.setattr(vlm_finetune, "filter_forward_kwargs", lambda model, batch: batch)
     monkeypatch.setattr(vlm_finetune, "calculate_loss", lambda *a, **k: torch.tensor(2.0))
 
@@ -534,7 +570,7 @@ def test_run_validation_epoch_cp_active_runs_pre_embed(monkeypatch):
     from nemo_automodel.recipes.vlm.finetune import FinetuneRecipeForVLM
 
     monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda *a, **k: nullcontext())
-    monkeypatch.setattr(cp_utils_mod, "_make_cp_batch_and_ctx", lambda mesh, batch, *a, **k: (nullcontext, batch, None))
+    monkeypatch.setattr(vlm_finetune.ContextParallelSharder, "shard", _identity_cp_shard)
     monkeypatch.setattr(vlm_finetune, "filter_forward_kwargs", lambda model, batch: batch)
     monkeypatch.setattr(vlm_finetune, "calculate_loss", lambda *a, **k: torch.tensor(2.0))
 
diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py
index 756d44267a..46e8891b6b 100644
--- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py
+++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py
@@ -421,7 +421,7 @@ def fake_calculate_loss(*args, **kwargs):
         return torch.tensor(1.0, requires_grad=True)
 
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
     )
     monkeypatch.setattr(
@@ -476,9 +476,10 @@ def test_forward_backward_step_routes_thd_batch_through_te(monkeypatch):
 
     def make_thd_batch(model, device_mesh, batch, **kwargs):
         captured.update(kwargs)
-        return nullcontext, batch, None
+        captured["qkv_format"] = batch.get("qkv_format")
+        return SimpleNamespace(shard=lambda actual: (nullcontext, actual))
 
-    monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.prepare_cp_forward", make_thd_batch)
+    monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.ContextParallelSharder", make_thd_batch)
     monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.get_sync_ctx", lambda *args, **kwargs: nullcontext())
     monkeypatch.setattr(
         "nemo_automodel.recipes.vlm.finetune.calculate_loss",
@@ -497,7 +498,9 @@ def make_thd_batch(model, device_mesh, batch, **kwargs):
         num_batches=1,
     )
 
-    assert captured["use_te"] is True
+    assert captured["qkv_format"] == "thd"
+    assert "use_te" not in captured
+    assert "magi" not in captured
     assert captured["padding_token_id"] == 7
 
 
@@ -1578,7 +1581,7 @@ def test_pp_skips_validation_forward(self, pp_recipe, monkeypatch):
         pp_recipe.pp = _MockAutoPipeline()
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -1606,7 +1609,7 @@ def test_pp_vlm_chunking_equal_images_and_batch(self, pp_recipe, monkeypatch):
         pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2)
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -1667,7 +1670,7 @@ def test_pp_vlm_chunking_videos_uses_video_grid_and_counts(self, pp_recipe, monk
         pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2)
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -1722,7 +1725,7 @@ def test_pp_vlm_chunking_image_and_video_mixed(self, pp_recipe, monkeypatch):
         pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2)
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -1806,7 +1809,7 @@ def test_pp_vlm_chunking_with_image_grid_thw(self, pp_recipe, monkeypatch):
         pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2)
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -1851,7 +1854,7 @@ def test_pp_vlm_chunking_qwen35_ep4_pp2_local_batch_images(self, pp_recipe, monk
         pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2)
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -1917,7 +1920,7 @@ def test_pp_vlm_chunking_with_image_sizes(self, pp_recipe, monkeypatch):
         pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2)
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -1957,7 +1960,7 @@ def test_pp_vlm_chunking_4d_pixel_values(self, pp_recipe, monkeypatch):
         pp_recipe.pp = _MockAutoPipeline(has_first_stage=True, has_last_stage=True, n_microbatches=2)
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -2005,7 +2008,7 @@ def mock_schedule_step(*args, **kwargs):
         pp_recipe.pp = pp
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -2034,7 +2037,7 @@ def test_pp_non_last_stage_returns_zero_loss(self, pp_recipe, monkeypatch):
         pp_recipe.pp = pp
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -2071,7 +2074,7 @@ def mock_schedule_step(*args, **kwargs):
         pp_recipe.pp = pp
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -2237,7 +2240,7 @@ def prepare_model_inputs_for_cp(self, batch, *, num_chunks=1):
         return {}
 
     def forward(self, **kwargs):
-        raise AssertionError("forward should not run: _make_cp_batch_and_ctx raises first")
+        raise AssertionError("forward should not run: ContextParallelSharder.shard raises first")
 
 
 class _CPPreEmbedStop(RuntimeError):
@@ -2250,7 +2253,7 @@ class TestForwardBackwardStepNonPP:
     def test_non_pp_cp_invokes_sharder_only_hook_and_keeps_inputs(self, monkeypatch):
         # Sunk contract: the non-PP CP path invokes the sharder-only hook, which
         # consumes nothing, so input_ids / pixel_values / mm_token_type_ids all
-        # reach _make_cp_batch_and_ctx intact (the model embeds + shards them in
+        # reach ContextParallelSharder.shard intact (the model embeds + shards them in
         # its own forward, not here).
         model = _CPPreEmbedModel()
         non_pp_recipe = _create_non_pp_recipe(model)
@@ -2258,7 +2261,15 @@ def test_non_pp_cp_invokes_sharder_only_hook_and_keeps_inputs(self, monkeypatch)
 
         mm_token_type_ids = torch.tensor([[1, 1, 0, 0]])
 
-        def _capture_cp_batch(device_mesh, batch, loss_mask=None, **kwargs):
+        def _capture_cp_batch(sharder, batch):
+            """Validate the global model-input mapping before CP transport.
+
+            Args:
+                sharder: Sharder configured by the VLM recipe.
+                batch: Mutable model-input mapping whose tensor values have
+                    global batch and sequence extents.
+            """
+            del sharder
             assert "input_ids" in batch
             assert "pixel_values" in batch
             assert "mm_token_type_ids" in batch
@@ -2267,7 +2278,7 @@ def _capture_cp_batch(device_mesh, batch, loss_mask=None, **kwargs):
             raise _CPPreEmbedStop
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.recipes.vlm.finetune.ContextParallelSharder.shard",
             _capture_cp_batch,
         )
 
@@ -2329,7 +2340,7 @@ def get_output_embeddings(self):
         non_pp_recipe.__dict__["loss_fn"] = FusedLinearCrossEntropy()
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
         monkeypatch.setattr(
@@ -2376,7 +2387,7 @@ def forward(self, logits_to_keep=None, **kwargs):
         non_pp_recipe.__dict__["loss_fn"] = FusedLinearCrossEntropy()
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
         monkeypatch.setattr(
@@ -2419,7 +2430,7 @@ def forward(self, **kwargs):
         non_pp_recipe.__dict__["loss_fn"] = MaskedCrossEntropy()
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
         monkeypatch.setattr(
@@ -2462,7 +2473,7 @@ def forward(self, **kwargs):
         non_pp_recipe.__dict__["loss_fn"] = MaskedCrossEntropy()
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
@@ -2496,7 +2507,7 @@ def forward(self, **kwargs):
         non_pp_recipe.__dict__["loss_fn"] = MaskedCrossEntropy()
 
         monkeypatch.setattr(
-            "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+            "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
             lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None),
         )
 
diff --git a/tests/unit_tests/recipes/test_kd_separate_mesh_recipe_helpers.py b/tests/unit_tests/recipes/test_kd_separate_mesh_recipe_helpers.py
index 11c8cf77e6..650ab52df9 100644
--- a/tests/unit_tests/recipes/test_kd_separate_mesh_recipe_helpers.py
+++ b/tests/unit_tests/recipes/test_kd_separate_mesh_recipe_helpers.py
@@ -178,7 +178,9 @@ def __exit__(self, exc_type, exc_value, traceback):
 @pytest.mark.parametrize("recipe_module,recipe_cls,_", _RECIPE_CASES)
 def test_teacher_forward_separate_materializes_logits(monkeypatch, recipe_module, recipe_cls, _):
     monkeypatch.setattr(
-        recipe_module, "prepare_cp_forward", lambda model, mesh, batch, **kwargs: (nullcontext, batch, None)
+        recipe_module,
+        "ContextParallelSharder",
+        lambda model, mesh, batch, **kwargs: SimpleNamespace(shard=lambda actual: (nullcontext, actual)),
     )
     materialized = []
     monkeypatch.setattr(
diff --git a/tests/unit_tests/recipes/test_train_ft.py b/tests/unit_tests/recipes/test_train_ft.py
index f1318135fa..860a3c0765 100644
--- a/tests/unit_tests/recipes/test_train_ft.py
+++ b/tests/unit_tests/recipes/test_train_ft.py
@@ -325,33 +325,69 @@ def test_validation_dataloaders_skip_packing_without_pack_size():
 @pytest.mark.parametrize("attn", ["magi", "te", "sdpa"])
 def test_should_pack_validation_for_explicit_thd_collater(attn):
     collate_fn = "nemo_automodel.components.datasets.utils.packed_sequence_thd_collater"
+    dataset = {"_target_": "tests.unit_tests.recipes.test_train_ft.DummyMapDataset"}
     cfg = RecipeConfig(
         ConfigNode(
             {
                 "model": {"backend": {"attn": attn}},
+                "dataset": dataset,
                 "dataloader": {"collate_fn": collate_fn},
+                "validation_dataset": dataset,
                 "validation_dataloader": {"collate_fn": collate_fn},
                 "packed_sequence": {"packed_sequence_size": 1024},
             }
         )
     )
 
-    assert _should_pack_validation(cfg, nn.Module()) is True
+    assert _should_pack_validation(cfg.dataloader, cfg.validation_dataloaders["default"], nn.Module()) is True
 
 
 def test_should_not_pack_validation_without_pack_size():
-    cfg = RecipeConfig(ConfigNode({"model": {}, "packed_sequence": {"packed_sequence_size": 0}}))
+    dataset = {"_target_": "tests.unit_tests.recipes.test_train_ft.DummyMapDataset"}
+    cfg = RecipeConfig(
+        ConfigNode(
+            {
+                "dataset": dataset,
+                "validation_dataset": dataset,
+                "packed_sequence": {"packed_sequence_size": 0},
+            }
+        )
+    )
 
-    assert _should_pack_validation(cfg, nn.Module()) is False
+    assert _should_pack_validation(cfg.dataloader, cfg.validation_dataloaders["default"], nn.Module()) is False
+
+
+@pytest.mark.parametrize(("attn", "expected"), [("magi", True), ("te", True), ("sdpa", False)])
+def test_should_pack_validation_for_live_model_backend(attn, expected):
+    collate_fn = "nemo_automodel.components.datasets.utils.packed_sequence_thd_collater"
+    dataset = {"_target_": "tests.unit_tests.recipes.test_train_ft.DummyMapDataset"}
+    cfg = RecipeConfig(
+        ConfigNode(
+            {
+                "dataset": dataset,
+                "dataloader": {"collate_fn": collate_fn},
+                "validation_dataset": dataset,
+                "validation_dataloader": {},
+                "packed_sequence": {"packed_sequence_size": 1024},
+            }
+        )
+    )
+    model = nn.Module()
+    model.backend = SimpleNamespace(attn=attn)
+
+    assert _should_pack_validation(cfg.dataloader, cfg.validation_dataloaders["default"], model) is expected
 
 
 def test_should_pack_validation_when_model_requires_training_layout():
     collate_fn = "nemo_automodel.components.datasets.utils.packed_sequence_thd_collater"
+    dataset = {"_target_": "tests.unit_tests.recipes.test_train_ft.DummyMapDataset"}
     cfg = RecipeConfig(
         ConfigNode(
             {
                 "model": {"backend": {"attn": "sdpa"}},
+                "dataset": dataset,
                 "dataloader": {"collate_fn": collate_fn},
+                "validation_dataset": dataset,
                 "validation_dataloader": {},
                 "packed_sequence": {"packed_sequence_size": 1024},
             }
@@ -362,7 +398,14 @@ class ModelRequiresPackedValidation(nn.Module):
         def should_pack_validation_with_training(self):
             return True
 
-    assert _should_pack_validation(cfg, ModelRequiresPackedValidation()) is True
+    assert (
+        _should_pack_validation(
+            cfg.dataloader,
+            cfg.validation_dataloaders["default"],
+            ModelRequiresPackedValidation(),
+        )
+        is True
+    )
 
 
 class DummyLinear(nn.Module):
@@ -1251,10 +1294,6 @@ def _create_minimal_recipe_for_pp_test(monkeypatch, pp_info):
     )
     monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.setup_logging", lambda: None)
 
-    # Mock helper functions to avoid needing full config
-    monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._uses_te_dot_product_attention", lambda cfg: False)
-    monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._uses_thd_collater", lambda cfg: False)
-
     # Create the recipe without calling setup
     recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg)
 
@@ -1263,7 +1302,16 @@ def _create_minimal_recipe_for_pp_test(monkeypatch, pp_info):
     object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True))
     object.__setattr__(recipe, "device_mesh", None)
     object.__setattr__(recipe, "pp_enabled", True)
-    object.__setattr__(recipe, "pp", SimpleNamespace(info=pp_info, update_seq_len=lambda seq_len: None))
+    object.__setattr__(
+        recipe,
+        "pp",
+        SimpleNamespace(
+            info=pp_info,
+            pp_batch_size=1,
+            pp_microbatch_size=1,
+            update_seq_len=lambda seq_len: None,
+        ),
+    )
     object.__setattr__(recipe, "tokenizer", SimpleNamespace(pad_token_id=0))
     object.__setattr__(recipe, "te_fp8", None)
 
@@ -1279,7 +1327,7 @@ def test_forward_backward_step_pp_uses_eval_for_validation(monkeypatch):
 
     # Mock _make_cp_batch_and_ctx to return a no-op context manager
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None),
     )
 
@@ -1313,7 +1361,7 @@ def test_forward_backward_step_pp_uses_step_for_training(monkeypatch):
 
     # Mock _make_cp_batch_and_ctx to return a no-op context manager
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None),
     )
 
@@ -1347,7 +1395,7 @@ def test_forward_backward_step_pp_non_first_stage_uses_eval_for_validation(monke
 
     # Mock _make_cp_batch_and_ctx to return a no-op context manager
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None),
     )
 
@@ -1383,7 +1431,7 @@ def test_forward_backward_step_pp_non_first_stage_uses_step_for_training(monkeyp
 
     # Mock _make_cp_batch_and_ctx to return a no-op context manager
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None),
     )
 
@@ -1445,7 +1493,7 @@ def mock_dp_allreduce(val, include_cp=False):
 
     # Mock _make_cp_batch_and_ctx
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None),
     )
 
@@ -1505,7 +1553,7 @@ def mock_dp_allreduce(val, include_cp=False):
     monkeypatch.setattr(recipe, "_dp_allreduce", mock_dp_allreduce)
 
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *args, **kwargs: (nullcontext, batch, None),
     )
 
@@ -1981,9 +2029,6 @@ def _make_recipe(self, monkeypatch, pp_enabled, dp_group_size=4):
             lambda *a, **k: SimpleNamespace(world_size=1, is_main=True, device=torch.device("cpu"), rank=0),
         )
         monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.setup_logging", lambda: None)
-        monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._uses_te_dot_product_attention", lambda cfg: False)
-        monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._uses_thd_collater", lambda cfg: False)
-
         recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg)
 
         object.__setattr__(recipe, "dist_env", SimpleNamespace(device=torch.device("cpu"), rank=0, is_main=True))
@@ -2377,8 +2422,6 @@ def test_forward_backward_step_model_cp_hook(monkeypatch, cp_size, uses_thd, sup
         lambda *a, **k: SimpleNamespace(world_size=1, is_main=True, device=torch.device("cpu"), rank=0),
     )
     monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.setup_logging", lambda: None)
-    monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._uses_te_dot_product_attention", lambda cfg: False)
-    monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft._uses_thd_collater", lambda cfg: uses_thd)
     recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg)
 
     class _CPModel(nn.Module):
@@ -2390,13 +2433,13 @@ def __init__(self):
         def prepare_model_inputs_for_cp(self, batch, **kwargs):
             self.prepared = True
             self.num_chunks = kwargs.get("num_chunks")
-            from nemo_automodel.components.distributed.cp_sharder import (
-                ContextParallelismSharder,
+            from nemo_automodel.components.distributed.context_parallel.sharder import (
+                ContextParallelSharder,
                 contiguous_local_indices,
             )
 
             return {
-                "cp_sharder": ContextParallelismSharder(
+                "cp_sharder": ContextParallelSharder(
                     shard_batch=lambda cp_mesh, tp_mesh, batch, **k: (nullcontext, batch, None),
                     local_token_global_indices=contiguous_local_indices,
                 )
@@ -2445,7 +2488,7 @@ def _fake_calc_loss(loss_fn, *, logits, labels, model, hidden_states, lm_weight,
         return logits.mean()
 
     monkeypatch.setattr(
-        "nemo_automodel.components.distributed.cp_utils._make_cp_batch_and_ctx",
+        "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx",
         lambda device_mesh, batch, *a, **k: (nullcontext, batch, None),
     )
     monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.calculate_loss", _fake_calc_loss)
@@ -2454,6 +2497,8 @@ def _fake_calc_loss(loss_fn, *, logits, labels, model, hidden_states, lm_weight,
     monkeypatch.setattr("nemo_automodel.recipes.llm.train_ft.filter_forward_kwargs", lambda model, batch: batch)
 
     batch = {"input_ids": torch.randn(1, 4, 4), "labels": torch.zeros(1, 4, dtype=torch.long)}
+    if uses_thd:
+        batch["qkv_format"] = "thd"
     loss_buffer = []
     recipe._forward_backward_step(
         idx=0, batch=batch, loss_buffer=loss_buffer, num_label_tokens=None, num_batches=1, is_train=True
diff --git a/tests/unit_tests/recipes/test_vlm_kd_tp_cp_correctness.py b/tests/unit_tests/recipes/test_vlm_kd_tp_cp_correctness.py
index 490ae38fcd..e9800cffb0 100644
--- a/tests/unit_tests/recipes/test_vlm_kd_tp_cp_correctness.py
+++ b/tests/unit_tests/recipes/test_vlm_kd_tp_cp_correctness.py
@@ -21,7 +21,6 @@
 import torch
 import torch.nn as nn
 
-from nemo_automodel.components.distributed import cp_utils as cp_utils_mod
 from nemo_automodel.components.loss import kd_loss as kd_loss_module
 from nemo_automodel.components.loss.kd_loss import KDLoss
 from nemo_automodel.recipes.vlm import kd as vlm_kd
@@ -162,11 +161,21 @@ def wrapped_kl_forward_tp(t_logits, s_logits, tp_group):
 def test_vlm_kd_cp_prepare_shards_input_ids_and_teacher_embeds_them(monkeypatch):
     make_cp_calls = []
 
-    def fake_make_cp_batch_and_ctx(device_mesh, batch, *args, **kwargs):
-        make_cp_calls.append((device_mesh, dict(batch)))
-        return nullcontext, batch, None
+    def fake_shard(sharder, batch):
+        """Capture the unsharded model-input mapping at the public CP seam.
 
-    monkeypatch.setattr(cp_utils_mod, "_make_cp_batch_and_ctx", fake_make_cp_batch_and_ctx)
+        Args:
+            sharder: Sharder configured with the recipe's device mesh.
+            batch: Mutable model-input mapping whose tensors have global batch
+                and sequence extents.
+
+        Returns:
+            The null context factory and the same input mapping.
+        """
+        make_cp_calls.append((sharder, dict(batch)))
+        return nullcontext, batch
+
+    monkeypatch.setattr(vlm_kd.ContextParallelSharder, "shard", fake_shard)
 
     student = _StudentVLM(hidden_size=8)
     teacher = _TeacherVLM(hidden_size=8)
diff --git a/tools/dsv4_cp_loss_parity.py b/tools/dsv4_cp_loss_parity.py
index 93323c91ae..ce0b0b4269 100644
--- a/tools/dsv4_cp_loss_parity.py
+++ b/tools/dsv4_cp_loss_parity.py
@@ -32,7 +32,7 @@
 from torch.distributed.tensor import DTensor
 
 from nemo_automodel.components.distributed.config import FSDP2Config
-from nemo_automodel.components.distributed.cp_utils import make_cp_batch_and_ctx
+from nemo_automodel.components.distributed.context_parallel.utils import make_cp_batch_and_ctx
 from nemo_automodel.components.distributed.mesh import ParallelismSizes
 from nemo_automodel.components.distributed.mesh_utils import _create_device_meshes
 from nemo_automodel.components.models.common import BackendConfig