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

Runtime knobs and activation-checkpoint-safe state for block-diagonal CP.

The knob normalization (synonym maps and defaults) lives here, in a
dependency-free leaf, so every consumer (the kernel-side
:func:`configure_cp_varlen` entry point and any policy-side config parser)
derives the accepted values from the same table and can never disagree on
synonyms or defaults.

## Module Contents

### Classes

| Name                                                                                             | Description                                                                   |
| ------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------- |
| [`_ThreadSharedVar`](#nemo_automodel-components-distributed-blockdiag_cp-state-_ThreadSharedVar) | A `contextvars.ContextVar`-style get/set/reset holder visible ACROSS threads. |

### Functions

| Name                                                                                                             | Description                                                                |
| ---------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_clean`](#nemo_automodel-components-distributed-blockdiag_cp-state-_clean)                                     | -                                                                          |
| [`configure_cp_varlen`](#nemo_automodel-components-distributed-blockdiag_cp-state-configure_cp_varlen)           | Configure the block-diagonal CP attention path from parsed runtime config. |
| [`cp_attn_fire_count`](#nemo_automodel-components-distributed-blockdiag_cp-state-cp_attn_fire_count)             | Number of block-diagonal CP attention calls since the last reset.          |
| [`cp_varlen_runtime_config`](#nemo_automodel-components-distributed-blockdiag_cp-state-cp_varlen_runtime_config) | Return the currently configured block-diagonal CP runtime settings.        |
| [`normalize_attn_backend`](#nemo_automodel-components-distributed-blockdiag_cp-state-normalize_attn_backend)     | Canonicalize an attention-backend knob to one of `ATTN_BACKEND_VALUES`.    |
| [`normalize_kv_exchange`](#nemo_automodel-components-distributed-blockdiag_cp-state-normalize_kv_exchange)       | Canonicalize a KV-exchange knob to one of `KV_EXCHANGE_VALUES`.            |
| [`reset_cp_attn_fire_count`](#nemo_automodel-components-distributed-blockdiag_cp-state-reset_cp_attn_fire_count) | Zero the CP-attention fire counter (call before each forward).             |

### Data

[`ATTN_BACKEND_DEFAULT`](#nemo_automodel-components-distributed-blockdiag_cp-state-ATTN_BACKEND_DEFAULT)

[`ATTN_BACKEND_SYNONYMS`](#nemo_automodel-components-distributed-blockdiag_cp-state-ATTN_BACKEND_SYNONYMS)

[`ATTN_BACKEND_VALUES`](#nemo_automodel-components-distributed-blockdiag_cp-state-ATTN_BACKEND_VALUES)

[`KV_EXCHANGE_DEFAULT`](#nemo_automodel-components-distributed-blockdiag_cp-state-KV_EXCHANGE_DEFAULT)

[`KV_EXCHANGE_SYNONYMS`](#nemo_automodel-components-distributed-blockdiag_cp-state-KV_EXCHANGE_SYNONYMS)

[`KV_EXCHANGE_VALUES`](#nemo_automodel-components-distributed-blockdiag_cp-state-KV_EXCHANGE_VALUES)

[`_ATTN_BACKEND_REVERSE`](#nemo_automodel-components-distributed-blockdiag_cp-state-_ATTN_BACKEND_REVERSE)

[`_CP_ATTN_BACKEND`](#nemo_automodel-components-distributed-blockdiag_cp-state-_CP_ATTN_BACKEND)

[`_CP_ATTN_FIRE_COUNT`](#nemo_automodel-components-distributed-blockdiag_cp-state-_CP_ATTN_FIRE_COUNT)

[`_CP_BLOCKDIAG_STATE`](#nemo_automodel-components-distributed-blockdiag_cp-state-_CP_BLOCKDIAG_STATE)

[`_CP_KV_EXCHANGE`](#nemo_automodel-components-distributed-blockdiag_cp-state-_CP_KV_EXCHANGE)

[`_KV_EXCHANGE_REVERSE`](#nemo_automodel-components-distributed-blockdiag_cp-state-_KV_EXCHANGE_REVERSE)

### API

```python
class nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar()
```

A `contextvars.ContextVar`-style get/set/reset holder visible ACROSS threads.

The block-diagonal CP state must be readable inside the autograd worker thread
that runs activation-checkpointing recompute during backward. A real
`ContextVar` is per-thread and reads its default (None) there, which would
silently drop the CP state mid-recompute -- the softmax SDPA would fall back to
local-only attention. Training steps are sequential, so a single shared slot
with token-based restore (supporting nesting) is safe; backward only ever READS
the slot.

```python
nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar.get()
```

Return the current value (`None` when unset).

```python
nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar.reset(
    token
)
```

Restore the value captured by a previous :meth:`set`.

```python
nemo_automodel.components.distributed.blockdiag_cp.state._ThreadSharedVar.set(
    value
)
```

Set the value; returns a token (the previous value) for :meth:`reset`.

```python
nemo_automodel.components.distributed.blockdiag_cp.state._clean(
    value: typing.Any
) -> str
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state.configure_cp_varlen(
    attn_backend: str = 'flash',
    kv_exchange: str = 'allgather'
) -> None
```

Configure the block-diagonal CP attention path from parsed runtime config.

**Parameters:**

Varlen kernel selection; any synonym accepted by
:func:`normalize_attn_backend` (default `"flash"`).

K/V delivery mode; any synonym accepted by
:func:`normalize_kv_exchange` (default `"allgather"`).

```python
nemo_automodel.components.distributed.blockdiag_cp.state.cp_attn_fire_count() -> int
```

Number of block-diagonal CP attention calls since the last reset.

```python
nemo_automodel.components.distributed.blockdiag_cp.state.cp_varlen_runtime_config() -> dict[str, str]
```

Return the currently configured block-diagonal CP runtime settings.

**Returns:** `dict[str, str]`

A dict with keys `"attn_backend"` and `"kv_exchange"` holding the

```python
nemo_automodel.components.distributed.blockdiag_cp.state.normalize_attn_backend(
    value: typing.Any
) -> str
```

Canonicalize an attention-backend knob to one of `ATTN_BACKEND_VALUES`.

`None`/`""`/`auto`/`default` map to the default (`flash`).

**Parameters:**

The raw user-facing knob value (string-like or `None`).

**Returns:** `str`

The canonical backend name (`"flash"`, `"te"`, or `"dense"`).

**Raises:**

* `ValueError`: If the value is not a recognized backend or synonym.

```python
nemo_automodel.components.distributed.blockdiag_cp.state.normalize_kv_exchange(
    value: typing.Any
) -> str
```

Canonicalize a KV-exchange knob to one of `KV_EXCHANGE_VALUES`.

`None`/`""`/`auto`/`default` map to the default (`allgather`).

**Parameters:**

The raw user-facing knob value (string-like or `None`).

**Returns:** `str`

The canonical exchange name (`"allgather"`, `"halo"`, or `"a2a"`).

**Raises:**

* `ValueError`: If the value is not a recognized exchange mode or synonym.

```python
nemo_automodel.components.distributed.blockdiag_cp.state.reset_cp_attn_fire_count() -> None
```

Zero the CP-attention fire counter (call before each forward).

```python
nemo_automodel.components.distributed.blockdiag_cp.state.ATTN_BACKEND_DEFAULT = 'flash'
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state.ATTN_BACKEND_SYNONYMS: dict[str, tuple[str, ...]] = {'flash': ('flash', 'flash_attn', 'flash_attention', 'flash_attention_2', 'fa2')...
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state.ATTN_BACKEND_VALUES: tuple[str, ...] = tuple(ATTN_BACKEND_SYNONYMS)
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state.KV_EXCHANGE_DEFAULT = 'allgather'
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state.KV_EXCHANGE_SYNONYMS: dict[str, tuple[str, ...]] = {'allgather': ('allgather', 'all_gather'), 'halo': ('halo', 'neighbor', 'needed'...
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state.KV_EXCHANGE_VALUES: tuple[str, ...] = tuple(KV_EXCHANGE_SYNONYMS)
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state._ATTN_BACKEND_REVERSE = {syn: canon for canon, syns in (ATTN_BACKEND_SYNONYMS.items()) for syn in syns}
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state._CP_ATTN_BACKEND = ATTN_BACKEND_DEFAULT
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state._CP_ATTN_FIRE_COUNT: list[int] = [0]
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state._CP_BLOCKDIAG_STATE = _ThreadSharedVar()
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state._CP_KV_EXCHANGE = KV_EXCHANGE_DEFAULT
```

```python
nemo_automodel.components.distributed.blockdiag_cp.state._KV_EXCHANGE_REVERSE = {syn: canon for canon, syns in (KV_EXCHANGE_SYNONYMS.items()) for syn in syns}
```