> 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.distributed.blockdiag_cp.batch

Batch padding, sequential sharding, and block-diagonal CP context setup.

## Module Contents

### Functions

| Name                                                                                                                           | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------- |
| [`_cp_blockdiag_doc_ids`](#nemo_automodel-components-distributed-blockdiag_cp-batch-_cp_blockdiag_doc_ids)                     | Resolve per-position document ids `[B, S]` (0 == padding) for the mask. |
| [`make_cp_blockdiag_batch_and_ctx`](#nemo_automodel-components-distributed-blockdiag_cp-batch-make_cp_blockdiag_batch_and_ctx) | Sequentially shard a pre-embedded batch for block-diagonal CP.          |

### API

```python
nemo_automodel.components.distributed.blockdiag_cp.batch._cp_blockdiag_doc_ids(
    batch: dict,
    seq_len: int,
    device,
    batch_size: int
) -> torch.Tensor
```

Resolve per-position document ids `[B, S]` (0 == padding) for the mask.

Prefers the collator's `_packed_seq_ids` (1-based document index per token,
present when a pack holds >1 document). Otherwise falls back to the 4-D
block-causal `attention_mask` diagonal (valid positions) or, lacking both,
treats the whole sequence as a single document.

**Parameters:**

The training batch; may contain `_packed_seq_ids` `[B, S]`
(int document index per token) or `attention_mask` (`[B, 1, S, S]`
block-causal bool, or `[B, S]` validity/indexed mask).

`S`, the (unpadded) sequence length.

Device for the returned tensor.

`B`, used for the all-ones fallback.

**Returns:** `torch.Tensor`

Per-position document ids `[B, S]` (int64, 0 == padding).

```python
nemo_automodel.components.distributed.blockdiag_cp.batch.make_cp_blockdiag_batch_and_ctx(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    batch: dict[str, typing.Any],
    loss_mask: torch.Tensor | None = None,
    padding_token_id: int = 0
) -> tuple[typing.Callable[[], typing.ContextManager], dict[str, typing.Any]]
```

Sequentially shard a pre-embedded batch for block-diagonal CP.

Pads the sequence to a multiple of the CP world size, slices each
sequence-aligned tensor to this rank's contiguous chunk (a differentiable
slice for `inputs_embeds`, so gradients flow back to the trainable vision
tower / embedding table), and returns `(train_ctx, batch)` where entering
`train_ctx` activates the per-document CP SDPA state for the step.

Softmax attention must route through
:func:`~nemo_automodel.components.distributed.blockdiag_cp.runtime.cp_blockdiag_sdpa`
while this context is active. A model opts in by attaching this callable to
the batch as `_cp_make_batch_fn` (the model-owned CP hook honored by
:func:`nemo_automodel.components.distributed.cp_utils.make_cp_batch_and_ctx`)
and rebinding its attention's SDPA call for the step.

**Parameters:**

The context-parallel device (sub)mesh.

Accepted for `_cp_make_batch_fn` signature compatibility;
unused (block-diagonal CP shards only the sequence dimension).

The training batch. Must contain pre-embedded `inputs_embeds`
`[B, S, H]` (multimodal token replacement happens pre-shard) and is
mutated in place: `attention_mask` is dropped, `padding_mask`
`[B, S]` (bool, True == pad) is added, and every sequence-aligned
tensor is padded then sliced to this rank's `[row_offset,
row_offset + S_full/cp)` chunk.

Optional per-token loss mask `[B, S]`; padded with 0 and
sharded like the other sequence-aligned tensors (stored back into
`batch["loss_mask"]`).

Accepted for signature compatibility; unused (the
batch is pre-embedded, so there are no token ids to pad).

**Returns:** `Callable[[], ContextManager]`

`(train_ctx, batch)`: a zero-arg callable returning the per-step