> 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.deepseek_v4.cp

Context-parallel helpers for the DeepSeek V4 custom model.

This implements the Miles-style training path: each CP rank owns a contiguous
query shard, while K/V and compressed K/V are all-gathered with autograd-aware
collectives before DSV4 sparse attention consumes them.

## Module Contents

### Functions

| Name                                                                                                                                          | Description                                                                   |
| --------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`_lcm`](#nemo_automodel-components-models-deepseek_v4-cp-_lcm)                                                                               | -                                                                             |
| [`_pad_1d`](#nemo_automodel-components-models-deepseek_v4-cp-_pad_1d)                                                                         | -                                                                             |
| [`_pad_length`](#nemo_automodel-components-models-deepseek_v4-cp-_pad_length)                                                                 | -                                                                             |
| [`_pad_position_ids_seq_dim_`](#nemo_automodel-components-models-deepseek_v4-cp-_pad_position_ids_seq_dim_)                                   | -                                                                             |
| [`_pad_tensor_seq_dim_`](#nemo_automodel-components-models-deepseek_v4-cp-_pad_tensor_seq_dim_)                                               | -                                                                             |
| [`_repad_dsv4_packed_batch`](#nemo_automodel-components-models-deepseek_v4-cp-_repad_dsv4_packed_batch)                                       | Insert DSV4 compression-safe padding into packed BSHD rows before CP slicing. |
| [`_valid_packed_lengths`](#nemo_automodel-components-models-deepseek_v4-cp-_valid_packed_lengths)                                             | -                                                                             |
| [`build_dsv4_cp_causal_padding_mask`](#nemo_automodel-components-models-deepseek_v4-cp-build_dsv4_cp_causal_padding_mask)                     | Build local-query/global-key additive mask for Miles-style DSV4 CP.           |
| [`build_dsv4_cp_packed_causal_padding_mask`](#nemo_automodel-components-models-deepseek_v4-cp-build_dsv4_cp_packed_causal_padding_mask)       | Build local-query/global-key additive mask for packed DSV4 CP.                |
| [`build_packed_seq_ids`](#nemo_automodel-components-models-deepseek_v4-cp-build_packed_seq_ids)                                               | Build per-token packed sequence IDs from padded packed lengths.               |
| [`dsv4_cp_all_gather`](#nemo_automodel-components-models-deepseek_v4-cp-dsv4_cp_all_gather)                                                   | All-gather activation tensors across CP ranks and concatenate on `dim`.       |
| [`dsv4_cp_all_gather_metadata`](#nemo_automodel-components-models-deepseek_v4-cp-dsv4_cp_all_gather_metadata)                                 | All-gather non-differentiable metadata such as padding masks.                 |
| [`dsv4_cp_enabled`](#nemo_automodel-components-models-deepseek_v4-cp-dsv4_cp_enabled)                                                         | Return whether a real CP process group is active.                             |
| [`dsv4_cp_local_seq_multiple`](#nemo_automodel-components-models-deepseek_v4-cp-dsv4_cp_local_seq_multiple)                                   | Required per-CP-rank sequence-length multiple for DSV4 Miles-style CP.        |
| [`dsv4_cp_rank`](#nemo_automodel-components-models-deepseek_v4-cp-dsv4_cp_rank)                                                               | Return this rank's index in the DSV4 CP group, or 0 without CP.               |
| [`dsv4_cp_size`](#nemo_automodel-components-models-deepseek_v4-cp-dsv4_cp_size)                                                               | Return the DSV4 CP group size, or 1 without CP.                               |
| [`make_dsv4_contiguous_shard_cp_batch_and_ctx`](#nemo_automodel-components-models-deepseek_v4-cp-make_dsv4_contiguous_shard_cp_batch_and_ctx) | Contiguously shard a batch for DeepSeek V4 Miles-style context parallelism.   |

### Data

[`_SEQ_LENS_PADDING_VALUE`](#nemo_automodel-components-models-deepseek_v4-cp-_SEQ_LENS_PADDING_VALUE)

### API

```python
nemo_automodel.components.models.deepseek_v4.cp._lcm(
    a: int,
    b: int
) -> int
```

```python
nemo_automodel.components.models.deepseek_v4.cp._pad_1d(
    values: list[int],
    width: int,
    padding_value: int = _SEQ_LENS_PADDING_VALUE
) -> torch.Tensor
```

```python
nemo_automodel.components.models.deepseek_v4.cp._pad_length(
    length: int,
    multiple: int
) -> int
```

```python
nemo_automodel.components.models.deepseek_v4.cp._pad_position_ids_seq_dim_(
    position_ids: torch.Tensor,
    seq_dim: int,
    pad_len: int
) -> torch.Tensor
```

```python
nemo_automodel.components.models.deepseek_v4.cp._pad_tensor_seq_dim_(
    tensor: torch.Tensor,
    seq_dim: int,
    pad_len: int,
    value
) -> torch.Tensor
```

```python
nemo_automodel.components.models.deepseek_v4.cp._repad_dsv4_packed_batch(
    batch: dict,
    cp_size: int,
    pad_multiple: int,
    padding_token_id: int,
    sync_packed_length: bool = False,
    loss_mask: torch.Tensor | None = None
) -> tuple[dict, torch.Tensor | None]
```

Insert DSV4 compression-safe padding into packed BSHD rows before CP slicing.

The generic packed dataset may pad each packed sequence only for TE CP. DSV4
compression additionally needs document boundaries to align to compressor
windows; CSA then uses `packed_seq_ids` to reset the previous-window overlap.
This routine rebuilds each row from real sequence spans, pads every span to
`pad_multiple`, and appends row-level pack padding with sequence ID 0. When
requested for HybridEP, the final physical length is max-reduced across ranks
before CP slicing so every rank in a flattened DP x CP expert group is uniform.

```python
nemo_automodel.components.models.deepseek_v4.cp._valid_packed_lengths(
    row: torch.Tensor,
    padding_value: int = _SEQ_LENS_PADDING_VALUE
) -> list[int]
```

```python
nemo_automodel.components.models.deepseek_v4.cp.build_dsv4_cp_causal_padding_mask(
    position_ids: torch.Tensor,
    key_len: int,
    dtype: torch.dtype,
    device: torch.device,
    cp_group,
    padding_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor
```

Build local-query/global-key additive mask for Miles-style DSV4 CP.

`position_ids` are the local query positions after contiguous CP slicing.
Keys are in global sequence order because DSV4 gathers K/V along sequence.
`padding_mask` follows the internal convention True=padding.

```python
nemo_automodel.components.models.deepseek_v4.cp.build_dsv4_cp_packed_causal_padding_mask(
    position_ids: torch.Tensor,
    packed_seq_ids: torch.Tensor,
    dtype: torch.dtype,
    device: torch.device,
    cp_group,
    padding_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor
```

Build local-query/global-key additive mask for packed DSV4 CP.

`packed_seq_ids` is local to the CP rank. The function all-gathers sequence
IDs and document-local positions for the key side, then applies same-document,
causal, padding, and optional sliding-window constraints.

```python
nemo_automodel.components.models.deepseek_v4.cp.build_packed_seq_ids(
    seq_lens_padded: torch.Tensor,
    seq_len: int,
    device: torch.device,
    padding_value: int = _SEQ_LENS_PADDING_VALUE
) -> torch.Tensor
```

Build per-token packed sequence IDs from padded packed lengths.

IDs are 1-based within each batch row; 0 marks trailing pack padding that
belongs to no sequence. `seq_lens_padded` may be right-padded with
`padding_value` to make rows rectangular.

```python
nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_all_gather(
    tensor: torch.Tensor,
    dim: int,
    cp_group
) -> torch.Tensor
```

All-gather activation tensors across CP ranks and concatenate on `dim`.

The distributed.nn functional collective preserves autograd, so backward
routes gradients for gathered remote slices back to their owning ranks.

```python
nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_all_gather_metadata(
    tensor: torch.Tensor | None,
    dim: int,
    cp_group
) -> torch.Tensor | None
```

All-gather non-differentiable metadata such as padding masks.

```python
nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_enabled(
    cp_group
) -> bool
```

Return whether a real CP process group is active.

```python
nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_local_seq_multiple(
    model_or_config
) -> int
```

Required per-CP-rank sequence-length multiple for DSV4 Miles-style CP.

Compress-ratio layers constrain how the sequence may be split across CP ranks:
a ratio-R layer needs each local shard divisible by R, and ratio-4 layers use
cross-window overlap so they need `2*R`. The returned value is the LCM across
all configured `compress_ratios` (1 when none are configured).

```python
nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_rank(
    cp_group
) -> int
```

Return this rank's index in the DSV4 CP group, or 0 without CP.

```python
nemo_automodel.components.models.deepseek_v4.cp.dsv4_cp_size(
    cp_group
) -> int
```

Return the DSV4 CP group size, or 1 without CP.

```python
nemo_automodel.components.models.deepseek_v4.cp.make_dsv4_contiguous_shard_cp_batch_and_ctx(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    pad_multiple: int | None = None,
    sync_packed_length: bool = False
)
```

Contiguously shard a batch for DeepSeek V4 Miles-style context parallelism.

Attached to the batch as `_cp_make_batch_fn` (via `functools.partial` to bind
`pad_multiple`) and invoked by `cp_utils.make_cp_batch_and_ctx`. HybridEP can
first max-reduce packed lengths so every rank contributes a uniform token count.
Each CP rank then keeps one `seq_start:seq_end` slice; DSV4 attention all-gathers
K/V across CP ranks during forward. Returns `(nullcontext, batch)`.

`pad_multiple` is the required *per-CP-rank* shard multiple (from
`dsv4_cp_local_seq_multiple`); the global sequence is padded so it is divisible
by `cp_size` and each local shard is divisible by `pad_multiple` (>= 2).
At CP size one, the native THD route only marks packed input as THD and leaves
its tensors and packing metadata unchanged.

```python
nemo_automodel.components.models.deepseek_v4.cp._SEQ_LENS_PADDING_VALUE = -1000
```