> 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

Block-diagonal (per-document) varlen context parallelism for packed sequences.

Packed-sequence training concatenates many documents into one long sequence; a
correct attention mask is block-causal per document. The stock DTensor
`context_parallel` path assumes a single causal document, so packed VLM/LLM
batches need a CP implementation that reshards by contiguous chunks and rebuilds
per-document masking on every rank.

This package provides that implementation, split by responsibility:

* :mod:`.state`   -- runtime knob normalization + activation-checkpoint-safe step state.
* :mod:`.kernels` -- dense-mask and varlen (FlashAttention / TransformerEngine) kernels.
* :mod:`.exchange` -- differentiable K/V collectives (all-gather, left-halo, all-to-all-v).
* :mod:`.runtime` -- the SDPA entry point and collective-safe path selection.
* :mod:`.batch`   -- batch padding/sharding and the per-step train context.
* :mod:`.packed`  -- the cp\_size==1 packed-sequence varlen SDPA hook.

Integration follows the model-owned CP convention of
:func:`nemo_automodel.components.distributed.cp_utils.make_cp_batch_and_ctx`: a model
attaches :func:`make_cp_blockdiag_batch_and_ctx` to the batch as `_cp_make_batch_fn`
and routes its softmax attention through :func:`cp_blockdiag_sdpa` while the returned
context is active. For `cp_size == 1` packed runs, the model scopes the varlen SDPA
patch to its attention forwards with :func:`attach_cp1_packed_varlen_hooks` and arms
the per-forward state with :func:`enable_cp1_packed_varlen` /
:func:`disable_cp1_packed_varlen`.

Only these integration entry points are exported; everything else (knob
normalization, varlen metadata precompute, fire counters, kernels) is an internal
detail of the package's modules. Model wiring lands in follow-up PRs.

## Submodules

* **[`nemo_automodel.components.distributed.blockdiag_cp.batch`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/batch)**
* **[`nemo_automodel.components.distributed.blockdiag_cp.exchange`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/exchange)**
* **[`nemo_automodel.components.distributed.blockdiag_cp.kernels`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/kernels)**
* **[`nemo_automodel.components.distributed.blockdiag_cp.packed`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/packed)**
* **[`nemo_automodel.components.distributed.blockdiag_cp.runtime`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/runtime)**
* **[`nemo_automodel.components.distributed.blockdiag_cp.state`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/state)**

## Package Contents

### Data

[`__all__`](#nemo_automodel-components-distributed-blockdiag_cp-__all__)

### API

```python
nemo_automodel.components.distributed.blockdiag_cp.__all__ = ['attach_cp1_packed_varlen_hooks', 'configure_cp_varlen', 'cp_blockdiag_sdpa', '...
```