> 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.exchange

Differentiable K/V collectives and needed-only exchange plans.

## Module Contents

### Classes

| Name                                                                                                  | Description                                                          |
| ----------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [`_AllGatherSeqDiff`](#nemo_automodel-components-distributed-blockdiag_cp-exchange-_AllGatherSeqDiff) | All-gather a per-rank sequence shard along `seq_dim` over `group`.   |
| [`_LeftHaloExchange`](#nemo_automodel-components-distributed-blockdiag_cp-exchange-_LeftHaloExchange) | Uniform-size single-step neighbor exchange for node-local halo mode. |
| [`_NeededKVExchange`](#nemo_automodel-components-distributed-blockdiag_cp-exchange-_NeededKVExchange) | Differentiable multi-cast all-to-all-v K/V delivery.                 |

### Functions

| Name                                                                                                                    | Description                                                                    |
| ----------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`_blockdiag_a2a_attention`](#nemo_automodel-components-distributed-blockdiag_cp-exchange-_blockdiag_a2a_attention)     | Needed-only block-diagonal CP attention via the general all-to-all-v exchange. |
| [`_blockdiag_halo_attention`](#nemo_automodel-components-distributed-blockdiag_cp-exchange-_blockdiag_halo_attention)   | Needed-only block-diagonal CP attention via the left-halo exchange.            |
| [`_compute_blockdiag_kv_plan`](#nemo_automodel-components-distributed-blockdiag_cp-exchange-_compute_blockdiag_kv_plan) | Global left-halo exchange plan, computed from replicated `doc_ids`.            |
| [`_needed_kv_a2a_plan`](#nemo_automodel-components-distributed-blockdiag_cp-exchange-_needed_kv_a2a_plan)               | Per-(src,dst) split sizes + local gather index for the general needed-only     |

### API

```python
class nemo_automodel.components.distributed.blockdiag_cp.exchange._AllGatherSeqDiff()
```

**Bases:** `Function`

All-gather a per-rank sequence shard along `seq_dim` over `group`.

Forward concatenates every rank's shard `[B, H, L, D]` in rank order
(sequential sharding), producing the full sequence `[B, H, L*world, D]`.
Backward reduce-scatters the incoming gradient: each K/V position is read by
every rank's queries, so its gradient is the sum over ranks of the per-rank
grad slice; reduce-scatter(SUM) hands each rank the summed gradient for the
shard it owns. All shards have equal length `L` (the sequence is padded to
a multiple of the CP world size before sharding).

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._AllGatherSeqDiff.backward(
    ctx,
    grad_out: torch.Tensor
)
```

staticmethod

Reduce-scatter(SUM) the full-sequence gradient back to this rank's shard.

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._AllGatherSeqDiff.forward(
    ctx,
    x: torch.Tensor,
    group,
    seq_dim: int
) -> torch.Tensor
```

staticmethod

All-gather `x` (this rank's shard) into the full sequence along `seq_dim`.

```python
class nemo_automodel.components.distributed.blockdiag_cp.exchange._LeftHaloExchange()
```

**Bases:** `Function`

Uniform-size single-step neighbor exchange for node-local halo mode.

Every rank sends the last `halo_size` tokens of its chunk `[B, C, L, D]` to
`next_peer` = (rank+1)%world and receives `halo_size` tokens from
`prev_peer` = (rank-1)%world, send-before-recv, in one `batch_isend_irecv`.
The caller slices the last `recv_count` of the received block as its real
straddle (rank 0's wraparound block and any over-send are ignored).
Differentiable: backward is the reverse exchange (send grad to prev, recv from
next).

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._LeftHaloExchange.backward(
    ctx,
    grad_recv
)
```

staticmethod

Reverse exchange: route `grad_recv` `[B, C, halo_size, D]` back to the sending suffix.

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._LeftHaloExchange.forward(
    ctx,
    x_local,
    group,
    halo_size,
    prev_peer,
    next_peer
)
```

staticmethod

Exchange the chunk suffix; `x_local` is `[B, C, L, D]`, returns `[B, C, halo_size, D]`.

```python
class nemo_automodel.components.distributed.blockdiag_cp.exchange._NeededKVExchange()
```

**Bases:** `Function`

Differentiable multi-cast all-to-all-v K/V delivery.

Delivers each rank exactly the K/V range it attends to, zero-redundantly (a
source token is sent once per rank that needs it). Forward maps the local
chunk `[B, C, L, D]` to the received needed range `[B, C, R_recv, D]`.
Backward scatter-ADDs the per-destination grads back to the source token
(index\_add), matching what the all-gather's reduce\_scatter would sum.

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._NeededKVExchange.backward(
    ctx,
    grad_out
)
```

staticmethod

Reverse all-to-all-v of `grad_out` `[B, C, R_recv, D]`; multi-cast grads accumulate.

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._NeededKVExchange.forward(
    ctx,
    x_local,
    group,
    in_splits,
    out_splits,
    send_index
)
```

staticmethod

All-to-all-v `x_local` `[B, C, L, D]` into the needed range `[B, C, R_recv, D]`.

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._blockdiag_a2a_attention(
    query,
    key,
    value,
    doc_ids,
    group,
    plan,
    gm,
    row_offset,
    scale,
    backend,
    dropout_p
)
```

Needed-only block-diagonal CP attention via the general all-to-all-v exchange.

Handles documents spanning >2 ranks (single long doc / unpacked long context)
that the left-halo can't, still zero-redundantly (vs the full all-gather
fallback). Same argument contract as :func:`_blockdiag_halo_attention`.

**Parameters:**

Local query shard `[B, Hq, L, D]`.

Local key shard `[B, Hkv, L, D]`.

Local value shard `[B, Hkv, L, D]`.

Replicated per-position document ids `[B, S_full]` (0 == pad).

The CP process group.

Replicated plan from :func:`_compute_blockdiag_kv_plan`
(augmented with `rank`/`world`).

This rank's per-step varlen metadata (global `s_first`/`real_end`).

Global position of this rank's first local query row.

Softmax scale (`None` -> kernel default).

Varlen kernel backend, `"flash"` or `"te"`.

Attention dropout probability.

**Returns:**

Attention output `[B, Hq, L, D]`, or `None` when the kernel is

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._blockdiag_halo_attention(
    query,
    key,
    value,
    doc_ids,
    group,
    plan,
    gm,
    row_offset,
    scale,
    backend,
    dropout_p
)
```

Needed-only block-diagonal CP attention via the left-halo exchange.

Each rank attends its local queries against `[boundary-doc halo from rank-1] + [local chunk]` instead of the full all-gathered sequence. The caller must fail
fast when this returns `None`: the halo collective has already executed, so a
rank-local all-gather fallback would fork the CP group's collective order.

**Parameters:**

Local query shard `[B, Hq, L, D]`.

Local key shard `[B, Hkv, L, D]`.

Local value shard `[B, Hkv, L, D]`.

Replicated per-position document ids `[B, S_full]` (0 == pad).

The CP process group.

Replicated halo plan from :func:`_compute_blockdiag_kv_plan`
(augmented with `rank`/`world`).

This rank's per-step varlen metadata (global `s_first`/`real_end`).

Global position of this rank's first local query row.

Softmax scale (`None` -> kernel default).

Varlen kernel backend, `"flash"` or `"te"`.

Attention dropout probability.

**Returns:**

Attention output `[B, Hq, L, D]`, or `None` when the kernel is

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._compute_blockdiag_kv_plan(
    doc_ids_full: torch.Tensor,
    world: int,
    local_len: int,
    dev
) -> dict
```

Global left-halo exchange plan, computed from replicated `doc_ids`.

For each rank r the block-diagonal path needs the contiguous K/V range
`[s_first_r, real_end_r)` = its boundary document's straddle (owned by ranks \< r)
plus its own local chunk. `back_r = r*local_len - s_first_r` is the straddle
length. If `back_r &lt;= local_len` for ALL ranks, every straddle fits entirely in
the LEFT neighbor (rank r-1), so a single neighbor p2p suffices (`use_halo`). If
any document spans >2 ranks (`back_r &gt; local_len` -- e.g. one doc across the whole
sequence) the caller falls back to the general all-to-all-v. Every rank computes
this from the same replicated `doc_ids` and therefore agrees on `use_halo` and
on all send/recv counts without communication.

**Parameters:**

Replicated per-position document ids `[B, S_full]` (row 0
used) or `[S_full]` (0 == padding) on the full padded sequence.

CP world size.

Per-rank local sequence length `L` (`S_full == world * L`).

Device used for the intermediate segmentation tensors.

**Returns:** `dict`

A plan dict with per-rank lists `recv` (tokens received from r-1 ==

```python
nemo_automodel.components.distributed.blockdiag_cp.exchange._needed_kv_a2a_plan(
    plan: dict,
    rank: int,
    world: int,
    local_len: int,
    dev
)
```

Per-(src,dst) split sizes + local gather index for the general needed-only
all-to-all-v exchange (the >2-rank / single-doc case the halo can't cover).

Each rank r needs the contiguous global range `[s_first_r, real_end_r)`. As a
SENDER, this rank s sends to each dst r the intersection of its owned chunk
`[s*L,(s+1)*L)` with r's needed range (a token may go to several dsts, so the
send buffer duplicates it once per destination). As a receiver, the pieces arrive
in source order and concatenate into the contiguous needed range. All counts derive
from the replicated plan, so every rank agrees.

**Parameters:**

Replicated plan from :func:`_compute_blockdiag_kv_plan`.

This rank's index within the CP group.

CP world size.

Per-rank local sequence length `L`.

Device for the produced `send_index` tensor.

**Returns:**

`(in_splits, out_splits, send_index)`: per-destination send counts,