> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.gemma4_moe.cp_batch

Gemma4's aux-only contiguous-shard context-parallel batch prep.

Gemma4 runs its own p2p ring FlexAttention over contiguous per-rank sequence
slices (no collective -- the transport lives in Gemma4's attention, see
`cp_attention.py`). Under the sunk (Megatron-style per-microbatch) CP path the
model embeds, splices vision, builds `per_layer_inputs` and the flex-ring mask
metadata inside its forward and contiguously slices them there; the dispatch-time
sharder therefore only touches the no-grad auxiliary streams.

The generic slicing lives in `components/distributed/context_parallel/sharder.py`; this
module owns the one Gemma4-specific piece the aux-only shard still needs: the
`_packed_seq_ids` document-boundary synthesis its manual CP attention mask
builder requires (its pad-region zeros depend on the global pad tail, which the
forward -- holding only this rank's slice -- cannot reconstruct). Gemma4's
`prepare_model_inputs_for_cp` exposes it through the `ContextParallelSharder`
it returns under the `"cp_sharder"` batch key, which the CP dispatch invokes in
place of the default load-balanced `context_parallel` path.

## Module Contents

### Functions

| Name                                                                                                                                                       | Description                                                            |
| ---------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`_synthesize_single_document_seq_ids`](#nemo_automodel-components-models-gemma4_moe-cp_batch-_synthesize_single_document_seq_ids)                         | Materialize the trivial single-document `_packed_seq_ids` map.         |
| [`make_contiguous_aux_only_shard_cp_batch_and_ctx`](#nemo_automodel-components-models-gemma4_moe-cp_batch-make_contiguous_aux_only_shard_cp_batch_and_ctx) | Aux-only contiguous CP shard for Gemma4's sunk (in-forward) pre-embed. |

### API

```python
nemo_automodel.components.models.gemma4_moe.cp_batch._synthesize_single_document_seq_ids(
    batch: dict,
    seq_len: int
) -> None
```

Materialize the trivial single-document `_packed_seq_ids` map.

Collates emit `_packed_seq_ids` only when 2+ documents are packed, but
Gemma4's manual CP attention mask builder needs document boundaries even
for one document (1 = real token, 0 = pad). Derived from `padding_mask`
when present, else all-ones. A no-op when `_packed_seq_ids` already
exists (genuinely packed input).

**Parameters:**

**`batch`** `dict`

The CP batch dict; mutated in place to add `_packed_seq_ids`.

---

**`seq_len`** `int`

The pre-pad sequence length.

---

```python
nemo_automodel.components.models.gemma4_moe.cp_batch.make_contiguous_aux_only_shard_cp_batch_and_ctx(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    extra_seq_keys: dict[str, int] | None = None,
    extra_pad_values: dict[str, typing.Any] | None = None
)
```

Aux-only contiguous CP shard for Gemma4's sunk (in-forward) pre-embed.

Exposed as `ContextParallelSharder.shard_batch` by Gemma4's sharder-only
`prepare_model_inputs_for_cp`. It shards only the no-grad auxiliary streams
(`labels` / `position_ids` / `loss_mask` / `padding_mask` plus the
synthesized `_packed_seq_ids` document map) and leaves `input_ids` /
`pixel_values` / `mm_token_type_ids` FULL-length in the batch. The model
forward then embeds, splices vision, builds `per_layer_inputs` and the
`_gemma4_vision_group_ids` / `mm_token_type_ids` ring metadata on the full
sequence and contiguously slices them per microbatch (see
`shard_sequence_for_cp_contiguous`), so the embeddings and vision tower are
trainable under CP and the PP×CP shared pre-embed graph no longer exists.

`_packed_seq_ids` is synthesized full (from the full `padding_mask`) and
sharded here rather than in the forward: its pad-region zeros depend on the
global pad tail, which the forward — holding only this rank's slice — cannot
reconstruct. Every other flex-ring mask input the forward owns is a pure
per-token or cumsum-over-full-then-slice quantity, so it slices to the same
contiguous layout this sharder applies.