diff --git a/docs/scaling.md b/docs/scaling.md
index 229beaa8278..44bc226c57a 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -122,6 +122,16 @@ cp = 2                       # CP degree
 cp_style = "ulysses"         # "ring"
 ```
 
+RL distributes the original packed tokens into balanced CP shards. For example, nine tokens across four ranks produce lengths `[3, 2, 2, 2]`; rows shorter than the CP degree include empty shards. The same partition applies to positions, labels, temperatures, sampling masks, and router replay. Packed sample boundaries remain independent of rank boundaries. Ulysses retains its head divisibility requirements.
+
+The packer splits at whole-sample boundaries to balance distributed iterations. When a lane has no sample for an iteration, it executes an explicit zero-token slot. FSDP parameter hooks, expert communication, and vision participation still run, so other lanes can use parameters and experts owned by that rank. Empty slots contribute no loss members or sample annotations. A globally empty training step is rejected before updating the optimizer.
+
+These behaviors are enabled by `model.cp_unpadded` and `model.inactive_micro_batches`, both defaulting to `true`. Managed RL configures the orchestrator accordingly; a standalone orchestrator must use matching trainer settings. `perf/input_tokens`, `perf/real_micro_batches`, and `perf/inactive_micro_batches` report the executed data separately from empty slots. Sampling-set `-1` entries, kernel workspace alignment, and SFT dataset packing retain their own storage contracts.
+
+Changing token counts can create additional compiled variants, particularly for zero- and one-token shards. CUDA Graph capture depends on the chosen compile mode and kernel support. Measure useful-token throughput and memory on the target workload; removing input tokens does not guarantee lower step latency.
+
+Ring CP gathers a combined buffer of real keys and values before applying causal packed attention. Its temporary memory scales with the full row's KV width. Recurrent layers redistribute features or gather their real sequence as required by their kernel. Compilation still follows each kernel's existing restrictions; FA4 and multi-LoRA contain graph breaks in the current dependencies and require the default `fullgraph=false` setting.
+
 ### Activation Checkpointing and Offloading
 
 | Knob | Memory ↓ | Throughput ↓ |
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
index 0f57abcced2..867fda6851f 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
@@ -661,6 +661,9 @@ class OrchestratorConfig(BaseConfig):
     pad_to_multiple_of: int = Field(1, ge=1)
     """Pad each packed micro batch to a multiple of this value (the trainer's cp degree). Auto-filled by the ``rl`` entrypoint; set explicitly for standalone orchestrator runs with cp > 1."""
 
+    inactive_micro_batches: bool = True
+    """Send explicit zero-token slots for distributed iteration alignment. Managed RL copies this setting from the trainer."""
+
     max_steps: int | None = None
     """Maximum training steps. If None, runs indefinitely."""
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
index 65e6ed74c92..a55438f330c 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py
@@ -686,7 +686,8 @@ def validate_mooncake_offload_requires_slurm(self):
 
     @model_validator(mode="after")
     def auto_setup_deployment(self):
-        self.orchestrator.pad_to_multiple_of = self.trainer.model.cp
+        self.orchestrator.pad_to_multiple_of = 1 if self.trainer.model.cp_unpadded else self.trainer.model.cp
+        self.orchestrator.inactive_micro_batches = self.trainer.model.inactive_micro_batches
         if self.deployment.type == "single_node":  # single-node
             # set num_train_workers to the number of data replicas
             non_data_parallel_size = self.trainer.model.cp
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 70100bee117..7acd267e41a 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -315,6 +315,12 @@ class ModelConfig(BaseModelConfig):
     cp: int = 1
     """Context parallelism degree. 1 disables CP."""
 
+    cp_unpadded: bool = True
+    """Distribute original tokens into balanced CP shards, including empty shards for short rows."""
+
+    inactive_micro_batches: bool = True
+    """Align distributed iterations with zero-token slots that retain parameter and expert participation. Requires cp_unpadded when CP is enabled."""
+
     cp_style: Literal["ring", "ulysses"] = "ring"
     """CP communication style. ``ring`` uses ring-attention all-gather/reduce-scatter (requires custom kernels per attention type). ``ulysses`` uses all-to-all to redistribute Q/K/V from sequence-sharded to head-sharded, runs vanilla attention locally on the full sequence, then all-to-all back — works out-of-the-box with any attention kernel (softmax FA, linear attention, mamba, etc.)."""
 
@@ -369,6 +375,8 @@ def vlm_cp_requires_ulysses(self):
 
     @model_validator(mode="after")
     def validate_cp(self):
+        if self.inactive_micro_batches and self.cp > 1 and not self.cp_unpadded:
+            raise ValueError("inactive_micro_batches with CP requires cp_unpadded")
         if self.cp > 1 and self.attn not in ["flash_attention_2", "flash_attention_3", "flash_attention_4", "auto"]:
             raise ValueError("CP is only supported with flash attention 2, 3, or 4")
         if self.cp > 1 and self.impl not in ("custom", "auto"):
diff --git a/src/prime_rl/orchestrator/packing.py b/src/prime_rl/orchestrator/packing.py
index c3716eb0701..cbdc24bc803 100644
--- a/src/prime_rl/orchestrator/packing.py
+++ b/src/prime_rl/orchestrator/packing.py
@@ -13,6 +13,7 @@ def __init__(self, config: OrchestratorConfig):
         self.seq_len = config.seq_len
         self.num_train_workers = config.num_train_workers
         self.pad_to_multiple_of = config.pad_to_multiple_of
+        self.inactive_micro_batches = config.inactive_micro_batches
         try:
             model_config = AutoConfig.from_pretrained(
                 config.model.name, trust_remote_code=config.tokenizer.trust_remote_code
@@ -32,4 +33,5 @@ def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]:
             num_train_workers=self.num_train_workers,
             bin_cost=self.bin_cost,
             pad_to_multiple_of=self.pad_to_multiple_of,
+            inactive_micro_batches=self.inactive_micro_batches,
         )
diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py
index 925e65a507d..14a96549281 100644
--- a/src/prime_rl/trainer/batch.py
+++ b/src/prime_rl/trainer/batch.py
@@ -667,11 +667,19 @@ def packed_samples_into_micro_bs(
             bins.append(_MicroBatchBin.from_sample(sample))
 
     if num_train_workers > 1:
-        target_count = max(
-            ((len(bins) + num_train_workers - 1) // num_train_workers) * num_train_workers,
-            num_train_workers,
-        )
-        _expand_bins_by_splitting(bins, target_count, bin_cost)
+        while True:
+            count = len(bins)
+            groups = [[], []]
+            for bin_content in bins:
+                has_images = any(_is_multimodal_sample(sample) for sample in bin_content.samples)
+                groups[int(has_images)].append(bin_content)
+            for group in groups:
+                target_count = ((len(group) + num_train_workers - 1) // num_train_workers) * num_train_workers
+                _expand_bins_by_splitting(group, target_count, bin_cost)
+            bins = groups[1] + groups[0]
+            # Splitting a mixed bin can create a text-only bin in another execution group.
+            if len(bins) == count:
+                break
 
     return [_materialize_bin(bin_content) for bin_content in bins]
 
@@ -809,11 +817,28 @@ def _make_dummy_batch(source: MicroBatch) -> MicroBatch:
     return dummy
 
 
-def _pad_group_for_distribution(group: list[MicroBatch], num_train_workers: int) -> list[MicroBatch]:
+def _pad_group_for_distribution(
+    group: list[MicroBatch], num_train_workers: int, inactive_micro_batches: bool = False
+) -> list[MicroBatch]:
     """Pad a group of micro batches so its length is divisible by num_train_workers."""
     num_padding = -len(group) % num_train_workers
     if num_padding > 0 and len(group) > 0:
-        dummy = _make_dummy_batch(group[0])
+        dummy = (
+            MicroBatch(
+                input_ids=[],
+                loss_mask=[],
+                advantages=[],
+                inference_logprobs=[],
+                position_ids=[],
+                sequence_lengths=[],
+                temperatures=[],
+                env_names=[],
+                seq_lens=[],
+                inactive=True,
+            )
+            if inactive_micro_batches
+            else _make_dummy_batch(group[0])
+        )
         group.extend([dummy] * num_padding)
     return group
 
@@ -824,6 +849,7 @@ def prepare_batch(
     num_train_workers: int,
     bin_cost: Callable[[Sequence[int]], int],
     pad_to_multiple_of: int = 1,
+    inactive_micro_batches: bool = True,
 ) -> list[list[MicroBatch]]:
     """
     Prepare a batch of problems for each GPU. Each batch is a list of micro batches.
@@ -844,8 +870,8 @@ def prepare_batch(
     text_batches = [b for b in micro_batches if not _is_multimodal_sample(b)]
 
     # Pad each group independently so its count is divisible by num_train_workers
-    mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers)
-    text_batches = _pad_group_for_distribution(text_batches, num_train_workers)
+    mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers, inactive_micro_batches)
+    text_batches = _pad_group_for_distribution(text_batches, num_train_workers, inactive_micro_batches)
 
     # Alignment check after distribution padding so the dummy batches are covered too
     for micro_batch in (*mm_batches, *text_batches):
diff --git a/src/prime_rl/trainer/distributed/collectives.py b/src/prime_rl/trainer/distributed/collectives.py
index 5a576ca7f81..4595405ba26 100644
--- a/src/prime_rl/trainer/distributed/collectives.py
+++ b/src/prime_rl/trainer/distributed/collectives.py
@@ -3,6 +3,8 @@
 import prime_kernels
 import torch
 import torch.distributed as dist
+import torch.distributed._functional_collectives as funcol
+from torch._dynamo.comptime import comptime
 from torch.distributed import ProcessGroup
 
 
@@ -254,8 +256,34 @@ def all_gather(x: torch.Tensor, dim: int, group: ProcessGroup) -> torch.Tensor:
     return _all_gather(x, dim, group.size(), group.group_name)
 
 
+def all_to_all_cp(x: torch.Tensor, total_tokens: int, group: ProcessGroup, sequence_to_head: bool) -> torch.Tensor:
+    """Exchange flattened token rows using the balanced CP partition."""
+    degree = group.size()
+    # Unit and empty shards need concrete broadcast strides in compiled backward.
+    if total_tokens < 2 * degree:
+        comptime.force_static(total_tokens)
+    lengths = [(total_tokens + degree - 1 - rank) // degree for rank in range(degree)]
+    repeated = [lengths[group.rank()]] * degree
+    outputs, inputs = (lengths, repeated) if sequence_to_head else (repeated, lengths)
+    exchanged = torch.ops._c10d_functional.all_to_all_single(x.contiguous(), outputs, inputs, group.group_name)
+    return funcol.wait_tensor(exchanged)
+
+
+def all_gather_cp(x: torch.Tensor, dim: int, total_tokens: int, group: ProcessGroup) -> torch.Tensor:
+    """Gather balanced token shards; backward sums every peer's contribution."""
+    if total_tokens % group.size() == 0:
+        rows = x.movedim(dim, 0).contiguous()
+        gathered = funcol.wait_tensor(funcol.all_gather_single(rows, 0, group))
+        return gathered.movedim(0, dim).contiguous()
+    rows = x.movedim(dim, 0)
+    repeated = rows.repeat(group.size(), *([1] * (rows.ndim - 1)))
+    return all_to_all_cp(repeated, total_tokens, group, True).movedim(0, dim).contiguous()
+
+
 __all__ = [
     "all_gather",
+    "all_gather_cp",
+    "all_to_all_cp",
     "all_to_all_single",
     "all_to_all_single_equal",
     "mxfp8_all_to_all_combine",
diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py
index b75a6fc4ee3..a2591966bcc 100644
--- a/src/prime_rl/trainer/model.py
+++ b/src/prime_rl/trainer/model.py
@@ -271,10 +271,13 @@ def get_load_balance_stats(
         num_routed_tokens = tokens_per_expert.sum() / block_mlp.router.top_k
         tokens_per_expert = tokens_per_expert.sort(dim=0, descending=True).values[block_mlp.router.top_k :]
         balanced_load = tokens_per_expert.mean()
-        max_vio = (tokens_per_expert.max() - balanced_load) / balanced_load
+        max_vio = (tokens_per_expert.max() - balanced_load) / balanced_load.clamp_min(1e-12)
+        max_vio = torch.where(num_routed_tokens > 0, max_vio, torch.nan)
         per_layer_max_vio.append(max_vio.detach())
 
-        routing_confidence = routing_confidence_sum / num_routed_tokens
+        routing_confidence = torch.where(
+            num_routed_tokens > 0, routing_confidence_sum / num_routed_tokens.clamp_min(1), torch.nan
+        )
         per_layer_routing_confidence.append(routing_confidence.detach())
 
         block_mlp.tokens_per_expert.zero_()
@@ -295,12 +298,15 @@ def get_global_moe_stats(
     for name, values in get_load_balance_stats(model, group=ep_group).items():
         if values is None:
             continue
-        value = values.max() if name == "max_vio" else values.mean()
+        # Empty routing groups have no observation to contribute to the global statistic.
+        valid = ~values.isnan()
         if name == "max_vio":
+            value = values.masked_fill(~valid, 0.0).max()
             dist.all_reduce(value, op=dist.ReduceOp.MAX, group=dp_cp_group)
         else:
-            dist.all_reduce(value, op=dist.ReduceOp.SUM, group=dp_cp_group)
-            value /= dist.get_world_size(dp_cp_group)
+            totals = torch.stack((values.masked_fill(~valid, 0.0).sum(), valid.sum()))
+            dist.all_reduce(totals, op=dist.ReduceOp.SUM, group=dp_cp_group)
+            value = totals[0] / totals[1].clamp_min(1)
         stats[name] = value.to("cpu")
     return stats
 
@@ -1095,6 +1101,7 @@ def forward(
     # True when seq_lens holds the full pre-CP-shard document boundaries
     # (kept global because documents can straddle the shard cut).
     seq_lens_are_pre_shard: bool = False,
+    cp_total_tokens: int | None = None,
 ) -> PrimeLmOutput:
     kwargs = {
         "input_ids": input_ids,
@@ -1124,6 +1131,8 @@ def forward(
     if isinstance(model, PreTrainedModelPrimeRL):
         kwargs["seq_lens"] = seq_lens
         kwargs["seq_lens_are_pre_shard"] = seq_lens_are_pre_shard
+        if cp_total_tokens is not None:
+            kwargs["cp_total_tokens"] = cp_total_tokens
 
     if routed_experts is not None:
         kwargs["routed_experts"] = routed_experts
diff --git a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
index 278c9745468..2e57e6dad37 100644
--- a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
+++ b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py
@@ -28,6 +28,7 @@
     RotaryEmbeddingConfig,
     apply_rotary_pos_emb,
 )
+from prime_rl.utils.cp import CPContext
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
 
 from .configuration_afmoe import AfmoeConfig
@@ -94,23 +95,49 @@ class AfmoeFlashAttention(AfmoeAttentionBase):
 
     def __init__(self, config: AfmoeAttentionConfig, flash_attn_version: int = 4):
         super().__init__(config)
+        self.cp_context = CPContext()
         self._flash_attn_version = flash_attn_version
         self.func = self._funcs[flash_attn_version]
         self._flash_attn_call = self.func
         if self._flash_attn_version == 4:
             self._flash_attn_call = torch._dynamo.disable(self.func)
 
-    def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen):
-        """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim]."""
-        args = [q, k, v, cu_seqlens, cu_seqlens]
-        if self._flash_attn_version != 4:
-            args.extend([max_seqlen, max_seqlen])
-        kwargs: dict = {"causal": True}
+    def _compute_attention(self, q, k, v, cu_seqlens, max_seqlen, cp_total_tokens=None):
+        if self.cp_context.cp_enabled and cp_total_tokens is not None:
+            from prime_rl.trainer.models.layers.cp_attn import context_parallel_attention
+
+            window = (self.sliding_window - 1, 0) if self.sliding_window is not None else (-1, -1)
+            return context_parallel_attention(
+                self._flash_attn_call,
+                q,
+                k,
+                v,
+                cu_seqlens,
+                max_seqlen,
+                cp_total_tokens,
+                self.cp_context,
+                self._flash_attn_version,
+                window_size=window,
+            )
+        if q.shape[0] == 0:
+            return q + (k.sum() + v.sum()).to(q.dtype)
+        kwargs = {"causal": True}
         if self.sliding_window is not None:
             kwargs["window_size"] = (self.sliding_window - 1, 0)
-        out = self._flash_attn_call(*args, **kwargs)
-        if isinstance(out, tuple):
-            out = out[0]
+        cu_k = cu_seqlens.clone()
+        if self._flash_attn_version == 4:
+            out, _ = self._flash_attn_call(
+                q,
+                k,
+                v,
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_k,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                **kwargs,
+            )
+        else:
+            out = self._flash_attn_call(q, k, v, cu_seqlens, cu_k, max_seqlen, max_seqlen, **kwargs)
         return out
 
     def forward(
@@ -120,13 +147,13 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
     ) -> tuple[torch.Tensor, None]:
         input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
 
-        query_states = self.q_proj(hidden_states).view(hidden_shape)
-        key_states = self.k_proj(hidden_states).view(hidden_shape)
-        value_states = self.v_proj(hidden_states).view(hidden_shape)
+        query_states = self.q_proj(hidden_states).view(*input_shape, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(*input_shape, self.num_key_value_heads, self.head_dim)
         gate_states = self.gate_proj(hidden_states)
 
         query_states = self.q_norm(query_states)
@@ -140,8 +167,15 @@ def forward(
             query_states = query_states.transpose(1, 2)
             key_states = key_states.transpose(1, 2)
 
-        attn_output = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen)
-        attn_output = attn_output.contiguous().view(*input_shape, -1)
+        attn_output = self._compute_attention(
+            query_states[0],
+            key_states[0],
+            value_states[0],
+            cu_seqlens,
+            max_seqlen,
+            **({"cp_total_tokens": cp_total_tokens} if cp_total_tokens is not None else {}),
+        )
+        attn_output = attn_output.contiguous().view(*input_shape, self.num_heads * self.head_dim)
         attn_output = attn_output * torch.sigmoid(gate_states)
         return self.o_proj(attn_output), None
 
@@ -248,6 +282,7 @@ def forward(
         position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
         routed_experts: Optional[torch.LongTensor] = None,
     ) -> torch.FloatTensor:
         residual = hidden_states
@@ -259,6 +294,7 @@ def forward(
             attention_mask=attention_mask,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = residual + hidden_states
@@ -342,6 +378,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> MoeModelOutputWithPast:
         """
         routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*):
@@ -358,9 +395,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
         causal_mask_mapping = None
 
         hidden_states = inputs_embeds
@@ -380,6 +421,7 @@ def forward(
                 position_embeddings=position_embeddings,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=routed_experts_layer,
             )
 
@@ -437,6 +479,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         r"""
@@ -459,6 +502,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
         hidden_states = outputs.last_hidden_state
diff --git a/src/prime_rl/trainer/models/deepseek_v4/attention.py b/src/prime_rl/trainer/models/deepseek_v4/attention.py
index 9dfd91deeb3..6ff869e4de2 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/attention.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/attention.py
@@ -125,7 +125,6 @@
 
 import torch
 import torch.distributed as dist
-import torch.distributed._functional_collectives as funcol
 from torch import Tensor, nn
 
 from prime_rl.trainer.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config
@@ -136,7 +135,7 @@
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
 from prime_rl.utils.cp import CPContext, gather_for_cp
-from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
+from prime_rl.utils.sequence import CPPartition, get_cu_seqlens_from_seq_lens
 
 # Guarded because tilelang ships in the linux-gated `gpu` extra, so some installs lack it.
 try:
@@ -176,7 +175,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         w = self.weight.view(self.n_groups, -1, hidden_dim).transpose(1, 2)
         x = x.reshape(-1, self.n_groups, hidden_dim).transpose(0, 1)
         y = torch.bmm(x, w).transpose(0, 1)
-        return y.reshape(*input_shape, self.n_groups, -1)
+        return y.reshape(*input_shape, self.n_groups, self.out_features // self.n_groups)
 
 
 @dataclass(frozen=True)
@@ -225,7 +224,7 @@ def build(cls, *, cu_seqlens: Tensor, compress_rate: int) -> "CompressionLayout"
             entry_doc_idx=entry_doc_idx,
             entry_local_idx=entry_local_idx,
             first_entry_of_doc=first_entry_of_doc,
-            max_entries_per_doc=int(counts.max()),
+            max_entries_per_doc=int(counts.max()) if counts.numel() else 0,
         )
 
 
@@ -250,6 +249,7 @@ class PackedContext:
     tok_doc_idx: Tensor  # (n_queries,) int64 - which document each query token belongs to
     window_indices: Tensor  # (n_queries, sliding_window) int32 - global token per window slot, IGNORE_SLOT if unused
     compression_layouts: dict[int, CompressionLayout]  # keyed by compress rate
+    total_tokens: int
 
     @classmethod
     def build(
@@ -265,7 +265,7 @@ def build(
 
         `rotary_emb` supplies the sliding window, the compress rates in use and the positions its
         RoPE cache covers, all read from its config. The sequence is as long as `seq_lens` says,
-        padding included: both packers fold their padding into the last document.
+        with every token assigned to its original document.
 
         `seq_lens` always describes the whole sequence. `cp_rank` and `cp_world_size` say which
         contiguous shard of it this rank holds the queries of; the keys, the entries and the index
@@ -274,17 +274,15 @@ def build(
         config = rotary_emb.config
         # Read the width before `seq_lens` moves: on a CPU `seq_lens` that costs no device sync.
         total_tokens = int(seq_lens.sum())
-        assert total_tokens % cp_world_size == 0, (
-            f"{total_tokens} tokens do not split evenly across {cp_world_size} CP ranks"
-        )
+        partition = CPPartition(total_tokens, cp_world_size)
         # Document-local positions stay below the row's length; the RoPE kernels index the cache unchecked.
         assert total_tokens <= config.max_position_embeddings, (
             f"{total_tokens} tokens exceed the {config.max_position_embeddings} positions of the RoPE cache"
         )
-        n_queries = total_tokens // cp_world_size
-        q_start = cp_rank * n_queries
+        n_queries = partition.lengths[cp_rank]
+        q_start = partition.offsets[cp_rank]
 
-        cu_seqlens, _ = get_cu_seqlens_from_seq_lens(seq_lens.to(device=device))
+        cu_seqlens, _ = get_cu_seqlens_from_seq_lens(seq_lens.to(device=device), total_tokens=total_tokens)
         compress_rates = {
             config.compress_rates[layer_type]
             for layer_type in set(config.layer_types)
@@ -307,6 +305,7 @@ def build(
 
         return cls(
             position_ids=position_ids,
+            total_tokens=total_tokens,
             tok_doc_idx=tok_doc_idx,
             compression_layouts={
                 rate: CompressionLayout.build(cu_seqlens=cu_seqlens, compress_rate=rate) for rate in compress_rates
@@ -472,7 +471,7 @@ def compress(
         width = self.n_series * self.head_dim
         proj = torch.cat([self.kv_proj(hidden_states), self.gate_proj(hidden_states)], dim=-1)
         if cp_world_size > 1:
-            proj = gather_for_cp(proj, cp_group)
+            proj = gather_for_cp(proj, cp_group, packed.total_tokens)
         kv, gate = proj.split(width, dim=-1)
 
         kv = kv[:, layout.entry_tok_idx]
@@ -486,6 +485,8 @@ def compress(
 
         entry_first_tok_pos = layout.entry_local_idx * self.compress_rate
         cos_sin_cache = self.rotary_emb.cos_sin_cache(self.rope_layer_type)
+        if compressed.shape[1] == 0:
+            return compressed
         return dsv4_rope(compressed.unsqueeze(2), cos_sin_cache, entry_first_tok_pos).squeeze(2)
 
     def causal_threshold(self, position_ids: torch.Tensor) -> torch.Tensor:
@@ -540,7 +541,9 @@ def forward(
         n_entries = compressed_kv.shape[1]
 
         cos_sin_cache = self.compressor.rotary_emb.cos_sin_cache(self.compressor.rope_layer_type)
-        q = self.q_b_proj(q_residual).view(batch, seq_len, -1, self.head_dim)
+        q = self.q_b_proj(q_residual).view(batch, seq_len, self.num_heads, self.head_dim)
+        if seq_len == 0:
+            return torch.empty((batch, 0, self.index_topk), device=q.device, dtype=torch.int64)
         q = dsv4_rope(q, cos_sin_cache, packed.position_ids)
         w = self.weights_proj(hidden_states)
 
@@ -722,25 +725,21 @@ def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[t
         # `hidden_states` is (b, t, hidden_size).
 
         input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)  # (b, t, h, d), the query view
+        hidden_shape = (*input_shape, self.num_heads, self.head_dim)  # (b, t, h, d), the query view
         cos_sin_cache = self.rotary_emb.cos_sin_cache(self.rope_layer_type)  # (max_position, qk_rope_head_dim)
 
         kv = self.kv_norm(self.kv_proj(hidden_states))  # (b, t, d)
         kv = kv.view(*kv.shape[:2], 1, self.head_dim)  # (b, t, 1, d)
-        kv = dsv4_rope(kv, cos_sin_cache, packed.position_ids)
+        if kv.shape[1]:
+            kv = dsv4_rope(kv, cos_sin_cache, packed.position_ids)
         if self.cp_context.cp_enabled:
-            # Launch on NCCL's communication stream; query/compressor work does not read KV.
-            kv = torch.ops._c10d_functional.all_gather_into_tensor(
-                kv.movedim(1, 0).contiguous(),
-                self.cp_context.cp_world_size,
-                self.cp_context.cp_group.group_name,
-            )
+            kv = gather_for_cp(kv, self.cp_context.cp_group, packed.total_tokens)
 
         q_residual = self.q_a_norm(self.q_a_proj(hidden_states))  # (b, t, r)
         # Keep the query in the sparse kernel's (batch, tokens, heads, dim) layout.
-        q = dsv4_q_norm_rope(
-            self.q_b_proj(q_residual).view(*hidden_shape), cos_sin_cache, packed.position_ids, self.q_b_norm.eps
-        )  # (b, t, h, d)
+        q = self.q_b_proj(q_residual).view(*hidden_shape)
+        if q.shape[1]:
+            q = dsv4_q_norm_rope(q, cos_sin_cache, packed.position_ids, self.q_b_norm.eps)
 
         compressed = (
             self.compressor(
@@ -754,8 +753,6 @@ def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[t
             else None
         )
         compressed_kv, top_k_indices = compressed if compressed is not None else (None, None)
-        if self.cp_context.cp_enabled:
-            kv = funcol.wait_tensor(kv).movedim(0, 1).contiguous()  # (b, T, 1, d)
         kv = kv.transpose(1, 2)  # (b, 1, T, d)
         inputs = SparseAttnInputs.build(
             kv=kv,
@@ -763,20 +760,28 @@ def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[t
             top_k_indices=top_k_indices,
             window_indices=packed.window_indices,
         )
-        attn_output, _ = dsv4_sparse_attn(
-            q,
-            inputs.kv_buf,
-            inputs.indices,
-            self.sinks,
-            self.scaling,
-        )  # (b, t, h, d)
+        if q.shape[1] == 0:
+            attn_output = q + (inputs.kv_buf.sum() + self.sinks.sum()).to(q.dtype)
+        else:
+            attn_output, _ = dsv4_sparse_attn(
+                q,
+                inputs.kv_buf,
+                inputs.indices,
+                self.sinks,
+                self.scaling,
+            )  # (b, t, h, d)
 
         # The value stream is the key stream, so it arrived rotated. Rotating the output
         # by the conjugate angle at the query position cancels that out.
-        attn_output = dsv4_rope(attn_output, cos_sin_cache, packed.position_ids, inverse=True)
+        if attn_output.shape[1]:
+            attn_output = dsv4_rope(attn_output, cos_sin_cache, packed.position_ids, inverse=True)
 
         # (b, t, g, h * d // g) -> (b, t, g, l) -> (b, t, g * l)
-        grouped = self.o_a_proj(attn_output.reshape(*input_shape, self.config.o_groups, -1)).flatten(2)
+        grouped = self.o_a_proj(
+            attn_output.reshape(
+                *input_shape, self.config.o_groups, self.num_heads * self.head_dim // self.config.o_groups
+            )
+        ).flatten(2)
         return self.o_b_proj(grouped), None  # (b, t, hidden_size)
 
     def init_weights(self, init_std: float) -> None:
diff --git a/src/prime_rl/trainer/models/deepseek_v4/hyperconnections.py b/src/prime_rl/trainer/models/deepseek_v4/hyperconnections.py
index 925e99aec82..8ae4e257c48 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/hyperconnections.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/hyperconnections.py
@@ -18,7 +18,7 @@ def __init__(self, eps: float = 1e-6, out_dtype: torch.dtype | None = None):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         out_dtype = self.out_dtype if self.out_dtype is not None else x.dtype
         quack_rms = norms.get_quack_rmsnorm()
-        if quack_rms is not None:
+        if quack_rms is not None and x.numel() > 0:
             return quack_rms(x, eps=self.eps, out_dtype=out_dtype)
         x = x.float()
         return (x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)).to(out_dtype)
@@ -64,7 +64,11 @@ def forward(self, mhc_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor,
         pre = torch.sigmoid(pre_w * pre_scale + pre_b) + self.hc_eps
         post = 2 * torch.sigmoid(post_w * post_scale + post_b)
         comb_logits = comb_w.view(*comb_w.shape[:-1], hc, hc) * comb_scale + comb_b.view(hc, hc)
-        comb = dsv4_mhc.fused_sinkhorn(comb_logits, self.hc_sinkhorn_iters, self.hc_eps)
+        comb = (
+            comb_logits
+            if comb_logits.numel() == 0
+            else dsv4_mhc.fused_sinkhorn(comb_logits, self.hc_sinkhorn_iters, self.hc_eps)
+        )
 
         collapsed = (pre.unsqueeze(-1) * mhc_states).sum(dim=2).to(mhc_states.dtype)
         return post, comb, collapsed
@@ -74,6 +78,8 @@ def update_states(
     ) -> torch.Tensor:
         """Broadcast the sublayer output over the streams via `post` and remix them via `comb`."""
         dtype = mhc_states.dtype
+        if mhc_states.numel() == 0:
+            return (mhc_states + sublayer_out.unsqueeze(2) + post.unsqueeze(-1) + comb.sum()).to(dtype)
         return dsv4_mhc.fused_post_bda(comb.to(dtype), mhc_states, post.to(dtype), sublayer_out)
 
     def init_weights(self, init_std: float) -> None:
diff --git a/src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py b/src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py
index 9f63f342b44..f7dd0e84ce1 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py
+++ b/src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py
@@ -29,6 +29,7 @@
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
+from prime_rl.utils.sequence import CPPartition
 
 
 class DeepseekV4DecoderLayer(GradientCheckpointingLayer):
@@ -204,6 +205,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> MoeModelOutputWithPast:
         """
         input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
@@ -246,9 +248,9 @@ def forward(
 
         # `seq_lens` describes the whole row and `inputs_embeds` carries this rank's shard of it.
         total_tokens = int(seq_lens.sum())
-        assert total_tokens == inputs_embeds.shape[1] * cp_world_size, (
-            f"seq_lens covers {total_tokens} tokens, but {cp_world_size} CP rank(s) holding "
-            f"{inputs_embeds.shape[1]} tokens each"
+        partition = CPPartition(total_tokens, cp_world_size)
+        assert inputs_embeds.shape[1] == partition.lengths[cp_rank], (
+            f"local tokens {inputs_embeds.shape[1]} disagree with CP partition {partition.lengths}"
         )
 
         # Every layer type attends over the same local window; the compressed variants add their
@@ -309,6 +311,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs,
     ) -> PrimeLmOutput:
         """
@@ -329,6 +332,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
         hidden_states = outputs.last_hidden_state
diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py
index 3b2dfa20601..b04755bbff8 100644
--- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py
+++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py
@@ -96,6 +96,7 @@ def forward(
         position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
         cu_seqlens: Optional[torch.LongTensor] = None,
         max_seqlen: Optional[int] = None,
+        cp_total_tokens: int | None = None,
         routed_experts: Optional[torch.LongTensor] = None,
     ) -> torch.Tensor:
         residual = hidden_states
@@ -106,6 +107,7 @@ def forward(
             position_embeddings=position_embeddings,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = residual + hidden_states
 
@@ -196,6 +198,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutputWithPast:
         """
         routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*):
@@ -204,6 +207,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced context parallelism.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -213,9 +218,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -231,6 +240,7 @@ def forward(
                 position_embeddings=position_embeddings,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=routed_experts_layer,
             )
 
@@ -269,6 +279,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         r"""
@@ -276,6 +287,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced context parallelism.
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
             Indices of input tokens in the KV cache. Accepted only for HuggingFace API
             compatibility — prime-rl asserts `use_cache is None` since training does not
@@ -320,6 +333,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
         hidden_states = outputs.last_hidden_state
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
index 08717e20654..39c22b5537e 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
@@ -21,6 +21,8 @@
 from prime_rl.trainer.models.layers.moe import MoE, MoEArgs
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig
+from prime_rl.utils.cp import gather_for_cp
+from prime_rl.utils.sequence import CPPartition
 
 
 def _sparse_mla_attention_args(config: GlmMoeDsaConfig, layer_idx: int) -> SparseMlaAttentionArgs:
@@ -182,7 +184,10 @@ def _gather_position_ids_for_cp(
         position_ids: torch.LongTensor,
         cp_group: dist.ProcessGroup,
         cp_world_size: int,
+        total_tokens: int | None = None,
     ) -> torch.LongTensor:
+        if total_tokens is not None:
+            return gather_for_cp(position_ids, cp_group, total_tokens)
         gathered_position_ids = [torch.empty_like(position_ids) for _ in range(cp_world_size)]
         dist.all_gather(gathered_position_ids, position_ids.contiguous(), group=cp_group)
         return torch.cat(gathered_position_ids, dim=1)
@@ -197,12 +202,15 @@ def forward(
         *,
         seq_lens: Optional[torch.LongTensor] = None,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutputWithPast:
         """
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced context parallelism.
         routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*):
             Routed experts for each token in the sequence. Only used for router replay.
         seq_lens (`torch.LongTensor` of shape `(num_documents,)`, *optional*):
-            Per-document lengths of the packed row (PrimeRL packed-batch contract). Unused.
+            Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries. Unused.
         """
@@ -212,9 +220,13 @@ def forward(
         if inputs_embeds is None:
             inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
 
+        if self.cp_context.cp_unpadded and cp_total_tokens is None:
+            cp_total_tokens = int(seq_lens.sum().item())
         cp_rank, cp_world_size = self.cp_context.cp_rank, self.cp_context.cp_world_size
         if self.cp_context.cp_enabled:
-            position_ids_full = self._gather_position_ids_for_cp(position_ids, self.cp_context.cp_group, cp_world_size)
+            position_ids_full = self._gather_position_ids_for_cp(
+                position_ids, self.cp_context.cp_group, cp_world_size, cp_total_tokens
+            )
         else:
             position_ids_full = position_ids
 
@@ -233,9 +245,9 @@ def forward(
         position_embeddings = self.rotary_emb(hidden_states, position_ids_full)
 
         if cp_world_size > 1:
-            s_local = S_full // cp_world_size
-            ks = ks_full[cp_rank * s_local : (cp_rank + 1) * s_local].contiguous()
-            ke = ke_full[cp_rank * s_local : (cp_rank + 1) * s_local].contiguous()
+            partition = CPPartition(S_full, cp_world_size)
+            ks = partition.shard(ks_full, cp_rank, dim=0).contiguous()
+            ke = partition.shard(ke_full, cp_rank, dim=0).contiguous()
         else:
             ks, ke = ks_full, ke_full
 
@@ -293,6 +305,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         r"""
@@ -308,6 +321,8 @@ def forward(
             Labels used by PrimeRL's wrapped LM head to optionally compute per-token logprobs/entropy.
         temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
             Per-token temperatures for logprobs/entropy computation when `labels` are provided.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced context parallelism.
         routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*):
             Routed experts for each token in the sequence. Only used for router replay.
         """
@@ -325,6 +340,9 @@ def forward(
             position_ids=position_ids,
             inputs_embeds=inputs_embeds,
             routed_experts=routed_experts,
+            seq_lens=seq_lens,
+            seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
         hidden_states = outputs.last_hidden_state
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 02673f6d3f0..36b92bb1501 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -8,6 +8,7 @@
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
 from prime_rl.utils.cp import CPContext, gather_for_cp
+from prime_rl.utils.sequence import CPPartition
 
 try:
     from prime_rl.trainer.models.kernels.sparse_mla_fwd import sparse_mla
@@ -78,14 +79,15 @@ def compute_sparse_indices(
         w = self.weights_proj(hidden_states_local[0])
 
         if cp_world_size > 1:
-            k_idx = gather_for_cp(k_idx_local.unsqueeze(0), cp_group).squeeze(0)
+            k_idx = gather_for_cp(k_idx_local.unsqueeze(0), cp_group, position_embeddings_full[0].shape[1]).squeeze(0)
         else:
             k_idx = k_idx_local
 
         cos_full, sin_full = position_embeddings_full
         if cp_world_size > 1:
-            cos_local = cos_full[:, cp_rank * s_local : (cp_rank + 1) * s_local, :]
-            sin_local = sin_full[:, cp_rank * s_local : (cp_rank + 1) * s_local, :]
+            partition = CPPartition(cos_full.shape[1], cp_world_size)
+            cos_local = partition.shard(cos_full, cp_rank)
+            sin_local = partition.shard(sin_full, cp_rank)
         else:
             cos_local, sin_local = cos_full, sin_full
 
@@ -105,6 +107,8 @@ def compute_sparse_indices(
         q_idx = torch.cat([q_pe, q_nope], dim=-1)
         k_idx = torch.cat([k_pe, k_nope], dim=-1)
 
+        if s_local == 0:
+            return torch.empty((1, 0, 1, index_topk), device=q_idx.device, dtype=torch.int32)
         indices = fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk, self.weight_scale)
         # indices shape: [S_local, topk] in K's coordinate space (sentinel = s_full)
         # KV passed to sparse MLA has length s_full + 1 (sentinel zeros at index s_full).
@@ -169,8 +173,9 @@ def mla_up_proj(
         cos_full, sin_full = position_embeddings_full
         if self.cp_context.cp_enabled:
             cp_rank = self.cp_context.cp_rank
-            cos_local = cos_full[:, cp_rank * s_local : (cp_rank + 1) * s_local, :]
-            sin_local = sin_full[:, cp_rank * s_local : (cp_rank + 1) * s_local, :]
+            partition = CPPartition(s_full, self.cp_context.cp_world_size)
+            cos_local = partition.shard(cos_full, cp_rank)
+            sin_local = partition.shard(sin_full, cp_rank)
         else:
             cos_local, sin_local = cos_full, sin_full
 
@@ -206,8 +211,10 @@ def forward(
         q_latent, k_compressed_normed, k_rope = self.mla_latents(hidden_states)
 
         if self.cp_context.cp_enabled:
-            k_compressed_normed = gather_for_cp(k_compressed_normed, self.cp_context.cp_group)
-            k_rope = gather_for_cp(k_rope, self.cp_context.cp_group)
+            k_compressed_normed = gather_for_cp(
+                k_compressed_normed, self.cp_context.cp_group, position_embeddings[0].shape[1]
+            )
+            k_rope = gather_for_cp(k_rope, self.cp_context.cp_group, position_embeddings[0].shape[1])
 
         indices = cached_indices
         if not self.skip_topk:
@@ -230,9 +237,12 @@ def forward(
             position_embeddings_full=position_embeddings,
         )
 
-        out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.scaling)
+        if sparse_q.shape[1] == 0:
+            out = sparse_q[..., : self.kv_lora_rank] + sparse_kv.sum().to(sparse_q.dtype)
+        else:
+            out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.scaling)
         out = torch.einsum("bshk,hdk->bshd", out, w_v)
         batch_size, total_tokens = out.shape[:2]
-        out = out.reshape(batch_size, total_tokens, -1)
+        out = out.reshape(batch_size, total_tokens, self.num_heads * self.v_head_dim)
         cached_indices = indices if self.use_index_cache else None
         return self.o_proj(out), cached_indices
diff --git a/src/prime_rl/trainer/models/gpt_oss/attention.py b/src/prime_rl/trainer/models/gpt_oss/attention.py
index 73cf2a6db2a..252e9c8a538 100644
--- a/src/prime_rl/trainer/models/gpt_oss/attention.py
+++ b/src/prime_rl/trainer/models/gpt_oss/attention.py
@@ -12,6 +12,7 @@
 from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig
 from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func
 from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb
+from prime_rl.utils.cp import CPContext
 
 
 class GptOssAttention(nn.Module):
@@ -50,6 +51,7 @@ def __init__(self, config: GptOssConfig, layer_idx: int) -> None:
         self.sinks = nn.Parameter(torch.empty(config.num_attention_heads))
         nn.init.normal_(self.sinks, mean=0.0, std=config.initializer_range)
         self.flash_attn = flash_attn_4_varlen_func
+        self.cp_context = CPContext()
 
     def compute_attention(
         self,
@@ -58,8 +60,28 @@ def compute_attention(
         value: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        cp_total_tokens: int | None = None,
     ) -> torch.Tensor:
         window_size = (self.sliding_window - 1, 0) if self.sliding_window is not None else (None, None)
+        if self.cp_context.cp_enabled and cp_total_tokens is not None:
+            from prime_rl.trainer.models.layers.cp_attn import context_parallel_attention
+
+            return context_parallel_attention(
+                self.flash_attn,
+                query,
+                key,
+                value,
+                cu_seqlens,
+                max_seqlen,
+                cp_total_tokens,
+                self.cp_context,
+                4,
+                window_size=window_size,
+                softmax_scale=self.scaling,
+                learnable_sink=self.sinks,
+            )
+        if query.shape[0] == 0:
+            return query + (key.sum() + value.sum() + self.sinks.sum()).to(query.dtype)
         output, _ = self.flash_attn(
             query,
             key,
@@ -81,6 +103,7 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        cp_total_tokens: int | None = None,
     ) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
         if self.qkv_proj is None:
@@ -101,8 +124,17 @@ def forward(
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)
 
-        output = self.compute_attention(query[0], key[0], value[0], cu_seqlens, max_seqlen)
-        return self.o_proj(output.contiguous().view(batch_size, sequence_length, -1))
+        output = self.compute_attention(
+            query[0],
+            key[0],
+            value[0],
+            cu_seqlens,
+            max_seqlen,
+            **({"cp_total_tokens": cp_total_tokens} if cp_total_tokens is not None else {}),
+        )
+        return self.o_proj(
+            output.contiguous().view(batch_size, sequence_length, self.num_attention_heads * self.head_dim)
+        )
 
 
 def substitute_gpt_oss_ring_attention(
diff --git a/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py
index 36a60c780e3..f31cc98c01f 100644
--- a/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py
+++ b/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py
@@ -69,11 +69,14 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        cp_total_tokens: int | None = None,
         routed_experts: torch.Tensor | None = None,
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(hidden_states, position_embeddings, cu_seqlens, max_seqlen)
+        hidden_states = self.self_attn(
+            hidden_states, position_embeddings, cu_seqlens, max_seqlen, cp_total_tokens=cp_total_tokens
+        )
         hidden_states = residual + hidden_states
 
         residual = hidden_states
@@ -135,15 +138,20 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutput:
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -154,6 +162,7 @@ def forward(
                 position_embeddings,
                 cu_seqlens,
                 max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=layer_routed_experts,
             )
         return BaseModelOutput(last_hidden_state=self.norm(hidden_states))
@@ -185,6 +194,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> PrimeLmOutput:
         outputs = self.model(
             input_ids=input_ids,
@@ -193,6 +203,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = outputs.last_hidden_state
         if isinstance(logits_to_keep, int):
diff --git a/src/prime_rl/trainer/models/laguna/modeling_laguna.py b/src/prime_rl/trainer/models/laguna/modeling_laguna.py
index cd919d3b116..45118869e73 100644
--- a/src/prime_rl/trainer/models/laguna/modeling_laguna.py
+++ b/src/prime_rl/trainer/models/laguna/modeling_laguna.py
@@ -126,9 +126,9 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
 
         query_states, key_states, value_states = self.project_qkv(hidden_states)
 
@@ -136,9 +136,9 @@ def forward(
             query_states = self.q_norm(query_states)
             key_states = self.k_norm(key_states)
 
-        query_states = query_states.view(hidden_shape)
-        key_states = key_states.view(hidden_shape)
-        value_states = value_states.view(hidden_shape)
+        query_states = query_states.view(*input_shape, self.qkv_sizes[0] // self.head_dim, self.head_dim)
+        key_states = key_states.view(*input_shape, self.qkv_sizes[1] // self.head_dim, self.head_dim)
+        value_states = value_states.view(*input_shape, self.qkv_sizes[2] // self.head_dim, self.head_dim)
 
         if self.use_qk_norm and self.qk_norm_type == "per_head":
             query_states = self.q_norm(query_states)
@@ -156,7 +156,14 @@ def forward(
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
 
-        attn_output = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen)
+        attn_output = self._compute_attention(
+            query_states[0],
+            key_states[0],
+            value_states[0],
+            cu_seqlens,
+            max_seqlen,
+            **({"cp_total_tokens": cp_total_tokens} if cp_total_tokens is not None else {}),
+        )
         attn_output = attn_output.contiguous().view(*input_shape, self.num_heads, self.head_dim)
         if self.gating:
             gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
@@ -164,7 +171,7 @@ def forward(
             # one-per-channel and line up with the [..., num_heads, head_dim] view.
             gate = gate.unsqueeze(-1) if self.gate_per_head else gate.view(*attn_output.shape)
             attn_output = attn_output * gate
-        attn_output = attn_output.view(*input_shape, -1)
+        attn_output = attn_output.view(*input_shape, self.num_heads * self.head_dim)
         return self.o_proj(attn_output), None
 
 
@@ -232,6 +239,7 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
         routed_experts: Optional[torch.LongTensor] = None,
     ) -> torch.Tensor:
         residual = hidden_states
@@ -242,6 +250,7 @@ def forward(
             position_embeddings=position_embeddings,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = residual + hidden_states
 
@@ -306,6 +315,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> MoeModelOutputWithPast:
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -317,9 +327,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
         causal_mask_mapping = dict.fromkeys(set(self.config.layer_types), None)
 
         hidden_states = inputs_embeds
@@ -336,6 +350,7 @@ def forward(
                 position_embeddings=position_embeddings[self.config.layer_types[layer_idx]],
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=routed_experts_layer,
             )
 
@@ -382,6 +397,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         assert use_cache is None, "use_cache is not supported for custom Laguna"
@@ -395,6 +411,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = outputs.last_hidden_state
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py
index 97e59000f9b..8d0d9ce3043 100644
--- a/src/prime_rl/trainer/models/layers/attn.py
+++ b/src/prime_rl/trainer/models/layers/attn.py
@@ -6,6 +6,7 @@
 from torch import nn
 
 from prime_rl.trainer.models.fusions import fuse_qkv_projections
+from prime_rl.utils.cp import CPContext
 
 from .norms import RMSNorm, RMSNormConfig
 from .rotary_emb import apply_rotary_pos_emb
@@ -91,6 +92,7 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):
 
         self._flash_attn_version = flash_attn_version
         self.func = self._funcs[flash_attn_version]
+        self.cp_context = CPContext()
 
     def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Query, key and value projections, from one packed GEMM when qkv is fused."""
@@ -98,8 +100,29 @@ def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.
             return self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states)
         return self.qkv_proj(hidden_states).split(self.qkv_sizes, dim=-1)
 
-    def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen):
+    def _compute_attention(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen, cp_total_tokens=None
+    ):
         """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim]."""
+        if self.cp_context.cp_enabled and cp_total_tokens is not None:
+            from .cp_attn import context_parallel_attention
+
+            sliding_window = getattr(self, "sliding_window", None)
+            window = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
+            return context_parallel_attention(
+                self.func,
+                q,
+                k,
+                v,
+                cu_seqlens,
+                max_seqlen,
+                cp_total_tokens,
+                self.cp_context,
+                self._flash_attn_version,
+                window_size=window,
+            )
+        if q.shape[0] == 0:
+            return q + (k.sum() + v.sum()).to(q.dtype)
         kwargs: dict = {"causal": True}
         sliding_window = getattr(self, "sliding_window", None)
         if sliding_window is not None:
@@ -111,7 +134,8 @@ def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
             kwargs["cu_seqlens_k"] = cu_seqlens
             out, _ = self.func(q, k, v, **kwargs)
         else:
-            out = self.func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, **kwargs)
+            cu_k = cu_seqlens.clone()
+            out = self.func(q, k, v, cu_seqlens, cu_k, max_seqlen, max_seqlen, **kwargs)
         return out
 
     def forward(
@@ -120,9 +144,9 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
 
         query_states, key_states, value_states = self.project_qkv(hidden_states)
 
@@ -130,9 +154,9 @@ def forward(
             query_states = self.q_norm(query_states)
             key_states = self.k_norm(key_states)
 
-        query_states = query_states.view(hidden_shape)
-        key_states = key_states.view(hidden_shape)
-        value_states = value_states.view(hidden_shape)
+        query_states = query_states.view(*input_shape, self.qkv_sizes[0] // self.head_dim, self.head_dim)
+        key_states = key_states.view(*input_shape, self.qkv_sizes[1] // self.head_dim, self.head_dim)
+        value_states = value_states.view(*input_shape, self.qkv_sizes[2] // self.head_dim, self.head_dim)
 
         if self.use_qk_norm and self.qk_norm_type == "per_head":
             query_states = self.q_norm(query_states)
@@ -151,8 +175,18 @@ def forward(
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
 
-        out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen)
-        attn_output = out.contiguous().view(1, out.shape[0], -1)
+        if cp_total_tokens is None:
+            out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen)
+        else:
+            out = self._compute_attention(
+                query_states[0],
+                key_states[0],
+                value_states[0],
+                cu_seqlens,
+                max_seqlen,
+                cp_total_tokens=cp_total_tokens,
+            )
+        attn_output = out.contiguous().view(1, out.shape[0], self.qkv_sizes[0])
         attn_output = self.o_proj(attn_output)
         return attn_output, None
 
diff --git a/src/prime_rl/trainer/models/layers/cp_attn.py b/src/prime_rl/trainer/models/layers/cp_attn.py
new file mode 100644
index 00000000000..c994f33010c
--- /dev/null
+++ b/src/prime_rl/trainer/models/layers/cp_attn.py
@@ -0,0 +1,77 @@
+"""Context-parallel attention over an explicit, balanced token partition."""
+
+import torch
+
+from prime_rl.trainer.distributed.collectives import all_gather_cp
+from prime_rl.trainer.models.layers.ulysses_attn import ulysses_flash_attn_varlen_func
+from prime_rl.utils.cp import CPContext
+from prime_rl.utils.sequence import CPPartition
+
+
+def context_parallel_attention(
+    flash_fn,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    max_seqlen: int,
+    total_tokens: int,
+    context: CPContext,
+    flash_attn_version: int,
+    window_size: tuple[int, int] = (-1, -1),
+    softmax_scale: float | None = None,
+    learnable_sink: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if context.cp_style == "ulysses":
+        return ulysses_flash_attn_varlen_func(
+            flash_fn,
+            q,
+            k,
+            v,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            True,
+            context.cp_group,
+            context.cp_world_size,
+            flash_attn_version=flash_attn_version,
+            window_size=window_size,
+            softmax_scale=softmax_scale,
+            learnable_sink=learnable_sink,
+            total_tokens=total_tokens,
+        )
+
+    partition = CPPartition(total_tokens, context.cp_world_size)
+    start, end = partition.offsets[context.cp_rank : context.cp_rank + 2]
+    cu_q = (cu_seqlens - start).clamp(min=0, max=end - start)
+    cu_k = cu_seqlens.clamp(max=end)
+    kv = all_gather_cp(torch.cat((k, v), dim=-1), 0, total_tokens, context.cp_group)[:end]
+    key, value = kv.split((k.shape[-1], v.shape[-1]), dim=-1)
+    if q.shape[0] == 0:
+        output = q + kv.sum().to(q.dtype)
+        if learnable_sink is not None:
+            output = output + learnable_sink.sum().to(output.dtype)
+        return output
+
+    kwargs = {"causal": True}
+    if window_size != (-1, -1):
+        kwargs["window_size"] = window_size
+    if softmax_scale is not None:
+        kwargs["softmax_scale"] = softmax_scale
+    if learnable_sink is not None:
+        kwargs["learnable_sink"] = learnable_sink
+    if flash_attn_version == 4:
+        output, _ = flash_fn(
+            q,
+            key,
+            value,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            max_seqlen_q=min(max_seqlen, end - start),
+            max_seqlen_k=max_seqlen,
+            **kwargs,
+        )
+    else:
+        output = flash_fn(q, key, value, cu_q, cu_k, min(max_seqlen, end - start), max_seqlen, **kwargs)
+    return output
diff --git a/src/prime_rl/trainer/models/layers/fp8_linear.py b/src/prime_rl/trainer/models/layers/fp8_linear.py
index b7ea80fb705..48e746d040e 100644
--- a/src/prime_rl/trainer/models/layers/fp8_linear.py
+++ b/src/prime_rl/trainer/models/layers/fp8_linear.py
@@ -17,6 +17,8 @@
 
 @torch.library.custom_op("prime_rl::fp8_blockwise_mm", mutates_args=())
 def _fp8_blockwise_mm(x: torch.Tensor, weight: torch.Tensor, block_size: int) -> torch.Tensor:
+    if x.numel() == 0:
+        return x.new_empty((*x.shape[:-1], weight.shape[0]), dtype=torch.bfloat16)
     import deep_gemm
 
     x_2d = x.reshape(-1, x.shape[-1]).contiguous()
@@ -43,6 +45,8 @@ def _fp8_blockwise_mm_backward(
     needs_grad_x: bool,
     needs_grad_weight: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    if x.numel() == 0:
+        return torch.empty_like(x), torch.zeros_like(weight)
     import deep_gemm
 
     x_2d = x.reshape(-1, x.shape[-1]).contiguous()
diff --git a/src/prime_rl/trainer/models/layers/lm_head.py b/src/prime_rl/trainer/models/layers/lm_head.py
index 46820fd674d..83289aa89db 100644
--- a/src/prime_rl/trainer/models/layers/lm_head.py
+++ b/src/prime_rl/trainer/models/layers/lm_head.py
@@ -7,9 +7,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
+from transformers.modeling_outputs import BaseModelOutputWithPast
 
 from prime_rl.utils.logger import get_logger
-from prime_rl.utils.vlm import get_final_logit_softcapping
+from prime_rl.utils.vlm import get_final_logit_softcapping, get_language_model
 
 # Same as torch's cross entropy loss
 IGNORE_INDEX = -100
@@ -385,6 +386,11 @@ def inject_prime_lm_head(
 
     logger = get_logger()
 
+    from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
+
+    if not isinstance(model, PreTrainedModelPrimeRL):
+        _enable_empty_backbone(model)
+
     # Check for Gemma-style softcapping - dispatch to specialized implementation.
     final_logit_softcapping = get_final_logit_softcapping(model.config)
     if final_logit_softcapping:
@@ -409,6 +415,42 @@ def inject_prime_lm_head(
     _patch_model_forward(model)
 
 
+def _enable_empty_backbone(model: nn.Module) -> None:
+    """Execute empty HF rows through the same embedding, decoder, and norm FSDP units."""
+    backbone = model.model
+    language_model = get_language_model(model)
+    language_path = next(name for name, module in backbone.named_modules() if module is language_model)
+    for layer in language_model.layers:
+        original_layer_forward = type(layer).forward
+
+        def layer_forward(self, hidden_states, *args, _empty_input=False, _forward=original_layer_forward, **kwargs):
+            if not _empty_input:
+                return _forward(self, hidden_states, *args, **kwargs)
+            for parameter in self.parameters():
+                if parameter.requires_grad:
+                    hidden_states = hidden_states + parameter.reshape(-1)[:0].sum().to(hidden_states.dtype)
+            return hidden_states
+
+        layer.forward = types.MethodType(layer_forward, layer)
+
+    original_forward = type(backbone).forward
+
+    def backbone_forward(self, input_ids=None, inputs_embeds=None, **kwargs):
+        inputs = input_ids if inputs_embeds is None else inputs_embeds
+        if inputs.shape[1] != 0:
+            return original_forward(self, input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs)
+        language_model = self.get_submodule(language_path)
+        if inputs_embeds is None:
+            embedding = getattr(language_model, "embed_tokens", None) or language_model.embeddings
+            inputs_embeds = embedding(input_ids)
+        for layer in language_model.layers:
+            inputs_embeds = layer(inputs_embeds, _empty_input=True)
+        norm = getattr(language_model, "norm", None) or language_model.norm_f
+        return BaseModelOutputWithPast(last_hidden_state=norm(inputs_embeds))
+
+    backbone.forward = types.MethodType(backbone_forward, backbone)
+
+
 def _patch_model_forward(model: nn.Module) -> None:
     # Patch the forward method to use the new lm_head with labels and temperature
     def new_forward(
diff --git a/src/prime_rl/trainer/models/layers/norms.py b/src/prime_rl/trainer/models/layers/norms.py
index 8658c92fa37..2e6d78f51e9 100644
--- a/src/prime_rl/trainer/models/layers/norms.py
+++ b/src/prime_rl/trainer/models/layers/norms.py
@@ -35,7 +35,7 @@ def __init__(self, config: RMSNormConfig) -> None:
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         quack_fn = get_quack_rmsnorm() if hidden_states.is_cuda else None
-        if quack_fn is not None:
+        if quack_fn is not None and hidden_states.numel() > 0:
             return quack_fn(hidden_states, self.weight, eps=self.variance_epsilon)
         input_dtype = hidden_states.dtype
         hidden_states = hidden_states.to(torch.float32)
diff --git a/src/prime_rl/trainer/models/layers/ulysses_attn.py b/src/prime_rl/trainer/models/layers/ulysses_attn.py
index 4706049e83a..efbcd5b25f5 100644
--- a/src/prime_rl/trainer/models/layers/ulysses_attn.py
+++ b/src/prime_rl/trainer/models/layers/ulysses_attn.py
@@ -29,7 +29,7 @@
 - cp_size must divide num_attention_heads.
 - cp_size must divide num_key_value_heads, OR num_key_value_heads must divide
   cp_size (the KV-replication path).
-- Sequence length must be divisible by cp_size.
+- Uneven sequence shards require the full token count at each redistribution.
 """
 
 from __future__ import annotations
@@ -37,7 +37,7 @@
 import torch
 import torch.distributed as dist
 
-from prime_rl.trainer.distributed.collectives import all_to_all_single_equal
+from prime_rl.trainer.distributed.collectives import all_to_all_cp, all_to_all_single_equal
 
 # Populated by `update_ulysses_params` before each forward pass. Mirrors
 # ring_flash_attn's DATA_PARAMS pattern so the patched attention path can
@@ -68,7 +68,9 @@ def _replicate_kv_heads(t: torch.Tensor, cp_size: int) -> torch.Tensor:
     return t.repeat_interleave(cp_size // h, dim=1)
 
 
-def _all_to_all_seq_to_head(t: torch.Tensor, cp_size: int, cp_group: dist.ProcessGroup) -> torch.Tensor:
+def _all_to_all_seq_to_head(
+    t: torch.Tensor, cp_size: int, cp_group: dist.ProcessGroup, total_tokens: int | None = None
+) -> torch.Tensor:
     """Redistribute [S_local, H, D] -> [S_global, H_local, D].
 
     Splits the head dim into cp_size groups and exchanges them so each rank
@@ -83,15 +85,26 @@ def _all_to_all_seq_to_head(t: torch.Tensor, cp_size: int, cp_group: dist.Proces
 
     # [S_local, cp_size, H_local, D] -> [cp_size, S_local, H_local, D]
     t = t.reshape(s_local, cp_size, h_local, d).transpose(0, 1).contiguous()
+    if total_tokens is not None:
+        return all_to_all_cp(t.reshape(cp_size * s_local, h_local, d), total_tokens, cp_group, True)
     out = all_to_all_single_equal(t, cp_group)
     # out[i] is the chunk that source-rank i had at position my_rank, i.e.
     # source-rank i's local sequence shard for *my* head slice.
     return out.reshape(cp_size * s_local, h_local, d)
 
 
-def _all_to_all_head_to_seq(t: torch.Tensor, cp_size: int, cp_group: dist.ProcessGroup) -> torch.Tensor:
+def _all_to_all_head_to_seq(
+    t: torch.Tensor,
+    cp_size: int,
+    cp_group: dist.ProcessGroup,
+    total_tokens: int | None = None,
+) -> torch.Tensor:
     """Inverse of `_all_to_all_seq_to_head`: [S_global, H_local, D] -> [S_local, H, D]."""
     s_global, h_local, d = t.shape
+    if total_tokens is not None:
+        out = all_to_all_cp(t, total_tokens, cp_group, False)
+        s_local = out.shape[0] // cp_size
+        return out.reshape(cp_size, s_local, h_local, d).transpose(0, 1).reshape(s_local, h_local * cp_size, d)
     assert s_global % cp_size == 0
     s_local = s_global // cp_size
     h = h_local * cp_size
@@ -102,6 +115,26 @@ def _all_to_all_head_to_seq(t: torch.Tensor, cp_size: int, cp_group: dist.Proces
     return out.transpose(0, 1).contiguous().reshape(s_local, h, d)
 
 
+def sequence_to_head_parallel(tensor, process_group, world_size, total_tokens=None):
+    """Redistribute packed [B, local tokens, features] to full tokens and local features."""
+    return torch.stack(
+        [
+            _all_to_all_seq_to_head(row.unsqueeze(-1), world_size, process_group, total_tokens).squeeze(-1)
+            for row in tensor
+        ]
+    )
+
+
+def head_to_sequence_parallel(tensor, process_group, world_size, total_tokens=None):
+    """Restore token shards from full tokens and local features."""
+    return torch.stack(
+        [
+            _all_to_all_head_to_seq(row.unsqueeze(-1), world_size, process_group, total_tokens).squeeze(-1)
+            for row in tensor
+        ]
+    )
+
+
 def ulysses_flash_attn_varlen_func(
     flash_fn,
     q: torch.Tensor,
@@ -120,6 +153,7 @@ def ulysses_flash_attn_varlen_func(
     dropout_p: float = 0.0,
     deterministic: bool | None = None,
     learnable_sink: torch.Tensor | None = None,
+    total_tokens: int | None = None,
 ) -> torch.Tensor:
     """Run varlen flash attention under Ulysses CP.
 
@@ -130,12 +164,18 @@ def ulysses_flash_attn_varlen_func(
     cp_size before the all-to-all (see `_replicate_kv_heads`), so each rank runs
     grouped attention on its query-head slice with the single matching KV head.
     """
-    q = _all_to_all_seq_to_head(q, cp_size, cp_group)
+    q = _all_to_all_seq_to_head(q, cp_size, cp_group, total_tokens)
     if k.shape[1] < cp_size:
         k = _replicate_kv_heads(k, cp_size)
         v = _replicate_kv_heads(v, cp_size)
-    k = _all_to_all_seq_to_head(k, cp_size, cp_group)
-    v = _all_to_all_seq_to_head(v, cp_size, cp_group)
+    k = _all_to_all_seq_to_head(k, cp_size, cp_group, total_tokens)
+    v = _all_to_all_seq_to_head(v, cp_size, cp_group, total_tokens)
+
+    if total_tokens == 0:
+        out = q + (k.sum() + v.sum()).to(q.dtype)
+        if learnable_sink is not None:
+            out = out + learnable_sink.sum().to(out.dtype)
+        return _all_to_all_head_to_seq(out, cp_size, cp_group, total_tokens)
 
     kwargs: dict = {"causal": causal}
     if window_size != (-1, -1):
@@ -159,8 +199,10 @@ def ulysses_flash_attn_varlen_func(
         kwargs["max_seqlen_k"] = max_seqlen_k
         out, _ = flash_fn(q, k, v, **kwargs)
     else:
+        # FA2/FA3's autograd Functions require distinct inputs under Dynamo.
+        cu_seqlens_k = cu_seqlens_k.clone()
         out = flash_fn(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs)
-    return _all_to_all_head_to_seq(out, cp_size, cp_group)
+    return _all_to_all_head_to_seq(out, cp_size, cp_group, total_tokens)
 
 
 def substitute_ulysses_attn(
@@ -188,11 +230,11 @@ def substitute_ulysses_attn(
 
         flash_attn_version = 2
 
-    def _ulysses_compute_attention(self, q, k, v, cu_seqlens, max_seqlen):
+    def _ulysses_compute_attention(self, q, k, v, cu_seqlens, max_seqlen, cp_total_tokens=None):
         # cu_seqlens / max_seqlen passed in are for the *local* sharded sequence;
         # ulysses needs the *full* ones (each rank holds the full seq after a2a).
-        cu_seqlens_full = ULYSSES_PARAMS["cu_seqlens"]
-        max_seqlen_full = ULYSSES_PARAMS["max_seqlen"]
+        cu_seqlens_full = cu_seqlens if cp_total_tokens is not None else ULYSSES_PARAMS["cu_seqlens"]
+        max_seqlen_full = max_seqlen if cp_total_tokens is not None else ULYSSES_PARAMS["max_seqlen"]
 
         window_size = (-1, -1)
         sliding_window = getattr(self, "sliding_window", None)
@@ -213,6 +255,7 @@ def _ulysses_compute_attention(self, q, k, v, cu_seqlens, max_seqlen):
             cp_size=cp_size,
             flash_attn_version=flash_attn_version,
             window_size=window_size,
+            total_tokens=cp_total_tokens,
         )
 
     from prime_rl.trainer.models.layers.attn import FlashAttention
diff --git a/src/prime_rl/trainer/models/llama/modeling_llama.py b/src/prime_rl/trainer/models/llama/modeling_llama.py
index e17d7464d25..3e1de07250e 100644
--- a/src/prime_rl/trainer/models/llama/modeling_llama.py
+++ b/src/prime_rl/trainer/models/llama/modeling_llama.py
@@ -78,6 +78,7 @@ def forward(
         position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
         cu_seqlens: Optional[torch.LongTensor] = None,
         max_seqlen: Optional[int] = None,
+        cp_total_tokens: int | None = None,
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
@@ -87,6 +88,7 @@ def forward(
             position_embeddings=position_embeddings,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = residual + hidden_states
 
@@ -185,12 +187,15 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutputWithPast:
         r"""
         seq_lens (`torch.LongTensor` of shape `(num_documents,)`):
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced Ulysses partitioning.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -200,9 +205,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -213,6 +222,7 @@ def forward(
                 position_embeddings=position_embeddings,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
             )
 
         hidden_states = self.norm(hidden_states)
@@ -253,6 +263,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         r"""
@@ -260,6 +271,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced Ulysses partitioning.
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
             Indices of input tokens in the KV cache. Accepted only for HuggingFace API
             compatibility — prime-rl asserts `use_cache is None` since training does not
@@ -301,6 +314,7 @@ def forward(
             inputs_embeds=inputs_embeds,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
         hidden_states = outputs.last_hidden_state
diff --git a/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py b/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py
index 8465019b331..43d41ea088d 100644
--- a/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py
+++ b/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py
@@ -67,6 +67,7 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
         routed_experts: Optional[torch.LongTensor] = None,
     ) -> torch.FloatTensor:
         residual = hidden_states
@@ -76,6 +77,7 @@ def forward(
             position_embeddings=position_embeddings,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = residual + hidden_states
 
@@ -152,6 +154,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutputWithPast:
         """
         routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*):
@@ -160,6 +163,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced context parallelism.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -169,9 +174,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -183,6 +192,7 @@ def forward(
                 position_embeddings=position_embeddings,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=routed_experts_layer,
             )
 
@@ -222,6 +232,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         r"""
@@ -229,6 +240,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced context parallelism.
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
             Indices of input tokens in the KV cache. Accepted only for HuggingFace API
             compatibility — prime-rl asserts `use_cache is None` since training does not
@@ -256,6 +269,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
         hidden_states = outputs.last_hidden_state
diff --git a/src/prime_rl/trainer/models/nemotron_h/mamba.py b/src/prime_rl/trainer/models/nemotron_h/mamba.py
index 572ac9e257b..062233018fe 100644
--- a/src/prime_rl/trainer/models/nemotron_h/mamba.py
+++ b/src/prime_rl/trainer/models/nemotron_h/mamba.py
@@ -9,50 +9,13 @@
 import math
 
 import torch
-import torch.distributed as dist
 from torch import nn
 
-from prime_rl.trainer.distributed.collectives import all_to_all_single_equal
+from prime_rl.trainer.models.layers.ulysses_attn import head_to_sequence_parallel, sequence_to_head_parallel
 from prime_rl.trainer.models.nemotron_h.configuration_nemotron_h import NemotronHConfig
 from prime_rl.utils.cp import CPContext
 
 
-def sequence_to_head_parallel(
-    tensor: torch.Tensor,
-    process_group: dist.ProcessGroup,
-    world_size: int,
-) -> torch.Tensor:
-    """Redistribute ``[B, S/world, D]`` into ``[B, S, D/world]``."""
-    batch_size, local_sequence_length, feature_size = tensor.shape
-    local_feature_size = feature_size // world_size
-    tensor = tensor.reshape(batch_size, local_sequence_length, world_size, local_feature_size)
-    tensor = tensor.permute(2, 0, 1, 3).contiguous()
-    tensor = all_to_all_single_equal(tensor, process_group)
-    return tensor.permute(1, 0, 2, 3).reshape(
-        batch_size,
-        world_size * local_sequence_length,
-        local_feature_size,
-    )
-
-
-def head_to_sequence_parallel(
-    tensor: torch.Tensor,
-    process_group: dist.ProcessGroup,
-    world_size: int,
-) -> torch.Tensor:
-    """Redistribute ``[B, S, D/world]`` into ``[B, S/world, D]``."""
-    batch_size, sequence_length, local_feature_size = tensor.shape
-    local_sequence_length = sequence_length // world_size
-    tensor = tensor.reshape(batch_size, world_size, local_sequence_length, local_feature_size)
-    tensor = tensor.permute(1, 0, 2, 3).contiguous()
-    tensor = all_to_all_single_equal(tensor, process_group)
-    return tensor.permute(1, 2, 0, 3).reshape(
-        batch_size,
-        local_sequence_length,
-        world_size * local_feature_size,
-    )
-
-
 class GatedRMSNorm(nn.Module):
     def __init__(self, hidden_size: int, group_size: int, eps: float) -> None:
         super().__init__()
@@ -127,7 +90,9 @@ def __init__(self, config: NemotronHConfig) -> None:
 
         self.cp_context = CPContext()
 
-    def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, cp_total_tokens: int | None = None
+    ) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
         projected_states = self.in_proj(hidden_states)
         gate, convolution_input, time_step = torch.split(
@@ -150,8 +115,8 @@ def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torc
             cp_group = self.cp_context.cp_group
             world_size = self.cp_context.cp_world_size
             rank = self.cp_context.cp_rank
-            gate = sequence_to_head_parallel(gate, cp_group, world_size)
-            time_step = sequence_to_head_parallel(time_step, cp_group, world_size)
+            gate = sequence_to_head_parallel(gate, cp_group, world_size, cp_total_tokens)
+            time_step = sequence_to_head_parallel(time_step, cp_group, world_size, cp_total_tokens)
 
             recurrent_input, state_input, state_output = torch.split(
                 convolution_input,
@@ -160,9 +125,9 @@ def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torc
             )
             convolution_input = torch.cat(
                 [
-                    sequence_to_head_parallel(recurrent_input, cp_group, world_size),
-                    sequence_to_head_parallel(state_input, cp_group, world_size),
-                    sequence_to_head_parallel(state_output, cp_group, world_size),
+                    sequence_to_head_parallel(recurrent_input, cp_group, world_size, cp_total_tokens),
+                    sequence_to_head_parallel(state_input, cp_group, world_size, cp_total_tokens),
+                    sequence_to_head_parallel(state_output, cp_group, world_size, cp_total_tokens),
                 ],
                 dim=-1,
             )
@@ -198,38 +163,51 @@ def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torc
                 )
             sequence_length = convolution_input.shape[1]
 
-        convolution_output, _ = self.causal_conv1d(
-            x=convolution_input,
-            weight=conv_weight.squeeze(1),
-            bias=conv_bias,
-            activation=self.activation,
-            cu_seqlens=cu_seqlens,
-        )
-        local_group_state_size = num_groups * self.state_size
-        hidden_states, state_input, state_output = torch.split(
-            convolution_output,
-            [intermediate_size, local_group_state_size, local_group_state_size],
-            dim=-1,
-        )
+        if sequence_length == 0:
+            anchor = (
+                convolution_input.sum()
+                + time_step.sum()
+                + conv_weight.sum()
+                + state_decay.sum()
+                + skip.sum()
+                + dt_bias.sum()
+            )
+            if conv_bias is not None:
+                anchor = anchor + conv_bias.sum()
+            hidden_states = gate + anchor.to(gate.dtype)
+        else:
+            convolution_output, _ = self.causal_conv1d(
+                x=convolution_input,
+                weight=conv_weight.squeeze(1),
+                bias=conv_bias,
+                activation=self.activation,
+                cu_seqlens=cu_seqlens,
+            )
+            local_group_state_size = num_groups * self.state_size
+            hidden_states, state_input, state_output = torch.split(
+                convolution_output,
+                [intermediate_size, local_group_state_size, local_group_state_size],
+                dim=-1,
+            )
 
-        scan_kwargs = {}
-        if self.time_step_limit is not None:
-            scan_kwargs["dt_limit"] = self.time_step_limit
-        hidden_states = self.scan(
-            hidden_states.reshape(batch_size, sequence_length, num_heads, self.head_dim),
-            time_step,
-            state_decay,
-            state_input.reshape(batch_size, sequence_length, num_groups, self.state_size),
-            state_output.reshape(batch_size, sequence_length, num_groups, self.state_size),
-            chunk_size=self.chunk_size,
-            D=skip,
-            z=None,
-            seq_idx=self.prepare_sequence_ids(cu_seqlens).to(torch.int32).unsqueeze(0),
-            return_final_states=False,
-            dt_bias=dt_bias,
-            dt_softplus=True,
-            **scan_kwargs,
-        ).reshape(batch_size, sequence_length, intermediate_size)
+            scan_kwargs = {}
+            if self.time_step_limit is not None:
+                scan_kwargs["dt_limit"] = self.time_step_limit
+            hidden_states = self.scan(
+                hidden_states.reshape(batch_size, sequence_length, num_heads, self.head_dim),
+                time_step,
+                state_decay,
+                state_input.reshape(batch_size, sequence_length, num_groups, self.state_size),
+                state_output.reshape(batch_size, sequence_length, num_groups, self.state_size),
+                chunk_size=self.chunk_size,
+                D=skip,
+                z=None,
+                seq_idx=self.prepare_sequence_ids(cu_seqlens).to(torch.int32).unsqueeze(0),
+                return_final_states=False,
+                dt_bias=dt_bias,
+                dt_softplus=True,
+                **scan_kwargs,
+            ).reshape(batch_size, sequence_length, intermediate_size)
         hidden_states = self.norm(hidden_states, gate, weight=norm_weight)
 
         if self.cp_context.cp_enabled:
@@ -237,6 +215,7 @@ def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torc
                 hidden_states,
                 self.cp_context.cp_group,
                 self.cp_context.cp_world_size,
+                cp_total_tokens,
             )
         return self.out_proj(hidden_states)
 
diff --git a/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py b/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py
index 2dc74b0ae1a..52dd1b37afc 100644
--- a/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py
+++ b/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py
@@ -110,18 +110,20 @@ def forward(
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        cp_total_tokens: int | None = None,
         routed_experts: torch.Tensor | None = None,
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.norm(hidden_states)
 
         if self.layer_type == "mamba":
-            hidden_states = self.mamba(hidden_states, cu_seqlens)
+            hidden_states = self.mamba(hidden_states, cu_seqlens, cp_total_tokens=cp_total_tokens)
         elif self.layer_type == "attention":
             hidden_states, _ = self.self_attn(
                 hidden_states=hidden_states,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
             )
         else:
             hidden_states = self.mlp(hidden_states, routed_experts=routed_experts)
@@ -182,14 +184,19 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutput:
         hidden_states = self.embed_tokens(input_ids)
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=hidden_states.device),
-            total_tokens=None if seq_lens_are_pre_shard else hidden_states.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else hidden_states.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
 
         for layer_idx, decoder_layer in enumerate(self.layers):
             layer_routed_experts = routed_experts[:, :, layer_idx] if routed_experts is not None else None
@@ -197,6 +204,7 @@ def forward(
                 hidden_states,
                 cu_seqlens,
                 max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=layer_routed_experts,
             )
         return BaseModelOutput(last_hidden_state=self.norm(hidden_states))
@@ -219,12 +227,14 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> PrimeLmOutput:
         outputs = self.model(
             input_ids=input_ids,
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
         return self.lm_head(
             outputs.last_hidden_state,
diff --git a/src/prime_rl/trainer/models/qwen3/modeling_qwen3.py b/src/prime_rl/trainer/models/qwen3/modeling_qwen3.py
index 8aa7e0101d0..66c62fd3755 100644
--- a/src/prime_rl/trainer/models/qwen3/modeling_qwen3.py
+++ b/src/prime_rl/trainer/models/qwen3/modeling_qwen3.py
@@ -66,6 +66,7 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
     ) -> torch.FloatTensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
@@ -74,6 +75,7 @@ def forward(
             position_embeddings=position_embeddings,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = residual + hidden_states
 
@@ -158,12 +160,15 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutputWithPast:
         r"""
         seq_lens (`torch.LongTensor` of shape `(num_documents,)`):
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced Ulysses partitioning.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -175,9 +180,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -188,6 +197,7 @@ def forward(
                 position_embeddings=position_embeddings,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
             )
 
         hidden_states = self.norm(hidden_states)
@@ -231,6 +241,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         r"""
@@ -238,6 +249,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced Ulysses partitioning.
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
             Indices of input tokens in the KV cache. Accepted only for HuggingFace API
             compatibility; prime-rl asserts `use_cache is None` since training does not
@@ -257,6 +270,7 @@ def forward(
             inputs_embeds=inputs_embeds,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = outputs.last_hidden_state
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
diff --git a/src/prime_rl/trainer/models/qwen3_5/attention.py b/src/prime_rl/trainer/models/qwen3_5/attention.py
index c8c8acb24c2..aca7970b66c 100644
--- a/src/prime_rl/trainer/models/qwen3_5/attention.py
+++ b/src/prime_rl/trainer/models/qwen3_5/attention.py
@@ -43,6 +43,7 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.LongTensor,
         max_seqlen: int,
+        cp_total_tokens: int | None = None,
     ) -> tuple[torch.Tensor, None]:
         batch_size, sequence_length, _ = hidden_states.shape
 
@@ -51,7 +52,7 @@ def forward(
             query_states, output_gate = query_states.view(
                 batch_size, sequence_length, self.num_heads, self.head_dim * 2
             ).chunk(2, dim=-1)
-            output_gate = output_gate.reshape(batch_size, sequence_length, -1)
+            output_gate = output_gate.reshape(batch_size, sequence_length, self.num_heads * self.head_dim)
         else:
             query_states = query_states.view(batch_size, sequence_length, self.num_heads, self.head_dim)
             output_gate = None
@@ -77,7 +78,8 @@ def forward(
             value_states[0],
             cu_seqlens,
             max_seqlen,
-        ).reshape(batch_size, sequence_length, -1)
+            **({"cp_total_tokens": cp_total_tokens} if cp_total_tokens is not None else {}),
+        ).reshape(batch_size, sequence_length, self.num_heads * self.head_dim)
         if output_gate is not None:
             attention_output = attention_output * output_gate.sigmoid()
         return self.o_proj(attention_output), None
diff --git a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
index 903ce6924c2..d5be64e9f42 100644
--- a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
+++ b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py
@@ -6,8 +6,10 @@
 from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 from torch import nn
 
+from prime_rl.trainer.models.layers.ulysses_attn import head_to_sequence_parallel, sequence_to_head_parallel
 from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
-from prime_rl.utils.cp import CPContext
+from prime_rl.utils.cp import CPContext, gather_for_cp
+from prime_rl.utils.sequence import CPPartition
 
 # Dynamo lowers all-gather to concatenation, then fails to copy the result into
 # FLA's stacked output buffer. Keep CP convolution eager until this is fixed:
@@ -54,8 +56,11 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         cu_seqlens: torch.LongTensor,
+        cp_total_tokens: int | None = None,
     ) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
+        # FLA caches metadata by tensor identity; checkpoint replay must execute the same operations.
+        cu_seqlens = cu_seqlens.clone()
         mixed_qkv = self.in_proj_qkv(hidden_states)
         output_gate = self.in_proj_z(hidden_states).reshape(
             batch_size, sequence_length, self.num_value_heads, self.value_head_dim
@@ -63,49 +68,115 @@ def forward(
         beta = self.in_proj_b(hidden_states).sigmoid()
         decay = -self.A_log.float().exp() * F.softplus(self.in_proj_a(hidden_states).float() + self.dt_bias)
 
+        partition = None
+        head_parallel = False
+        num_key_heads, num_value_heads = self.num_key_heads, self.num_value_heads
+        key_dim, value_dim = self.key_dim, self.value_dim
+        conv_weight = self.conv1d.weight.squeeze(1)
+        if self.cp_context.cp_enabled and cp_total_tokens is not None:
+            partition = CPPartition(cp_total_tokens, self.cp_context.cp_world_size)
+            group, degree, rank = self.cp_context.cp_group, self.cp_context.cp_world_size, self.cp_context.cp_rank
+            head_parallel = self.num_value_heads % degree == 0
+            if head_parallel:
+                heads_per_key = self.num_value_heads // self.num_key_heads
+                q, k, v = mixed_qkv.split((self.key_dim, self.key_dim, self.value_dim), dim=-1)
+                q = (
+                    q.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim)
+                    .repeat_interleave(heads_per_key, dim=2)
+                    .flatten(2)
+                )
+                k = (
+                    k.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim)
+                    .repeat_interleave(heads_per_key, dim=2)
+                    .flatten(2)
+                )
+                mixed_qkv = torch.cat(
+                    [sequence_to_head_parallel(x, group, degree, cp_total_tokens) for x in (q, k, v)], dim=-1
+                )
+                num_key_heads = num_value_heads = self.num_value_heads // degree
+                key_dim = num_key_heads * self.key_head_dim
+                value_dim = num_value_heads * self.value_head_dim
+                q_weight, k_weight, v_weight = conv_weight.split((self.key_dim, self.key_dim, self.value_dim))
+                head_slice = slice(rank * num_value_heads, (rank + 1) * num_value_heads)
+                weights = []
+                for weight, head_dim, repeats in (
+                    (q_weight, self.key_head_dim, heads_per_key),
+                    (k_weight, self.key_head_dim, heads_per_key),
+                    (v_weight, self.value_head_dim, 1),
+                ):
+                    weight = weight.reshape(-1, head_dim, self.conv_kernel_size).repeat_interleave(repeats, dim=0)
+                    weights.append(weight[head_slice].flatten(0, 1))
+                conv_weight = torch.cat(weights)
+                beta = sequence_to_head_parallel(beta, group, degree, cp_total_tokens)
+                decay = sequence_to_head_parallel(decay, group, degree, cp_total_tokens)
+                output_gate = sequence_to_head_parallel(output_gate.flatten(2), group, degree, cp_total_tokens).reshape(
+                    batch_size, cp_total_tokens, num_value_heads, self.value_head_dim
+                )
+            else:
+                mixed_qkv = gather_for_cp(mixed_qkv, group, cp_total_tokens)
+                beta = gather_for_cp(beta, group, cp_total_tokens)
+                decay = gather_for_cp(decay, group, cp_total_tokens)
+            sequence_length = cp_total_tokens
+
         context = None
-        if self.cp_context.cp_enabled:
+        if self.cp_context.cp_enabled and partition is None:
             context = build_cp_context(
                 cu_seqlens=cu_seqlens.to(device=hidden_states.device, dtype=torch.int32),
                 group=self.cp_context.cp_group,
                 conv1d_kernel_size=self.conv_kernel_size,
             )
 
-        convolution = causal_conv1d_with_context_parallelism if context is not None else causal_conv1d
-        mixed_qkv, _ = convolution(
-            x=mixed_qkv,
-            weight=self.conv1d.weight.squeeze(1),
-            bias=self.conv1d.bias,
-            activation=self.activation,
-            cu_seqlens=cu_seqlens,
-            cp_context=context,
-        )
+        if sequence_length == 0:
+            anchor = mixed_qkv.sum() + beta.sum() + decay.sum() + conv_weight.sum()
+            core_output = mixed_qkv[..., :value_dim].reshape(
+                batch_size, 0, num_value_heads, self.value_head_dim
+            ) + anchor.to(mixed_qkv.dtype)
+        else:
+            convolution = causal_conv1d_with_context_parallelism if context is not None else causal_conv1d
+            mixed_qkv, _ = convolution(
+                x=mixed_qkv,
+                weight=conv_weight,
+                bias=self.conv1d.bias,
+                activation=self.activation,
+                cu_seqlens=cu_seqlens,
+                cp_context=context,
+            )
 
-        query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
-        query = query.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim)
-        key = key.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim)
-        value = value.reshape(batch_size, sequence_length, self.num_value_heads, self.value_head_dim)
+            query, key, value = torch.split(mixed_qkv, [key_dim, key_dim, value_dim], dim=-1)
+            query = query.reshape(batch_size, sequence_length, num_key_heads, self.key_head_dim)
+            key = key.reshape(batch_size, sequence_length, num_key_heads, self.key_head_dim)
+            value = value.reshape(batch_size, sequence_length, num_value_heads, self.value_head_dim)
 
-        heads_per_key = self.num_value_heads // self.num_key_heads
-        if heads_per_key > 1:
-            query = query.repeat_interleave(heads_per_key, dim=2)
-            key = key.repeat_interleave(heads_per_key, dim=2)
+            heads_per_key = num_value_heads // num_key_heads
+            if heads_per_key > 1:
+                query = query.repeat_interleave(heads_per_key, dim=2)
+                key = key.repeat_interleave(heads_per_key, dim=2)
 
-        core_output, _ = chunk_gated_delta_rule(
-            q=query,
-            k=key,
-            v=value,
-            g=decay,
-            beta=beta,
-            use_qk_l2norm_in_kernel=True,
-            cu_seqlens=context.cu_seqlens if context is not None else cu_seqlens,
-            cp_context=context,
-        )
+            core_output, _ = chunk_gated_delta_rule(
+                q=query,
+                k=key,
+                v=value,
+                g=decay,
+                beta=beta,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=context.cu_seqlens if context is not None else cu_seqlens,
+                cp_context=context,
+            )
 
-        core_output = self.norm(
-            core_output.reshape(-1, self.value_head_dim),
-            output_gate.reshape(-1, self.value_head_dim),
-        ).reshape(batch_size, sequence_length, self.value_dim)
+        if partition is not None and not head_parallel:
+            core_output = partition.shard(core_output, self.cp_context.cp_rank)
+            sequence_length = core_output.shape[1]
+        if sequence_length == 0:
+            core_output = (core_output + output_gate + self.norm.weight.sum().to(core_output.dtype)).reshape(
+                batch_size, 0, value_dim
+            )
+        else:
+            core_output = self.norm(
+                core_output.reshape(-1, self.value_head_dim),
+                output_gate.reshape(-1, self.value_head_dim),
+            ).reshape(batch_size, sequence_length, value_dim)
+        if head_parallel:
+            core_output = head_to_sequence_parallel(core_output, group, degree, cp_total_tokens)
         return self.out_proj(core_output)
 
 
diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py
index ab269415c45..772dc8fe16f 100644
--- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py
+++ b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py
@@ -24,7 +24,7 @@
 )
 from prime_rl.trainer.models.qwen3_5.vision import Qwen3_5VisionModel
 from prime_rl.utils.cp import CPContext, setup_cp_attention_params, shard_for_cp, shard_position_ids_for_cp
-from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
+from prime_rl.utils.sequence import CPPartition, get_cu_seqlens_from_seq_lens
 
 
 class Qwen3_5SharedExpert(FeedForward):
@@ -93,6 +93,7 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.LongTensor,
         max_seqlen: int,
+        cp_total_tokens: int | None = None,
         routed_experts: torch.LongTensor | None = None,
     ) -> torch.Tensor:
         residual = hidden_states
@@ -101,6 +102,7 @@ def forward(
             hidden_states = self.linear_attn(
                 hidden_states,
                 cu_seqlens,
+                cp_total_tokens=cp_total_tokens,
             )
         else:
             hidden_states, _ = self.self_attn(
@@ -108,6 +110,7 @@ def forward(
                 position_embeddings,
                 cu_seqlens,
                 max_seqlen,
+                cp_total_tokens=cp_total_tokens,
             )
         hidden_states = residual + hidden_states
 
@@ -172,6 +175,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutput:
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
@@ -180,9 +184,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
         position_embeddings = self.rotary_emb(inputs_embeds, position_ids)
 
         hidden_states = inputs_embeds
@@ -193,6 +201,7 @@ def forward(
                 position_embeddings,
                 cu_seqlens,
                 max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=layer_routed_experts,
             )
         return BaseModelOutput(last_hidden_state=self.norm(hidden_states))
@@ -227,20 +236,18 @@ def prepare_inputs(
         if has_images:
             pixel_values = pixel_values.to(self.visual.dtype)
         else:
-            merge_size = self.config.vision_config.spatial_merge_size
-            num_patches = merge_size**2
             patch_dim = (
                 self.config.vision_config.in_channels
                 * self.config.vision_config.temporal_patch_size
                 * self.config.vision_config.patch_size**2
             )
             pixel_values = torch.zeros(
-                num_patches,
+                0,
                 patch_dim,
                 device=inputs_embeds.device,
                 dtype=self.visual.dtype,
             )
-            vision_grid = torch.tensor([[1, merge_size, merge_size]], device=inputs_embeds.device)
+            vision_grid = torch.empty((0, 3), device=inputs_embeds.device, dtype=torch.long)
 
         image_embeds = self.visual(pixel_values, vision_grid).pooler_output.to(inputs_embeds.dtype)
         if has_images:
@@ -276,6 +283,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> BaseModelOutput:
         inputs_embeds, position_ids = self.prepare_inputs(
             input_ids,
@@ -287,16 +295,25 @@ def forward(
         )
         if image_grid_thw is not None and self.cp_context.cp_enabled:
             rank, world_size = self.cp_context.cp_rank, self.cp_context.cp_world_size
-            setup_cp_attention_params(
-                position_ids,
-                cp_group=self.cp_context.cp_group,
-                cp_style=self.cp_context.cp_style,
-                seq_lens=seq_lens,
+            partition = None
+            if self.cp_context.cp_unpadded:
+                cp_total_tokens = inputs_embeds.shape[1]
+                partition = CPPartition(cp_total_tokens, world_size)
+            else:
+                setup_cp_attention_params(
+                    position_ids,
+                    cp_group=self.cp_context.cp_group,
+                    cp_style=self.cp_context.cp_style,
+                    seq_lens=seq_lens,
+                )
+            inputs_embeds = shard_for_cp(inputs_embeds, cp_rank=rank, cp_world_size=world_size, partition=partition)
+            position_ids = shard_position_ids_for_cp(
+                position_ids, cp_rank=rank, cp_world_size=world_size, partition=partition
             )
-            inputs_embeds = shard_for_cp(inputs_embeds, cp_rank=rank, cp_world_size=world_size)
-            position_ids = shard_position_ids_for_cp(position_ids, cp_rank=rank, cp_world_size=world_size)
             if routed_experts is not None:
-                routed_experts = shard_for_cp(routed_experts, cp_rank=rank, cp_world_size=world_size)
+                routed_experts = shard_for_cp(
+                    routed_experts, cp_rank=rank, cp_world_size=world_size, partition=partition
+                )
             seq_lens_are_pre_shard = True
 
         return self.language_model(
@@ -305,6 +322,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
 
@@ -343,6 +361,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> PrimeLmOutput:
         if self.is_vlm:
             outputs = self.model(
@@ -354,6 +373,7 @@ def forward(
                 routed_experts=routed_experts,
                 seq_lens=seq_lens,
                 seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+                cp_total_tokens=cp_total_tokens,
             )
         else:
             outputs = self.model(
@@ -363,6 +383,7 @@ def forward(
                 routed_experts=routed_experts,
                 seq_lens=seq_lens,
                 seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+                cp_total_tokens=cp_total_tokens,
             )
 
         return self.lm_head(
diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py
index 902c202839a..1d5bc7d90e2 100644
--- a/src/prime_rl/trainer/models/qwen3_5/vision.py
+++ b/src/prime_rl/trainer/models/qwen3_5/vision.py
@@ -54,7 +54,7 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
             self.patch_size,
             self.patch_size,
         )
-        return self.proj(pixel_values.to(self.proj.weight.dtype)).view(pixel_values.shape[0], -1)
+        return self.proj(pixel_values.to(self.proj.weight.dtype)).view(pixel_values.shape[0], self.proj.out_channels)
 
 
 class Qwen3_5VisionMLP(nn.Module):
@@ -100,7 +100,9 @@ def forward(
         query = ((query.float() * cos) + (rotate_half(query.float()) * sin)).to(query_dtype)
         key = ((key.float() * cos) + (rotate_half(key.float()) * sin)).to(query_dtype)
 
-        if self.attention_implementation == "flash_attention_4":
+        if sequence_length == 0:
+            attention_output = query + key + value
+        elif self.attention_implementation == "flash_attention_4":
             attention_output, _ = self.flash_attention(
                 query,
                 key,
@@ -116,12 +118,12 @@ def forward(
                 key,
                 value,
                 cu_seqlens,
-                cu_seqlens,
+                cu_seqlens.clone(),
                 max_sequence_length,
                 max_sequence_length,
                 causal=False,
             )
-        return self.proj(attention_output.reshape(sequence_length, -1))
+        return self.proj(attention_output.reshape(sequence_length, self.num_heads * self.head_dim))
 
 
 class Qwen3_5VisionBlock(nn.Module):
@@ -173,6 +175,8 @@ def dtype(self) -> torch.dtype:
         return self.patch_embed.proj.weight.dtype
 
     def rotary_embeddings(self, grid_thw: torch.Tensor) -> torch.Tensor:
+        if grid_thw.shape[0] == 0:
+            return self.rotary_pos_emb(0).repeat(1, 2)
         grids = grid_thw.tolist()
         frequency_table = self.rotary_pos_emb(max(max(height, width) for _, height, width in grids))
         coordinates = []
@@ -190,6 +194,8 @@ def rotary_embeddings(self, grid_thw: torch.Tensor) -> torch.Tensor:
         return frequency_table[torch.cat(coordinates)].flatten(1)
 
     def interpolated_position_embeddings(self, grid_thw: torch.Tensor) -> torch.Tensor:
+        if grid_thw.shape[0] == 0:
+            return self.pos_embed(torch.empty(0, dtype=torch.long, device=grid_thw.device))
         embeddings = []
         merge_size = self.spatial_merge_size
         for frames, height, width in grid_thw.tolist():
diff --git a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py
index f75f9cb387d..621dbe837ce 100644
--- a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py
+++ b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py
@@ -94,6 +94,7 @@ def forward(
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.LongTensor | None = None,
         max_seqlen: int | None = None,
+        cp_total_tokens: int | None = None,
         routed_experts: Optional[torch.LongTensor] = None,
     ) -> torch.FloatTensor:
         residual = hidden_states
@@ -104,6 +105,7 @@ def forward(
             position_embeddings=position_embeddings,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            cp_total_tokens=cp_total_tokens,
         )
         hidden_states = residual + hidden_states
 
@@ -185,6 +187,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
     ) -> MoeModelOutputWithPast:
         """
         routed_experts (`torch.LongTensor` of shape `(batch_size, sequence_length, num_hidden_layers, num_experts_per_tok)`, *optional*):
@@ -193,6 +196,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced Ulysses partitioning.
         """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
@@ -202,9 +207,13 @@ def forward(
 
         cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens(
             seq_lens.to(device=inputs_embeds.device),
-            total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1],
+            total_tokens=None if seq_lens_are_pre_shard and seq_lens.numel() else inputs_embeds.shape[1],
         )
         torch._dynamo.mark_dynamic(cu_seqlens, 0)
+        if not self.cp_context.cp_unpadded:
+            cp_total_tokens = None
+        elif cp_total_tokens is None:
+            cp_total_tokens = int(cu_seqlens[-1].item())
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -216,6 +225,7 @@ def forward(
                 position_embeddings=position_embeddings,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
+                cp_total_tokens=cp_total_tokens,
                 routed_experts=routed_experts_layer,
             )
 
@@ -267,6 +277,7 @@ def forward(
         *,
         seq_lens: torch.LongTensor,
         seq_lens_are_pre_shard: bool = False,
+        cp_total_tokens: int | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> PrimeLmOutput:
         r"""
@@ -274,6 +285,8 @@ def forward(
             Per-document lengths of the packed row (PrimeRL packed-batch contract).
         seq_lens_are_pre_shard (`bool`, *optional*, defaults to `False`):
             Whether `seq_lens` holds pre-CP-shard (global) document boundaries.
+        cp_total_tokens (`int`, *optional*):
+            Full packed row length for balanced Ulysses partitioning.
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
             Indices of input tokens in the KV cache. Accepted only for HuggingFace API
             compatibility — prime-rl asserts `use_cache is None` since training does not
@@ -323,6 +336,7 @@ def forward(
             routed_experts=routed_experts,
             seq_lens=seq_lens,
             seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+            cp_total_tokens=cp_total_tokens,
         )
 
         hidden_states = outputs.last_hidden_state
diff --git a/src/prime_rl/trainer/parallel_dims.py b/src/prime_rl/trainer/parallel_dims.py
index 86a4a255faa..de8c1208254 100644
--- a/src/prime_rl/trainer/parallel_dims.py
+++ b/src/prime_rl/trainer/parallel_dims.py
@@ -332,7 +332,7 @@ def get_parallel_dims(config: ModelConfig, seq_len: int | None = None) -> Parall
     )
 
     # Validate sequence length against parallel dimensions requirements
-    if seq_len is not None and seq_len % parallel_dims.seq_len_divisor != 0:
+    if seq_len is not None and not config.cp_unpadded and seq_len % parallel_dims.seq_len_divisor != 0:
         raise ValueError(
             f"Sequence length ({seq_len}) must be divisible by "
             f"seq_len_divisor ({parallel_dims.seq_len_divisor}) for the given parallel dimensions. "
diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py
index e2b360dc617..c71e39b5ec7 100644
--- a/src/prime_rl/trainer/rl/data.py
+++ b/src/prime_rl/trainer/rl/data.py
@@ -20,6 +20,8 @@
 class TensorMicroBatch(TypedDict):
     """A micro batch of data for training."""
 
+    inactive: bool
+
     # Token level
     input_ids: Int[Tensor, "batch seq"]
     position_ids: Int[Tensor, "batch seq"]
@@ -117,6 +119,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc
         inference_logprobs = torch.randn(input_ids.shape[0], generator=generator)
 
         return {
+            "inactive": False,
             "input_ids": input_ids.unsqueeze(0),
             "position_ids": position_ids.unsqueeze(0),
             "advantages": advantages.unsqueeze(0),
@@ -141,6 +144,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc
 
     def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch:
         return {
+            "inactive": False,
             "input_ids": torch.randint(
                 0,
                 100,
@@ -221,6 +225,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch:
             padded[np.arange(max_mask_size)[None, :] < counts[:, None]] = ids
             sampling_mask = torch.from_numpy(padded).unsqueeze(0)
         return TensorMicroBatch(
+            inactive=micro_batch.inactive,
             input_ids=torch.tensor(micro_batch.input_ids, dtype=torch.long).unsqueeze(0),
             position_ids=torch.tensor(micro_batch.position_ids, dtype=torch.long).unsqueeze(0),
             advantages=torch.tensor(micro_batch.advantages, dtype=torch.float).unsqueeze(0),
diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py
index b1f59b836bb..79e11a9fe0b 100644
--- a/src/prime_rl/trainer/rl/loss.py
+++ b/src/prime_rl/trainer/rl/loss.py
@@ -126,7 +126,7 @@ def shift_tensor_right(t: Float[Tensor, "batch seq"], pad_value: float | None =
     """
     if pad_value is None:
         pad_value = 0.0
-    return torch.cat([torch.full((t.shape[0], 1), pad_value, device=t.device, dtype=t.dtype), t[:, :-1]], dim=1)
+    return torch.cat([torch.full_like(t[:, :1], pad_value), t[:, :-1]], dim=1)
 
 
 def _safe_mean(values: Tensor, mask: Tensor) -> Tensor:
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 990004e2fae..102919f636f 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -19,6 +19,7 @@
 from prime_rl.configs.trainer import TrainerConfig
 from prime_rl.trainer.rl.data import DataLoader, FakeDataLoader
 from prime_rl.utils.cp import (
+    CPPartition,
     gather_for_cp,
     gather_for_cp_wo_grad,
     setup_context_parallel,
@@ -306,7 +307,13 @@ def train(config: TrainerConfig):
             memory_profiler = MemoryProfiler(progress.step, config.memory_profiler_path)
 
         forward_backward_start_time = time.perf_counter()
-        seq_len = micro_batches[0]["input_ids"].shape[1]
+        seq_len = max((batch["input_ids"].shape[1] for batch in micro_batches), default=0)
+        num_local_tokens = sum(batch["input_ids"].shape[1] for batch in micro_batches)
+        inactive_slots = sum(batch.get("inactive", False) for batch in micro_batches)
+        real_micro_batches = len(micro_batches) - inactive_slots
+        real_samples = sum(
+            len(batch["sequence_lengths"]) for batch in micro_batches if not batch.get("inactive", False)
+        )
 
         # Normalize each loss component by its own global (dp_cp) token count, so every rank
         # divides by the same denominator. With a per-rank denominator, ranks with fewer loss
@@ -326,11 +333,25 @@ def train(config: TrainerConfig):
             if micro_batch["ref_kl_weights"] is not None:
                 local_ref_kl_scale += int((micro_batch["ref_kl_weights"] != 0).sum())
         global_scales = torch.tensor(
-            [local_rl_scale, local_ce_scale, local_ref_kl_scale], dtype=torch.int64, device="cuda"
+            [
+                local_rl_scale,
+                local_ce_scale,
+                local_ref_kl_scale,
+                num_local_tokens,
+                inactive_slots,
+                real_micro_batches,
+                real_samples,
+            ],
+            dtype=torch.int64,
+            device="cuda",
         )
         dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group()
         dist.all_reduce(global_scales, op=dist.ReduceOp.SUM, group=dp_cp_group)
-        rl_scale, ce_scale, ref_kl_scale = (max(scale, 1) for scale in global_scales.tolist())
+        counts = global_scales.tolist()
+        if counts[3] == 0:
+            raise ValueError("An entirely empty global step cannot advance the optimizer")
+        rl_scale, ce_scale, ref_kl_scale = (max(scale, 1) for scale in counts[:3])
+        num_tokens = counts[3] // parallel_dims.cp
         prepare_gradient_offload(
             gradient_manager,
             parallel_dims.fsdp_gradient_divide_factor,
@@ -345,7 +366,13 @@ def train(config: TrainerConfig):
         cp_size = parallel_dims.cp
 
         for micro_step, micro_batch in enumerate(micro_batches):
+            inactive = micro_batch.get("inactive", False)
+            if inactive and not config.model.inactive_micro_batches:
+                raise ValueError("Received an inactive slot without model.inactive_micro_batches")
+            if inactive and micro_batch["input_ids"].numel():
+                raise ValueError("An inactive slot must not contain tokens")
             input_ids = micro_batch["input_ids"].to("cuda")
+            cp_partition = CPPartition(input_ids.shape[1], cp_size) if config.model.cp_unpadded else None
             position_ids = micro_batch["position_ids"].to("cuda")
             advantages = micro_batch["advantages"].to("cuda")
             loss_mask = micro_batch["loss_mask"].to("cuda")
@@ -360,7 +387,7 @@ def train(config: TrainerConfig):
                 micro_batch["routed_experts"].to("cuda") if micro_batch["routed_experts"] is not None else None
             )
 
-            if routed_experts is None and config.enable_router_replay:
+            if routed_experts is None and config.enable_router_replay and not inactive:
                 raise ValueError(
                     "You must set `enable_return_routed_experts=True` in the inference config or pass `--enable-return-routed-experts` to vLLM server to use router replay."
                 )
@@ -413,15 +440,20 @@ def train(config: TrainerConfig):
                         cp_group,
                         seq_lens=seq_lens,
                         cp_style=config.model.cp_style,
+                        partition=cp_partition,
                     )
                 seq_lens_are_pre_shard = True
-                labels = shard_for_cp(labels, cp_rank=cp_rank, cp_world_size=cp_size)
+                labels = shard_for_cp(labels, cp_rank=cp_rank, cp_world_size=cp_size, partition=cp_partition)
                 if routed_experts is not None and not defer_vlm_cp_to_model:
-                    routed_experts = shard_for_cp(routed_experts, cp_rank=cp_rank, cp_world_size=cp_size)
+                    routed_experts = shard_for_cp(
+                        routed_experts, cp_rank=cp_rank, cp_world_size=cp_size, partition=cp_partition
+                    )
                 if sampling_mask is not None:
                     # The LM head consumes masks after any deferred VLM sharding, so
                     # they must follow the label shard rather than the input shard.
-                    sampling_mask = shard_for_cp(sampling_mask, cp_rank=cp_rank, cp_world_size=cp_size)
+                    sampling_mask = shard_for_cp(
+                        sampling_mask, cp_rank=cp_rank, cp_world_size=cp_size, partition=cp_partition
+                    )
 
             if config.model.lora:
                 lora_num_tokens = micro_batch["lora_num_tokens"].to("cuda")
@@ -429,7 +461,8 @@ def train(config: TrainerConfig):
                     chunk_size = labels.shape[1]
                     # Convert to cumsum, adjust for CP chunk, convert back to num_tokens
                     cu_offsets = lora_num_tokens.cumsum(dim=0, dtype=torch.int32)
-                    adjusted_cu = torch.clip(cu_offsets - chunk_size * cp_rank, min=0, max=chunk_size)
+                    offset = cp_partition.offsets[cp_rank] if cp_partition is not None else chunk_size * cp_rank
+                    adjusted_cu = torch.clip(cu_offsets - offset, min=0, max=chunk_size)
                     lora_num_tokens = torch.diff(
                         adjusted_cu, prepend=torch.tensor([0], device=adjusted_cu.device, dtype=adjusted_cu.dtype)
                     )
@@ -439,7 +472,9 @@ def train(config: TrainerConfig):
 
             # Shard temperatures for context parallelism if enabled
             if cp_enabled:
-                temperatures = shard_for_cp(temperatures, cp_rank=cp_rank, cp_world_size=cp_size)
+                temperatures = shard_for_cp(
+                    temperatures, cp_rank=cp_rank, cp_world_size=cp_size, partition=cp_partition
+                )
 
             if sampling_mask is not None:
                 assert sampling_mask.shape[:2] == labels.shape, (
@@ -460,6 +495,7 @@ def train(config: TrainerConfig):
                     mm_token_type_ids=mm_token_type_ids,
                     seq_lens=seq_lens,
                     seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+                    cp_total_tokens=cp_partition.total_tokens if cp_partition is not None and cp_enabled else None,
                     routed_experts=routed_experts,
                     sampling_mask=sampling_mask,
                 )
@@ -478,8 +514,9 @@ def train(config: TrainerConfig):
             # else: FusedOutputLinear was used - logprobs already computed with per-token temperatures
 
             if cp_enabled:
-                out["logprobs"] = gather_for_cp(out["logprobs"], cp_group)
-                out["entropy"] = gather_for_cp_wo_grad(out["entropy"], cp_size, cp_group)
+                total_tokens = cp_partition.total_tokens if cp_partition is not None else None
+                out["logprobs"] = gather_for_cp(out["logprobs"], cp_group, total_tokens)
+                out["entropy"] = gather_for_cp_wo_grad(out["entropy"], cp_size, cp_group, total_tokens)
 
             vocab_size = getattr(model.config, "vocab_size", None) or model.config.text_config.vocab_size
             # This is not really necessary as the first token should be masked out, but we do it anyway to be sure
@@ -492,20 +529,25 @@ def train(config: TrainerConfig):
 
             # Compute loss
             sequence_lengths = micro_batch["sequence_lengths"]
-            loss, loss_tensors = compute_loss(
-                trainer_logprobs=out["logprobs"].squeeze().split(sequence_lengths),
-                inference_logprobs=inference_logprobs.squeeze().split(sequence_lengths),
-                ref_logprobs=ref_logprobs.squeeze().split(sequence_lengths) if ref_logprobs is not None else None,
-                advantages=advantages.squeeze().split(sequence_lengths),
-                loss_mask=loss_mask.squeeze().split(sequence_lengths),
-                rl_weights=rl_weights.squeeze().split(sequence_lengths) if rl_weights is not None else None,
-                ce_weights=ce_weights.squeeze().split(sequence_lengths) if ce_weights is not None else None,
-                ref_kl_weights=ref_kl_weights.squeeze().split(sequence_lengths) if ref_kl_weights is not None else None,
-                rl_loss_fn=rl_loss_fn,
-                rl_scale=rl_scale,
-                ce_scale=ce_scale,
-                ref_kl_scale=ref_kl_scale,
-            )
+            if inactive:
+                loss, loss_tensors = out["logprobs"].sum(), {}
+            else:
+                loss, loss_tensors = compute_loss(
+                    trainer_logprobs=out["logprobs"].squeeze(0).split(sequence_lengths),
+                    inference_logprobs=inference_logprobs.squeeze(0).split(sequence_lengths),
+                    ref_logprobs=ref_logprobs.squeeze(0).split(sequence_lengths) if ref_logprobs is not None else None,
+                    advantages=advantages.squeeze(0).split(sequence_lengths),
+                    loss_mask=loss_mask.squeeze(0).split(sequence_lengths),
+                    rl_weights=rl_weights.squeeze(0).split(sequence_lengths) if rl_weights is not None else None,
+                    ce_weights=ce_weights.squeeze(0).split(sequence_lengths) if ce_weights is not None else None,
+                    ref_kl_weights=ref_kl_weights.squeeze(0).split(sequence_lengths)
+                    if ref_kl_weights is not None
+                    else None,
+                    rl_loss_fn=rl_loss_fn,
+                    rl_scale=rl_scale,
+                    ce_scale=ce_scale,
+                    ref_kl_scale=ref_kl_scale,
+                )
 
             # Backward pass
             with maybe_record_function("backward"):
@@ -518,7 +560,8 @@ def train(config: TrainerConfig):
             # Add relevant tensors to tensor dict for logging purposes
             entropy = out["entropy"][loss_mask].detach().to("cpu")
             tensors["entropy/all"].append(entropy)
-            tensors["loss"].append(loss.detach().to("cpu").unsqueeze(0))
+            if not inactive:
+                tensors["loss"].append(loss.detach().to("cpu").unsqueeze(0))
 
             env_names = micro_batch["env_names"]
             masked_env_names = [env_name for env_name, keep in zip(env_names, loss_mask.flatten().tolist()) if keep]
@@ -534,7 +577,7 @@ def train(config: TrainerConfig):
             # whose action component is ce (frozen-model tokens).
             if rl_weights is None and ref_kl_weights is None:
                 mismatch_mask = loss_mask
-                has_mismatch_tokens = True
+                has_mismatch_tokens = not inactive
             else:
                 sampled_mask = (rl_weights != 0) if rl_weights is not None else loss_mask
                 if ref_kl_weights is not None:
@@ -567,7 +610,11 @@ def train(config: TrainerConfig):
                 tensors[key].append(loss_tensor.detach().to("cpu"))
 
             # Debug log with *local, micro step* stats
-            micro_step_message = f"Micro Step {micro_step + 1}/{len(micro_batches)} | Loss {tensors['loss'][-1].mean().item():.4f} | Entropy {tensors['entropy/all'][-1].mean().item():.4f}"
+            micro_step_message = f"Micro Step {micro_step + 1}/{len(micro_batches)}"
+            if inactive:
+                micro_step_message += " | Inactive"
+            else:
+                micro_step_message += f" | Loss {loss.item():.4f} | Entropy {entropy.mean().item():.4f}"
             if has_mismatch_tokens:
                 micro_step_message += f" | Mismatch KL {tensors['mismatch_kl/all'][-1].mean().item():.4f}"
             logger.debug(micro_step_message)
@@ -590,6 +637,8 @@ def train(config: TrainerConfig):
 
         # Update learning rate scheduler
         scheduler.step()
+        progress.total_tokens += num_tokens
+        progress.total_samples += counts[6] // parallel_dims.cp
 
         current_lr = optimizer.param_groups[0]["lr"]
         forward_backward_time = time.perf_counter() - forward_backward_start_time
@@ -636,10 +685,6 @@ def train(config: TrainerConfig):
         tensor_stats = tensors.compute_stats()
 
         # Compute step metrics
-        num_local_tokens = seq_len * batch_size
-        num_tokens = parallel_dims.get_mesh("dp").size() * num_local_tokens
-        progress.total_tokens += num_tokens
-        progress.total_samples += batch_size
         perf_counter = get_perf_counter(model, seq_len)
         throughput = perf_counter.get_step_tokens_per_second(num_tokens, forward_backward_time)
         mfu = perf_counter.get_step_mfu(num_tokens, forward_backward_time)
@@ -670,6 +715,9 @@ def train(config: TrainerConfig):
         # Log performance metrics
         perf_metrics = {
             "perf/throughput": throughput,
+            "perf/input_tokens": num_tokens,
+            "perf/inactive_micro_batches": counts[4] // parallel_dims.cp,
+            "perf/real_micro_batches": counts[5] // parallel_dims.cp,
             "perf/throughput_per_gpu": throughput / world.world_size,
             "perf/mfu": mfu,
             "perf/peak_memory": peak_memory,
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index 7ff71cd7155..a8f961a24e2 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -17,7 +17,7 @@
 from prime_rl.configs.sft import SFTConfig
 from prime_rl.configs.trainer import CheckpointConfig
 from prime_rl.transports.weights import prune_broadcasts_beyond, setup_weight_sender
-from prime_rl.utils.cp import setup_context_parallel, setup_cp_params, shard_for_cp
+from prime_rl.utils.cp import CPPartition, setup_context_parallel, setup_cp_params, shard_for_cp
 from prime_rl.trainer.lora import get_lora_state
 from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX
 from prime_rl.trainer.models.layers.lora import set_lora_num_tokens
@@ -253,10 +253,9 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
             mm_type_ids = mm_type_ids.to("cuda", non_blocking=True)
 
         seq_lens_are_pre_shard = False
+        cp_partition = CPPartition(input_ids.shape[1], cp_size) if config.model.cp_unpadded else None
 
         if cp_enabled:
-            # CP requires the sequence length to be divisible by cp_size. CatDataset
-            # pads every pack to seq_len; shard_for_cp raises on violations.
             defer_vlm_cp_to_model = (
                 mm_kwargs is not None and "image_grid_thw" in mm_kwargs and config.model.cp_style == "ulysses"
             )
@@ -269,10 +268,11 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
                     cp_group,
                     seq_lens=seq_lens,
                     cp_style=config.model.cp_style,
+                    partition=cp_partition,
                 )
             seq_lens_are_pre_shard = True
-            target_ids = shard_for_cp(target_ids, cp_rank=cp_rank, cp_world_size=cp_size)
-            loss_mask = shard_for_cp(loss_mask, cp_rank=cp_rank, cp_world_size=cp_size)
+            target_ids = shard_for_cp(target_ids, cp_rank=cp_rank, cp_world_size=cp_size, partition=cp_partition)
+            loss_mask = shard_for_cp(loss_mask, cp_rank=cp_rank, cp_world_size=cp_size, partition=cp_partition)
 
         if config.model.lora is not None:
             set_lora_num_tokens(torch.full((1,), loss_mask.numel(), dtype=torch.int32, device="cuda"))
@@ -292,6 +292,7 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
                 mm_kwargs=mm_kwargs,
                 mm_token_type_ids=mm_type_ids,
                 seq_lens_are_pre_shard=seq_lens_are_pre_shard,
+                cp_total_tokens=cp_partition.total_tokens if cp_partition is not None and cp_enabled else None,
             )
             loss_sum = out["loss"]
 
diff --git a/src/prime_rl/transports/batch/types.py b/src/prime_rl/transports/batch/types.py
index 4fe5225dd6c..f0992e84b56 100644
--- a/src/prime_rl/transports/batch/types.py
+++ b/src/prime_rl/transports/batch/types.py
@@ -115,3 +115,6 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
     # (e.g. a dummy micro batch). ``None`` when no packed sample carried one.
     trace_ids: list[str] | None = None
     branch_indices: list[int] | None = None
+
+    # Empty execution slot: no tokens or samples, but full distributed participation.
+    inactive: bool = False
diff --git a/src/prime_rl/utils/cp.py b/src/prime_rl/utils/cp.py
index ab5088fc783..8f0c5d36007 100644
--- a/src/prime_rl/utils/cp.py
+++ b/src/prime_rl/utils/cp.py
@@ -11,9 +11,9 @@
 import torch.nn as nn
 from ring_flash_attn import substitute_hf_flash_attn, update_ring_flash_attn_params
 
-from prime_rl.trainer.distributed.collectives import all_gather
+from prime_rl.trainer.distributed.collectives import all_gather, all_gather_cp
 from prime_rl.utils.logger import get_logger
-from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
+from prime_rl.utils.sequence import CPPartition, get_cu_seqlens_from_seq_lens
 
 if TYPE_CHECKING:
     # `prime_rl.trainer.models` imports this module, so importing the model base eagerly would
@@ -31,6 +31,7 @@ class CPContext:
     cp_rank: int = 0
     cp_world_size: int = 1
     cp_style: CPStyle | None = None
+    cp_unpadded: bool = False
 
     @property
     def cp_enabled(self) -> bool:
@@ -41,22 +42,22 @@ def setup_context_parallel(model: nn.Module, config: ModelConfig, parallel_dims:
     cp_group = parallel_dims.world_mesh["cp"].get_group()
     cp_rank = parallel_dims.world_mesh["cp"].get_local_rank()
 
-    if config.cp_style == "ring":
+    if not config.cp_unpadded and config.cp_style == "ring":
         # Delayed imports: both modules live under trainer.models, which imports back into
         # prime_rl.utils — a top-level import would deadlock at startup.
         from prime_rl.trainer.models.layers.attn import substitute_ring_attn
 
         substitute_hf_flash_attn(cp_group, heads_k_stride=1)
         substitute_ring_attn(cp_group, heads_k_stride=1, attn_impl=config.attn)
-    elif config.cp_style == "ulysses":
+    elif not config.cp_unpadded and config.cp_style == "ulysses":
         from prime_rl.trainer.models.layers.ulysses_attn import substitute_hf_ulysses_attn, substitute_ulysses_attn
 
         substitute_hf_ulysses_attn(cp_group)
         substitute_ulysses_attn(cp_group, attn_impl=config.attn)
-    else:
+    elif config.cp_style not in ("ring", "ulysses"):
         raise ValueError(f"Unknown cp_style: {config.cp_style}")
 
-    cp_context = CPContext(cp_group, cp_rank, parallel_dims.cp, config.cp_style)
+    cp_context = CPContext(cp_group, cp_rank, parallel_dims.cp, config.cp_style, config.cp_unpadded)
     for module in model.modules():
         if not hasattr(module, "cp_context"):
             continue
@@ -70,7 +71,9 @@ def setup_context_parallel(model: nn.Module, config: ModelConfig, parallel_dims:
     get_logger().info(f"Configured {config.cp_style} context parallelism (cp={parallel_dims.cp})")
 
 
-def shard_for_cp(t: torch.Tensor, cp_rank: int, cp_world_size: int, seq_dim: int = 1) -> torch.Tensor:
+def shard_for_cp(
+    t: torch.Tensor, cp_rank: int, cp_world_size: int, seq_dim: int = 1, *, partition: CPPartition | None = None
+) -> torch.Tensor:
     """
     Shard a tensor for context parallelism.
     Args:
@@ -83,6 +86,10 @@ def shard_for_cp(t: torch.Tensor, cp_rank: int, cp_world_size: int, seq_dim: int
 
     if seq_dim == 1 and t.shape[0] != 1:
         raise ValueError(f"For CP, tensor must have batch dimension 1, got shape={tuple(t.shape)}")
+    if partition is not None:
+        if partition.degree != cp_world_size:
+            raise ValueError("CP partition degree must match the process group")
+        return partition.shard(t, cp_rank, seq_dim)
     if t.shape[seq_dim] % cp_world_size != 0:
         raise ValueError(
             f"CP requires sequence dimension {seq_dim} to be divisible by cp size: "
@@ -95,17 +102,26 @@ def shard_for_cp(t: torch.Tensor, cp_rank: int, cp_world_size: int, seq_dim: int
     return chunked_t[cp_rank]
 
 
-def shard_position_ids_for_cp(position_ids: torch.Tensor, cp_rank: int, cp_world_size: int) -> torch.Tensor:
+def shard_position_ids_for_cp(
+    position_ids: torch.Tensor, cp_rank: int, cp_world_size: int, *, partition: CPPartition | None = None
+) -> torch.Tensor:
     if position_ids.ndim == 3:
-        return shard_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size, seq_dim=2)
-    return shard_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size)
+        return shard_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size, seq_dim=2, partition=partition)
+    return shard_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size, partition=partition)
 
 
-def gather_for_cp(t: torch.Tensor, cp_group: dist.ProcessGroup) -> torch.Tensor:
+def gather_for_cp(t: torch.Tensor, cp_group: dist.ProcessGroup, total_tokens: int | None = None) -> torch.Tensor:
+    if total_tokens is not None:
+        return all_gather_cp(t, 1, total_tokens, cp_group)
     return all_gather(t, 1, cp_group)
 
 
-def gather_for_cp_wo_grad(t: torch.Tensor, cp_world_size: int, cp_group: dist.ProcessGroup) -> torch.Tensor:
+@torch.no_grad()
+def gather_for_cp_wo_grad(
+    t: torch.Tensor, cp_world_size: int, cp_group: dist.ProcessGroup, total_tokens: int | None = None
+) -> torch.Tensor:
+    if total_tokens is not None:
+        return all_gather_cp(t, 1, total_tokens, cp_group)
     empty_like_t = [torch.empty_like(t) for _ in range(cp_world_size)]
     dist.all_gather(empty_like_t, t, group=cp_group)
     return torch.cat(empty_like_t, dim=1)
@@ -120,6 +136,7 @@ def setup_cp_params(
     *,
     seq_lens: torch.Tensor,
     cp_style: CPStyle = "ring",
+    partition: CPPartition | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Prepare the input for context parallelism and set required attention params.
@@ -132,10 +149,13 @@ def setup_cp_params(
     Returns the sequence-sharded input_ids and position_ids — the rest of the
     model still runs sequence-sharded; only attention sees the full sequence.
     """
-    setup_cp_attention_params(position_ids, cp_group=cp_group, cp_style=cp_style, seq_lens=seq_lens)
+    if partition is None:
+        setup_cp_attention_params(position_ids, cp_group=cp_group, cp_style=cp_style, seq_lens=seq_lens)
 
-    input_ids = shard_for_cp(input_ids, cp_rank=cp_rank, cp_world_size=cp_world_size)
-    position_ids = shard_position_ids_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size)
+    input_ids = shard_for_cp(input_ids, cp_rank=cp_rank, cp_world_size=cp_world_size, partition=partition)
+    position_ids = shard_position_ids_for_cp(
+        position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size, partition=partition
+    )
     return input_ids, position_ids
 
 
diff --git a/src/prime_rl/utils/sequence.py b/src/prime_rl/utils/sequence.py
index 9ef20f5574f..1f625b903a6 100644
--- a/src/prime_rl/utils/sequence.py
+++ b/src/prime_rl/utils/sequence.py
@@ -1,8 +1,39 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
+
 import torch
 
 
+@dataclass(frozen=True)
+class CPPartition:
+    """Contiguous balanced token ranges for one context-parallel forward."""
+
+    total_tokens: int
+    degree: int
+
+    def __post_init__(self):
+        if self.total_tokens < 0 or self.degree < 1:
+            raise ValueError("CP partition requires a nonnegative length and a positive degree")
+
+    @property
+    def lengths(self) -> tuple[int, ...]:
+        quotient, remainder = self.total_tokens // self.degree, self.total_tokens % self.degree
+        return tuple(quotient + (rank < remainder) for rank in range(self.degree))
+
+    @property
+    def offsets(self) -> tuple[int, ...]:
+        quotient, remainder = self.total_tokens // self.degree, self.total_tokens % self.degree
+        return tuple(rank * quotient + min(rank, remainder) for rank in range(self.degree + 1))
+
+    def shard(self, tensor: torch.Tensor, rank: int, dim: int = 1) -> torch.Tensor:
+        if tensor.shape[dim] != self.total_tokens:
+            raise ValueError(f"CP partition length {self.total_tokens} != tensor length {tensor.shape[dim]}")
+        if not 0 <= rank < self.degree:
+            raise ValueError(f"CP rank {rank} is outside degree {self.degree}")
+        return tensor.narrow(dim, self.offsets[rank], self.lengths[rank])
+
+
 def get_cu_seqlens_from_position_ids(position_ids: torch.Tensor) -> tuple[torch.Tensor, int]:
     """Build local-relative cumulative sequence lengths from packed position ids.
 
@@ -39,6 +70,8 @@ def get_cu_seqlens_from_seq_lens(seq_lens: torch.Tensor, total_tokens: int | Non
     if seq_lens.ndim != 1:
         raise ValueError(f"seq_lens must be 1D, got shape={tuple(seq_lens.shape)}")
     if seq_lens.numel() == 0:
+        if total_tokens == 0:
+            return torch.zeros(1, dtype=torch.int32, device=seq_lens.device), 0
         raise ValueError("seq_lens must not be empty")
     if bool((seq_lens <= 0).any().item()):
         raise ValueError(f"seq_lens must be positive, got {seq_lens.tolist()}")
diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py
index 4622b1146f6..5166a08b2b9 100644
--- a/tests/unit/orchestrator/test_batch.py
+++ b/tests/unit/orchestrator/test_batch.py
@@ -1,5 +1,6 @@
 from types import SimpleNamespace
 
+import msgspec
 import numpy as np
 import pytest
 
@@ -89,8 +90,9 @@ def test_training_sample_requires_env_name():
 @pytest.mark.parametrize(
     ("rollout_count", "num_train_workers", "expected_batches_per_worker"), [(4, 2, 2), (5, 2, 3), (7, 1, 7), (11, 4, 3)]
 )
+@pytest.mark.parametrize("inactive", [False, True])
 def test_prepare_batch_balances_micro_batches_across_workers(
-    make_training_example, rollout_count, num_train_workers, expected_batches_per_worker
+    make_training_example, rollout_count, num_train_workers, expected_batches_per_worker, inactive
 ):
     examples = [make_training_example() for i in range(rollout_count)]
 
@@ -99,11 +101,13 @@ def test_prepare_batch_balances_micro_batches_across_workers(
         seq_len=4,
         num_train_workers=num_train_workers,
         bin_cost=build_bin_cost(None),
+        inactive_micro_batches=inactive,
     )
 
     assert all(len(worker_batches) == expected_batches_per_worker for worker_batches in batches_per_gpu)
 
     flat_batches = _flatten_batches(batches_per_gpu)
+    flat_batches = msgspec.msgpack.decode(msgspec.msgpack.encode(flat_batches), type=list[MicroBatch])
     assert len(examples) <= len(flat_batches) < len(examples) + num_train_workers
 
     # Identify real vs padding batches by content, not position — the packer
@@ -120,6 +124,9 @@ def test_prepare_batch_balances_micro_batches_across_workers(
 
     # Verify padded batches have zero advantages and loss mask
     for batch in dummy_batches:
+        assert batch.inactive == inactive
+        if inactive:
+            assert batch.input_ids == batch.sequence_lengths == batch.seq_lens == []
         assert sum(1 for advantage in batch.advantages if advantage != 0.0) == 0
         assert sum(1 for loss_mask in batch.loss_mask if loss_mask) == 0
 
@@ -497,6 +504,27 @@ def test_prepare_batch_packs_multimodal_with_text():
     assert batch.env_names == ["mm-env"] * 3 + ["text-env"] * 2
 
 
+def test_split_to_align_groups_modalities_before_alignment():
+    samples = [make_sized_training_example(n) for n in [6, 6, 4, 4, 4, 4]]
+    for index, sample in enumerate(samples):
+        sample.trace_id = f"sample-{index}"
+        sample.branch_index = 0
+        if index < 2:
+            sample.mm_refs = MMRefs(images=[MMImageRef(url=f"fixture://{index}", offset=1, length=1)])
+            sample.mm_token_type_ids = [0, 1] + [0] * (len(sample.token_ids) - 2)
+    lanes = prepare_batch(samples, 10, 4, build_bin_cost(None))
+    batches = _flatten_batches(lanes)
+    real = [batch for batch in batches if batch.trace_ids is not None]
+    dummy = [batch for batch in batches if batch.trace_ids is None]
+    assert len(dummy) == 2
+    assert all(batch.inactive and not batch.input_ids for batch in dummy)
+    assert sum(len(batch.input_ids) for batch in real) == 28
+    assert sum(sum(batch.loss_mask) for batch in real) == 6
+    assert sorted(trace for batch in real for trace in batch.trace_ids) == [f"sample-{i}" for i in range(6)]
+    assert all(_is_multimodal_sample(lane[0]) or lane[0].inactive for lane in lanes)
+    assert all(not _is_multimodal_sample(lane[1]) and not lane[1].inactive for lane in lanes)
+
+
 def test_split_to_align_splits_multimodal_bins():
     def make_mm_sample(token: int) -> TrainingSample:
         return TrainingSample(
diff --git a/tests/unit/train/test_cp_collectives.py b/tests/unit/train/test_cp_collectives.py
new file mode 100644
index 00000000000..9bf9af91098
--- /dev/null
+++ b/tests/unit/train/test_cp_collectives.py
@@ -0,0 +1,543 @@
+"""Run on four GPUs with ``uv run torchrun --standalone --nproc-per-node=4 -m pytest -s``."""
+
+import copy
+import datetime
+import json
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from prime_rl.trainer.distributed.collectives import all_gather_cp
+from prime_rl.trainer.models.layers.ulysses_attn import (
+    _all_to_all_head_to_seq,
+    _all_to_all_seq_to_head,
+)
+from prime_rl.utils.sequence import CPPartition
+
+
+@pytest.fixture(autouse=True)
+def restore_attention(monkeypatch):
+    from prime_rl.trainer.models.layers.attn import FlashAttention
+
+    monkeypatch.setattr(FlashAttention, "_compute_attention", FlashAttention._compute_attention)
+
+
+@pytest.fixture(scope="module")
+def cp_group():
+    if int(os.environ.get("WORLD_SIZE", "1")) != 4:
+        pytest.skip("requires four distributed GPU ranks")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    torch._dynamo.config.recompile_limit = 64
+    dist.init_process_group("nccl", timeout=datetime.timedelta(seconds=180))
+    yield dist.group.WORLD
+    dist.destroy_process_group()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("compiled", [False, True])
+def test_cp_exchange_and_gather_gradients(cp_group, compiled):
+    rank, degree = cp_group.rank(), cp_group.size()
+
+    def exchange(x, shape):
+        total = shape.shape[0]
+        heads = _all_to_all_seq_to_head(x, degree, cp_group, total)
+        restored = _all_to_all_head_to_seq(heads, degree, cp_group, total)
+        return restored, all_gather_cp(x, 0, total, cp_group)
+
+    torch._dynamo.reset()
+    torch._dynamo.utils.counters.clear()
+    call = torch.compile(exchange, fullgraph=True, dynamic=True) if compiled else exchange
+    for total in [8, 9, 10, 11, 12, 13, 9, 8, 1, 2, 3, 0]:
+        print(f"collectives rank={rank} compiled={compiled} total={total}", flush=True)
+        partition = CPPartition(total, degree)
+        full = torch.arange(total * 8 * 3, device="cuda", dtype=torch.float64).reshape(total, 8, 3)
+        local = partition.shard(full, rank, 0).detach().requires_grad_()
+        shape = torch.empty(total, 0, device="cuda")
+        restored, gathered = call(local, shape)
+        torch.testing.assert_close(restored, local, rtol=0, atol=0)
+        torch.testing.assert_close(gathered, full, rtol=0, atol=0)
+        upstream = torch.arange(total, device="cuda", dtype=torch.float64).reshape(total, 1, 1) + 1
+        (restored.sum() * (rank + 1) + (gathered * upstream).sum() * (rank + 1)).backward()
+        expected = partition.shard(upstream, rank, 0) * (degree * (degree + 1) // 2) + rank + 1
+        torch.testing.assert_close(local.grad, expected.expand_as(local), rtol=0, atol=0)
+    print(
+        json.dumps(
+            {
+                "test": "collectives",
+                "rank": rank,
+                "compiled": compiled,
+                "graphs": dict(torch._dynamo.utils.counters["stats"]),
+                "breaks": dict(torch._dynamo.utils.counters["graph_break"]),
+            }
+        ),
+        flush=True,
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("compiled", [False, True])
+@pytest.mark.parametrize("style", ["ulysses", "ring"])
+@pytest.mark.parametrize("flash_attn_version", [2, 3, 4])
+def test_uneven_attention_matches_unsharded(cp_group, compiled, style, flash_attn_version):
+    if flash_attn_version == 4:
+        from flash_attn.cute import flash_attn_varlen_func as flash
+    elif flash_attn_version == 3:
+        from flash_attn_interface import flash_attn_varlen_func as flash
+    else:
+        from flash_attn import flash_attn_varlen_func as flash
+
+    rank, degree = cp_group.rank(), cp_group.size()
+
+    from prime_rl.trainer.models.layers.cp_attn import context_parallel_attention
+    from prime_rl.utils.cp import CPContext
+
+    context = CPContext(cp_group, rank, degree, style, True)
+
+    def attention(q, k, v, cu, shape, maximum):
+        return context_parallel_attention(flash, q, k, v, cu, maximum, shape.shape[0], context, flash_attn_version)
+
+    torch._dynamo.reset()
+    torch._dynamo.utils.counters.clear()
+    call = torch.compile(attention, fullgraph=flash_attn_version != 4) if compiled else attention
+    errors = torch.zeros(4, device="cuda")
+    for total in [8, 9, 10, 11, 12, 13, 9, 8, 1, 2, 3]:
+        print(f"attention rank={rank} compiled={compiled} total={total}", flush=True)
+        torch.manual_seed(42 + total)
+        partition = CPPartition(total, degree)
+        full = [
+            torch.randn(total, heads, 64, device="cuda", dtype=torch.bfloat16).requires_grad_() for heads in [8, 2, 2]
+        ]
+        boundaries = [0, 5, total] if total > 5 else [0, total]
+        maximum = max(b - a for a, b in zip(boundaries, boundaries[1:]))
+        cu = torch.tensor(boundaries, device="cuda", dtype=torch.int32)
+        if flash_attn_version == 4:
+            reference, _ = flash(
+                *full, cu_seqlens_q=cu, cu_seqlens_k=cu, max_seqlen_q=maximum, max_seqlen_k=maximum, causal=True
+            )
+        else:
+            reference = flash(*full, cu, cu, maximum, maximum, causal=True)
+        reference.float().square().sum().backward()
+        local = [partition.shard(x, rank, 0).detach().clone().requires_grad_() for x in full]
+        actual = call(*local, cu, torch.empty(total, 0, device="cuda"), maximum)
+        actual.float().square().sum().backward()
+        torch.testing.assert_close(actual, partition.shard(reference, rank, 0), rtol=0.03, atol=0.03)
+        pairs = [(actual.detach(), partition.shard(reference.detach(), rank, 0))]
+        pairs.extend((x.grad, partition.shard(ref.grad, rank, 0)) for x, ref in zip(local, full))
+        for index, (observed, expected) in enumerate(pairs):
+            torch.testing.assert_close(observed, expected, rtol=0.05, atol=0.05)
+            if observed.numel():
+                errors[index] = torch.maximum(errors[index], (observed.float() - expected.float()).abs().max())
+    dist.all_reduce(errors, op=dist.ReduceOp.MAX)
+    print(
+        json.dumps(
+            {
+                "test": "attention",
+                "version": flash_attn_version,
+                "rank": rank,
+                "compiled": compiled,
+                "max_abs_errors": errors.tolist(),
+                "graphs": dict(torch._dynamo.utils.counters["stats"]),
+                "breaks": dict(torch._dynamo.utils.counters["graph_break"]),
+            }
+        ),
+        flush=True,
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("compiled", [False, True])
+@pytest.mark.parametrize("objective", ["rl", "sft"])
+def test_qwen3_cp_model_and_checkpoint_gradients(cp_group, compiled, objective):
+    from transformers import Qwen3Config
+
+    from prime_rl.configs.trainer import ActivationCheckpointConfig, CompileConfig, ModelConfig
+    from prime_rl.trainer.model import apply_ac, apply_compile
+    from prime_rl.trainer.models.layers.attn import FlashAttention
+    from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head
+    from prime_rl.trainer.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM
+    from prime_rl.trainer.parallel_dims import ParallelDims
+    from prime_rl.utils.cp import setup_context_parallel
+
+    rank, degree = cp_group.rank(), cp_group.size()
+    torch.manual_seed(7)
+    config = Qwen3Config(
+        vocab_size=64,
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        max_position_embeddings=128,
+        attn_implementation="flash_attention_3",
+        use_cache=False,
+    )
+    model = Qwen3ForCausalLM(config).cuda().to(torch.bfloat16)
+    inject_prime_lm_head(model, chunk_size=8)
+    reference = copy.deepcopy(model)
+    original_attention = FlashAttention._compute_attention
+    for module in reference.modules():
+        if isinstance(module, FlashAttention):
+            module._compute_attention = original_attention.__get__(module)
+    dims = ParallelDims(dp_replicate=1, dp_shard=1, cp=degree, pp=1, ep=1, world_size=degree)
+    setup_context_parallel(
+        model,
+        ModelConfig(cp=degree, cp_style="ulysses", cp_unpadded=True, impl="custom", attn="flash_attention_3"),
+        dims,
+    )
+    apply_ac(model, ActivationCheckpointConfig(mode="full"))
+    if compiled:
+        apply_compile(model, CompileConfig(fullgraph=True))
+    for total in [9, 10, 1, 3, 12]:
+        print(f"model rank={rank} compiled={compiled} total={total}", flush=True)
+        model.zero_grad(set_to_none=True)
+        reference.zero_grad(set_to_none=True)
+        partition = CPPartition(total, degree)
+        ids = (torch.arange(total, device="cuda").reshape(1, total) + 3) % 64
+        lengths = [5, total - 5] if total > 5 else [total]
+        seq_lens = torch.tensor(lengths, device="cuda")
+        positions = torch.cat([torch.arange(n, device="cuda") for n in lengths]).unsqueeze(0)
+        labels = ids.roll(-1, 1)
+        temperature = torch.ones_like(ids, dtype=torch.float32)
+        ref = reference(ids, position_ids=positions, labels=labels, temperature=temperature, seq_lens=seq_lens)
+        if objective == "sft":
+            mask = positions.remainder(3) != 0
+            reference_loss = -ref["logprobs"][mask].sum()
+            labels = labels.masked_fill(~mask, -100)
+            reference_loss.backward()
+        else:
+            ref["logprobs"].sum().backward()
+        actual = model(
+            partition.shard(ids, rank),
+            position_ids=partition.shard(positions, rank),
+            labels=partition.shard(labels, rank),
+            temperature=partition.shard(temperature, rank) if objective == "rl" else None,
+            seq_lens=seq_lens,
+            seq_lens_are_pre_shard=True,
+            cp_total_tokens=total,
+        )
+        if objective == "sft":
+            global_loss = actual["loss"].detach().clone()
+            dist.all_reduce(global_loss)
+            torch.testing.assert_close(global_loss, reference_loss, rtol=0.03, atol=0.03)
+            actual["loss"].backward()
+        else:
+            gathered = all_gather_cp(actual["logprobs"], 1, total, cp_group)
+            torch.testing.assert_close(gathered, ref["logprobs"], rtol=0.03, atol=0.03)
+            (gathered.sum() / degree).backward()
+        max_error = 0.0
+        for (name, param), (ref_name, ref_param) in zip(model.named_parameters(), reference.named_parameters()):
+            assert name.replace("_checkpoint_wrapped_module.", "") == ref_name
+            assert param.grad is not None, name
+            dist.all_reduce(param.grad, group=cp_group)
+            torch.testing.assert_close(param.grad, ref_param.grad, rtol=0.05, atol=0.05, msg=name)
+            max_error = max(max_error, float((param.grad.float() - ref_param.grad.float()).abs().max()))
+        print(
+            json.dumps(
+                {"test": "model", "rank": rank, "compiled": compiled, "total": total, "max_param_grad_error": max_error}
+            ),
+            flush=True,
+        )
+    FlashAttention._compute_attention = original_attention
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "moe,compiled,ep",
+    [(False, False, 1), (False, True, 1), (True, False, 1), (True, True, 1), (True, False, 4), (True, True, 4)],
+)
+def test_fsdp_inactive_lane_preserves_global_gradient(cp_group, moe, compiled, ep):
+    from transformers import Qwen3Config
+
+    from prime_rl.configs.trainer import ActivationCheckpointConfig, CompileConfig, ModelConfig
+    from prime_rl.trainer.model import apply_ac, apply_compile, apply_fp32_moe_router, get_global_moe_stats, setup_fsdp
+    from prime_rl.trainer.models.layers.attn import FlashAttention
+    from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head
+    from prime_rl.trainer.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM
+    from prime_rl.trainer.moe_runtime import configure_moe_runtime
+    from prime_rl.trainer.parallel_dims import ParallelDims
+    from prime_rl.utils.cp import setup_context_parallel
+
+    torch.manual_seed(17)
+    config = Qwen3Config(
+        vocab_size=64,
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        max_position_embeddings=128,
+        attn_implementation="flash_attention_3",
+        use_cache=False,
+    )
+    if moe:
+        from prime_rl.trainer.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
+        from prime_rl.trainer.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeForCausalLM
+
+        config = Qwen3MoeConfig(
+            **(config.to_dict() | {"num_experts": 4, "num_experts_per_tok": 2, "moe_intermediate_size": 256})
+        )
+        model = Qwen3MoeForCausalLM(config).cuda()
+        for layer in model.model.layers:
+            layer.mlp.experts.init_weights(config.initializer_range)
+    else:
+        model = Qwen3ForCausalLM(config).cuda()
+    inject_prime_lm_head(model, chunk_size=8)
+    reference = copy.deepcopy(model).to(torch.bfloat16)
+    if moe:
+        apply_fp32_moe_router(model)
+        apply_fp32_moe_router(reference)
+    original_attention = FlashAttention._compute_attention
+    for module in reference.modules():
+        if isinstance(module, FlashAttention):
+            module._compute_attention = original_attention.__get__(module)
+    degree = ep
+    dims = ParallelDims(dp_replicate=1, dp_shard=2, cp=2, pp=1, ep=degree, world_size=4)
+    model_config = ModelConfig(
+        cp=2,
+        cp_style="ulysses",
+        cp_unpadded=True,
+        ep=degree,
+        impl="custom",
+        attn="flash_attention_3",
+        optim_cpu_offload=False,
+    )
+    configure_moe_runtime(model, model_config, dims)
+    setup_context_parallel(model, model_config, dims)
+    apply_ac(model, ActivationCheckpointConfig(mode="full"))
+    if compiled:
+        apply_compile(model, CompileConfig(fullgraph=not moe))
+    setup_fsdp(model, model_config, dims)
+    group = dims.get_mesh("cp").get_group()
+    cp_rank = group.rank()
+    lane = dist.get_rank() // 2
+
+    for lengths in json.loads(os.environ.get("CP_TEST_LENGTHS", "[[9,5],[9,0],[0,7],[3,1]]")):
+        print(f"fsdp rank={dist.get_rank()} lengths={lengths} moe={moe} compiled={compiled}", flush=True)
+        model.zero_grad(set_to_none=True)
+        reference.zero_grad(set_to_none=True)
+        denominator = sum(lengths)
+        for total in lengths:
+            if total == 0:
+                continue
+            ids = (torch.arange(total, device="cuda").reshape(1, total) + 3) % 64
+            ref = reference(
+                ids,
+                labels=ids.roll(-1, 1),
+                temperature=torch.ones_like(ids, dtype=torch.float32),
+                seq_lens=torch.tensor([total], device="cuda"),
+            )
+            (ref["logprobs"].sum() / denominator).backward()
+
+        total = lengths[lane]
+        partition = CPPartition(total, 2)
+        ids = (torch.arange(total, device="cuda").reshape(1, total) + 3) % 64
+        labels = ids.roll(-1, 1)
+        positions = torch.arange(total, device="cuda").reshape(1, total)
+        seq_lens = torch.tensor([total] if total else [], device="cuda", dtype=torch.int64)
+        actual = model(
+            partition.shard(ids, cp_rank),
+            position_ids=partition.shard(positions, cp_rank),
+            labels=partition.shard(labels, cp_rank),
+            temperature=torch.ones(1, partition.lengths[cp_rank], device="cuda"),
+            seq_lens=seq_lens,
+            seq_lens_are_pre_shard=True,
+        )
+        assert torch.isfinite(actual["logprobs"]).all()
+        gathered = all_gather_cp(actual["logprobs"], 1, total, group)
+        (gathered.sum() / (2 * denominator)).backward()
+        if moe:
+            metrics = get_global_moe_stats(
+                model,
+                dims.get_mesh("ep").get_group() if ep > 1 else None,
+                dims.get_mesh("dp_cp").get_group(),
+            )
+            assert metrics
+            assert all(torch.isfinite(value) for value in metrics.values())
+        max_error = 0.0
+        for (name, param), (_, ref_param) in zip(model.named_parameters(), reference.named_parameters()):
+            assert param.grad is not None, name
+            gradient = param.grad.full_tensor() * dims.fsdp_gradient_divide_factor
+            torch.testing.assert_close(
+                gradient,
+                ref_param.grad.float(),
+                rtol=0.05,
+                atol=0.025 if moe else 0.015,
+                msg=lambda msg: f"{name}: {msg}",
+            )
+            max_error = max(max_error, float((gradient - ref_param.grad.float()).abs().max()))
+        print(
+            json.dumps(
+                {
+                    "test": "fsdp-inactive",
+                    "moe": moe,
+                    "compiled": compiled,
+                    "rank": dist.get_rank(),
+                    "lengths": lengths,
+                    "max_param_grad_error": max_error,
+                }
+            ),
+            flush=True,
+        )
+    FlashAttention._compute_attention = original_attention
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("kind", ["mamba", "delta", "delta_gather"])
+def test_uneven_recurrence_matches_unsharded(cp_group, kind):
+    from prime_rl.utils.cp import CPContext
+
+    if kind == "mamba":
+        from prime_rl.trainer.models.nemotron_h import NemotronHConfig
+        from prime_rl.trainer.models.nemotron_h.mamba import NemotronHMamba2
+
+        config = NemotronHConfig(
+            hidden_size=256,
+            hybrid_override_pattern="M",
+            mamba_num_heads=8,
+            mamba_head_dim=64,
+            n_groups=4,
+            ssm_state_size=64,
+            conv_kernel=4,
+            chunk_size=64,
+        )
+        model = NemotronHMamba2(config)
+    else:
+        from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
+        from prime_rl.trainer.models.qwen3_5.gated_delta_net import Qwen3_5GatedDeltaNet
+
+        config = Qwen3_5TextConfig(
+            hidden_size=256,
+            linear_key_head_dim=32,
+            linear_value_head_dim=32,
+            linear_num_key_heads=1 if kind == "delta_gather" else 4,
+            linear_num_value_heads=2 if kind == "delta_gather" else 8,
+            linear_conv_kernel_dim=4,
+        )
+        model = Qwen3_5GatedDeltaNet(config)
+    model = model.cuda().to(torch.bfloat16)
+    for parameter in model.parameters():
+        dist.broadcast(parameter.data, 0)
+    reference = copy.deepcopy(model)
+    rank, degree = cp_group.rank(), cp_group.size()
+    model.cp_context = CPContext(cp_group, rank, degree, "ulysses", True)
+    for total in [9, 10, 1, 3, 0, 12]:
+        model.zero_grad(set_to_none=True)
+        reference.zero_grad(set_to_none=True)
+        torch.manual_seed(42 + total)
+        partition = CPPartition(total, degree)
+        full = torch.randn(1, total, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        local = partition.shard(full, rank).detach().clone().requires_grad_()
+        boundaries = [0, 5, total] if total > 5 else ([0, total] if total else [0])
+        cu = torch.tensor(boundaries, device="cuda", dtype=torch.int32)
+        expected = reference(full, cu)
+        actual = model(local, cu, cp_total_tokens=total)
+        expected.float().square().sum().backward()
+        actual.float().square().sum().backward()
+        torch.testing.assert_close(actual, partition.shard(expected, rank), rtol=0.03, atol=0.03)
+        torch.testing.assert_close(local.grad, partition.shard(full.grad, rank), rtol=0.05, atol=0.05)
+        for (name, parameter), (_, ref_parameter) in zip(model.named_parameters(), reference.named_parameters()):
+            assert parameter.grad is not None, name
+            dist.all_reduce(parameter.grad)
+            torch.testing.assert_close(
+                parameter.grad, ref_parameter.grad, rtol=0.05, atol=0.25, msg=lambda msg: f"{name}: {msg}"
+            )
+        print(json.dumps({"test": "recurrence", "kind": kind, "rank": rank, "total": total}), flush=True)
+
+
+@pytest.mark.gpu
+def test_vlm_images_and_inactive_lanes_with_ep(cp_group):
+    from prime_rl.configs.shared import VLMConfig
+    from prime_rl.configs.trainer import ActivationCheckpointConfig, CompileConfig, ModelConfig
+    from prime_rl.trainer.model import apply_ac, apply_compile, apply_fp32_moe_router, setup_fsdp
+    from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL
+    from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head
+    from prime_rl.trainer.moe_runtime import configure_moe_runtime
+    from prime_rl.trainer.parallel_dims import ParallelDims
+    from prime_rl.utils.cp import gather_for_cp, setup_context_parallel
+    from tests.unit.train.models.qwen3_5.test_vlm import get_image_inputs, get_vlm_config
+
+    torch.manual_seed(7)
+    config = get_vlm_config()
+    model = AutoModelForCausalLMPrimeRL.from_config(config, attn_implementation="flash_attention_3").cuda()
+    inject_prime_lm_head(model, chunk_size=8)
+    reference = copy.deepcopy(model).to(torch.bfloat16)
+    apply_fp32_moe_router(model)
+    apply_fp32_moe_router(reference)
+    dims = ParallelDims(dp_replicate=1, dp_shard=2, cp=2, pp=1, ep=4, world_size=4)
+    runtime = ModelConfig(
+        impl="custom",
+        attn="flash_attention_3",
+        cp=2,
+        cp_style="ulysses",
+        cp_unpadded=True,
+        inactive_micro_batches=True,
+        ep=4,
+        optim_cpu_offload=False,
+        vlm=VLMConfig(
+            freeze_vision_encoder=False, vision_encoder_attr="model.visual", language_model_attr="model.language_model"
+        ),
+    )
+    configure_moe_runtime(model, runtime, dims)
+    setup_context_parallel(model, runtime, dims)
+    apply_ac(model, ActivationCheckpointConfig(mode="full"))
+    apply_compile(model, CompileConfig(fullgraph=False))
+    setup_fsdp(model, runtime, dims)
+    group = dims.get_mesh("cp").get_group()
+    rank = group.rank()
+    lane = dist.get_rank() // 2
+
+    def inputs(total):
+        ids = (torch.arange(total, device="cuda").reshape(1, total) + 3) % 64
+        extras = {}
+        if total:
+            ids[0, 1] = config.image_token_id
+            torch.manual_seed(17 + total)
+            pixels, grid, _ = get_image_inputs(config)
+            extras = dict(
+                pixel_values=pixels, image_grid_thw=grid, mm_token_type_ids=(ids == config.image_token_id).long()
+            )
+        return ids, extras
+
+    for totals in [[9, 0], [0, 7], [3, 5]]:
+        model.zero_grad(set_to_none=True)
+        reference.zero_grad(set_to_none=True)
+        for total in totals:
+            if not total:
+                continue
+            ids, extras = inputs(total)
+            out = reference(
+                ids,
+                labels=ids.roll(-1, 1),
+                temperature=torch.ones_like(ids, dtype=torch.float32),
+                seq_lens=torch.tensor([total], device="cuda"),
+                **extras,
+            )
+            (out["logprobs"].sum() / sum(totals)).backward()
+        total = totals[lane]
+        ids, extras = inputs(total)
+        partition = CPPartition(total, 2)
+        out = model(
+            ids,
+            labels=partition.shard(ids.roll(-1, 1), rank),
+            temperature=torch.ones(1, partition.lengths[rank], device="cuda"),
+            seq_lens=torch.tensor([total] if total else [], device="cuda", dtype=torch.int64),
+            seq_lens_are_pre_shard=True,
+            cp_total_tokens=total,
+            **extras,
+        )
+        gathered = gather_for_cp(out["logprobs"], group, total)
+        (gathered.sum() / (2 * sum(totals))).backward()
+        maximum = 0.0
+        for (name, p), rp in zip(model.named_parameters(), reference.parameters()):
+            assert p.grad is not None, name
+            g = p.grad.full_tensor() * dims.fsdp_gradient_divide_factor
+            torch.testing.assert_close(g, rp.grad.float(), atol=0.02, rtol=0.06, msg=lambda msg: f"{name}: {msg}")
+            maximum = max(maximum, float((g - rp.grad.float()).abs().max()))
+        print(
+            json.dumps({"test": "vlm-inactive", "totals": totals, "rank": dist.get_rank(), "max_grad_error": maximum}),
+            flush=True,
+        )
diff --git a/tests/unit/utils/test_sequence.py b/tests/unit/utils/test_sequence.py
index 03af4418d0d..326dba2c723 100644
--- a/tests/unit/utils/test_sequence.py
+++ b/tests/unit/utils/test_sequence.py
@@ -2,11 +2,28 @@
 import torch
 
 from prime_rl.utils.sequence import (
+    CPPartition,
     get_cu_seqlens_from_position_ids,
     get_cu_seqlens_from_seq_lens,
 )
 
 
+@pytest.mark.parametrize("degree", [1, 2, 4])
+@pytest.mark.parametrize("total", [0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 32769])
+def test_cp_partition_roundtrip(total, degree):
+    partition = CPPartition(total, degree)
+    assert len(partition.lengths) == degree
+    assert partition.offsets[0] == 0
+    assert partition.offsets[-1] == total
+    assert sum(partition.lengths) == total
+    assert max(partition.lengths) - min(partition.lengths) <= 1
+    for dim, shape in [(1, (1, total, 2, 3)), (2, (3, 1, total))]:
+        tensor = torch.arange(torch.tensor(shape).prod()).reshape(shape)
+        shards = [partition.shard(tensor, rank, dim) for rank in range(degree)]
+        assert [shard.shape[dim] for shard in shards] == list(partition.lengths)
+        torch.testing.assert_close(torch.cat(shards, dim=dim), tensor)
+
+
 @pytest.mark.parametrize(
     ("position_ids", "expected_cu_seqlens", "expected_max_seqlen"),
     [