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

Block-diagonal SDPA routing and collective-safe fallback policy.

## Module Contents

### Functions

| Name                                                                                                                                                   | Description                                                          |
| ------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------- |
| [`_cp_group_spans_nodes`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_cp_group_spans_nodes)                                           | True if the CP group spans more than one node.                       |
| [`_needed_only_kernel_succeeded_on_all_ranks`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_needed_only_kernel_succeeded_on_all_ranks) | Make a rank-local varlen result safe to act on collectively.         |
| [`_needed_only_preflight`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_needed_only_preflight)                                         | Collectively decide whether every CP rank can run the varlen kernel. |
| [`_resolve_cp_varlen_config`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_resolve_cp_varlen_config)                                   | Return `(attn_backend, kv_exchange)` for this step.                  |
| [`_select_kv_exchange_path`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_select_kv_exchange_path)                                     | Decide this step's KV-exchange path and WHY.                         |
| [`cp_blockdiag_sdpa`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-cp_blockdiag_sdpa)                                                   | Block-diagonal context-parallel SDPA.                                |

### Data

[`_KV_EXCHANGE_PATH_LOGGED`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_KV_EXCHANGE_PATH_LOGGED)

[`_KV_XNODE_LOGGED`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_KV_XNODE_LOGGED)

[`_ORIGINAL_SDPA`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-_ORIGINAL_SDPA)

[`logger`](#nemo_automodel-components-distributed-blockdiag_cp-runtime-logger)

### API

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._cp_group_spans_nodes(
    group
) -> bool
```

True if the CP group spans more than one node.

The needed-only exchanges (halo neighbor-p2p, a2a all-to-all-v) are verified when
the CP group is node-local (cp\_size \<= GPUs/node, the common mesh layout where cp
is the fast dim). Cross-node needed-only exchange is still fabric-sensitive, so
the production default falls back to NCCL all-gather when the CP group spans nodes.
Heuristic: cp\_size > the local node's world size (LOCAL\_WORLD\_SIZE /
NPROC\_PER\_NODE / visible GPUs).

Set `NEMO_CP_ALLOW_XNODE=1` to force the needed-only path cross-node for targeted
validation on a specific cluster.

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._needed_only_kernel_succeeded_on_all_ranks(
    out,
    group,
    device,
    diagnostic = None
) -> bool
```

Make a rank-local varlen result safe to act on collectively.

In particular, an all-padding rank returns a zero tensor without invoking
FlashAttention/TE. It must nevertheless learn when a peer's real-token
kernel returned `None` before either rank advances to another collective.

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._needed_only_preflight(
    state: dict,
    group,
    backend: str,
    query_dtype: torch.dtype,
    device: torch.device,
    query_len: int,
    key_len: int
) -> tuple[bool, str | None]
```

Collectively decide whether every CP rank can run the varlen kernel.

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._resolve_cp_varlen_config(
    state: dict
) -> tuple[str, str]
```

Return `(attn_backend, kv_exchange)` for this step.

Prefers the per-step snapshot threaded into `state` by the batch/ctx builder (so
:func:`cp_blockdiag_sdpa` is a function of `(qkv, state)` on the production path),
and falls back to the runtime configuration owned by
:mod:`~nemo_automodel.components.distributed.blockdiag_cp.state` for callers that
set the step state manually (e.g. parity tests).

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._select_kv_exchange_path(
    state,
    group,
    doc_ids,
    local_len,
    device,
    offset,
    query_dtype: torch.dtype | None = None,
    dropout_p: float = 0.0
)
```

Decide this step's KV-exchange path and WHY.

Every downgrade to all-gather names its cause (mode, kernel, missing varlen meta,
or cross-node topology) so a silent fall-through can't hide a misconfiguration.

**Parameters:**

The per-step CP state dict (memoizes the plan and preflight).

The CP process group.

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

Per-rank local sequence length `L`.

Device for plan tensors and the preflight all-reduce.

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

Query dtype for the kernel preflight (skip when `None`).

Attention dropout probability.

**Returns:**

`(path, plan, reason)` where `path` is `"halo"`/`"a2a"`/`"allgather"`,

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime.cp_blockdiag_sdpa(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_mask: torch.Tensor | None = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: float | None = None,
    enable_gqa: bool = False,
    kwargs = {}
) -> torch.Tensor
```

Block-diagonal context-parallel SDPA.

Drop-in replacement for `torch.nn.functional.scaled_dot_product_attention`
while the context returned by
:func:`~nemo_automodel.components.distributed.blockdiag_cp.batch.make_cp_blockdiag_batch_and_ctx`
is active; a plain pass-through to stock SDPA otherwise. K/V are exchanged
across the CP group (all-gather, or a needed-only halo/a2a exchange) and one
local attention runs the local queries against the delivered keys with a
per-document causal mask. The passed `attn_mask` / `is_causal` are ignored
on the CP path -- masking is rebuilt from the document ids so packed sequences
never attend across document boundaries.

**Parameters:**

This rank's LOCAL query shard `[B, Hq, L, D]` (`B` = batch,
`Hq` = query heads, `L` = local sequence length, `D` = head dim).

This rank's LOCAL key shard `[B, Hkv, L, D]`.

This rank's LOCAL value shard `[B, Hkv, L, D]`.

Ignored on the CP path (forwarded to stock SDPA otherwise).

Dropout probability.

Ignored on the CP path (forwarded to stock SDPA otherwise).

Softmax scale (`None` -> `D**-0.5`).

Grouped-query attention flag as passed by HF's sdpa path.

Ignored; accepted for SDPA signature compatibility.

**Returns:** `torch.Tensor`

Attention output `[B, Hq, L, D]` for this rank's local rows.

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._KV_EXCHANGE_PATH_LOGGED = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._KV_XNODE_LOGGED = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime._ORIGINAL_SDPA = torch.nn.functional.scaled_dot_product_attention
```

```python
nemo_automodel.components.distributed.blockdiag_cp.runtime.logger = logging.getLogger(__name__)
```