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

Dense and varlen block-diagonal attention kernels.

## Module Contents

### Functions

| Name                                                                                                                                       | Description                                                                        |
| ------------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------- |
| [`_cp_blockdiag_mask`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_cp_blockdiag_mask)                                     | Per-document causal attention mask for block-diagonal CP, shape `[B, 1, L, S]`.    |
| [`_cp_blockdiag_varlen`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_cp_blockdiag_varlen)                                 | Block-diagonal CP attention via varlen (flash or TE) -- no dense `[B,1,L,S]` mask. |
| [`_flash_varlen_with_long_prefix_guard`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_flash_varlen_with_long_prefix_guard) | Run FlashAttention without an asymmetric-prefix varlen launch.                     |
| [`_te_varlen_dpa`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_te_varlen_dpa)                                             | Cached TE DotProductAttention module for thd/varlen block-diagonal CP attention.   |
| [`_varlen_backend_unavailable_reason`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_varlen_backend_unavailable_reason)     | Return a cheap, non-kernel-launching reason a varlen kernel cannot run.            |
| [`_varlen_metadata_unavailable_reason`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_varlen_metadata_unavailable_reason)   | Validate every index consumed by a varlen CUDA kernel.                             |
| [`_varlen_seg_for_rank`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_varlen_seg_for_rank)                                 | Per-rank block-diagonal varlen segmentation for ONE packed sequence.               |
| [`precompute_blockdiag_varlen_meta`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-precompute_blockdiag_varlen_meta)         | Precompute this rank's varlen segmentation once per step.                          |

### Data

[`_CP_FLASH_DETERMINISTIC`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_CP_FLASH_DETERMINISTIC)

[`_CP_FLASH_ENGAGED`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_CP_FLASH_ENGAGED)

[`_CP_FLASH_LONG_SEGMENT_WARNED`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_CP_FLASH_LONG_SEGMENT_WARNED)

[`_CP_FLASH_WARNED`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_CP_FLASH_WARNED)

[`_CP_TE_DROPOUT_WARNED`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_CP_TE_DROPOUT_WARNED)

[`_CP_VARLEN_SHAPE_LOGGED`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_CP_VARLEN_SHAPE_LOGGED)

[`_TE_DPA_CACHE`](#nemo_automodel-components-distributed-blockdiag_cp-kernels-_TE_DPA_CACHE)

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

### API

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._cp_blockdiag_mask(
    doc_ids: torch.Tensor,
    row_offset: int,
    local_len: int,
    full_len: int,
    batch_size: int
) -> torch.Tensor
```

Per-document causal attention mask for block-diagonal CP, shape `[B, 1, L, S]`.

`doc_ids` is the full (all-rank, padded) per-position document index `[B, S]`
(0 == padding). Query rows are this rank's local positions
`[row_offset, row_offset+local_len)`; key columns span the full sequence. A
query attends to a key iff they share a document, neither is padding, and the
key is causally visible (global query position >= key position) -- identical to
the block-causal mask a non-CP packed run would build for real tokens. The
diagonal is always allowed so a query row is never fully masked (which the
math/efficient SDPA backend turns into NaN); for padding rows this is a
harmless self-edge whose output is dropped by the -100 labels.

**Parameters:**

Per-position document ids `[B, S]` or `[S]` (0 == padding),
where `B` = batch and `S` = full padded sequence length.

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

`L`, the number of local query rows.

`S`, the number of key columns (full padded sequence).

`B`, used to expand a 1D `doc_ids`.

**Returns:** `torch.Tensor`

Boolean allow-mask `[B, 1, L, S]` (True == may attend).

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._cp_blockdiag_varlen(
    query,
    key_full,
    value_full,
    doc_ids,
    row_offset,
    scale = None,
    backend = 'flash',
    dropout_p = 0.0,
    meta = None
)
```

Block-diagonal CP attention via varlen (flash or TE) -- no dense `[B,1,L,S]` mask.

Equivalent to `_cp_blockdiag_mask` + SDPA for real (non-padding) query rows.
The sequence is sharded contiguously and packing puts padding (doc id 0) only
as a contiguous tail, so the local query's real rows are a prefix `[0, n_real)`.
For each document the local rows touch, we emit a varlen segment: q segment =
the local rows in that doc, k segment = `[doc_start, last_local_q+1)` (the only
doc with `sk > sq` straddles the left boundary). Bottom-right causal alignment
reproduces same-doc + global-causal exactly. Padding query rows are returned as
zeros (their loss is masked by -100 labels).

**Parameters:**

Local query shard `[B, Hq, L, D]` (`B` = batch, `Hq` = query
heads, `L` = local sequence length, `D` = head dim).

Keys `[B, Hkv, S, D]` covering the K/V range indexed by
`meta`'s `[s_first, real_end)` slice (full sequence for the
all-gather path, needed range for halo/A2A).

Values `[B, Hkv, S, D]`; same layout as `key_full`.

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

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

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

`"flash"` (flash\_attn\_varlen\_func) or `"te"`
(TransformerEngine DotProductAttention thd).

Attention dropout probability.

Optional per-step segmentation from
:func:`precompute_blockdiag_varlen_meta`; rebuilt inline when absent.

**Returns:**

Attention output `[B, Hq, L, D]`, or `None` to signal "fall back to

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._flash_varlen_with_long_prefix_guard(
    q_packed: torch.Tensor,
    k_packed: torch.Tensor,
    v_packed: torch.Tensor,
    cu_q: torch.Tensor,
    cu_k: torch.Tensor,
    max_q: int,
    max_k: int,
    local_query_len: int,
    scale,
    dropout_p: float,
    meta: dict
)
```

Run FlashAttention without an asymmetric-prefix varlen launch.

In CP, only the first local document can have `K > Q`: it may include a
left halo from the preceding rank. Some FlashAttention builds have produced
an asynchronous illegal access for long packs with this layout.
Fixed-sequence Flash accepts the same bottom-right causal Q/K layout without
consulting `cu_seqlens`. Peel off just that one asymmetric segment, then
process all remaining (ordinary `K == Q`) documents in one varlen call.

Normal packs stay on the original single-call path. A guarded pack adds at
most one kernel launch, independent of its number of packed documents.

**Parameters:**

Packed local queries `[n_real, Hq, D]` (`n_real` = real
local query tokens, `Hq` = query heads, `D` = head dim).

Packed keys `[T_k, Hkv, D]` covering the needed K/V range.

Packed values `[T_k, Hkv, D]`; same layout as `k_packed`.

Per-document cumulative query offsets, 1D int32 `[n_docs + 1]`.

Per-document cumulative key offsets, 1D int32 `[n_docs + 1]`.

Maximum per-document query segment length.

Maximum per-document key segment length.

This rank's local sequence length (for diagnostics).

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

Attention dropout probability.

Per-step metadata carrying `first_q`/`first_k`/`max_tail`.

**Returns:**

Packed attention output `[n_real, Hq, D]`.

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._te_varlen_dpa(
    num_q_heads,
    num_kv_heads,
    head_dim,
    scale,
    device,
    dtype
)
```

Cached TE DotProductAttention module for thd/varlen block-diagonal CP attention.

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._varlen_backend_unavailable_reason(
    backend: str,
    dtype: torch.dtype,
    device: torch.device
) -> str | None
```

Return a cheap, non-kernel-launching reason a varlen kernel cannot run.

Needed-only KV exchange cannot safely discover these failures after a
rank-local kernel call: an all-padding rank skips the call while a real-token
rank may fail and otherwise diverge in collective order. The runtime calls
this on every CP rank and reaches a small consensus before exchanging K/V.
Shape/data-dependent kernel failures are still handled by a post-call
consensus in :mod:`nemo_automodel.components.distributed.blockdiag_cp.runtime`.

**Parameters:**

Varlen backend name (`"flash"` or `"te"`).

Query/key/value dtype the kernel would run with.

Device the kernel would run on.

**Returns:** `str | None`

A human-readable reason string, or `None` when the kernel can run.

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._varlen_metadata_unavailable_reason(
    meta: dict | None,
    query_len: int,
    key_len: int,
    device: torch.device
) -> str | None
```

Validate every index consumed by a varlen CUDA kernel.

FlashAttention trusts `cu_seqlens` and the caller-provided maxima. A bad
terminal offset therefore does not reliably raise a Python exception; it can
become an asynchronous illegal memory access and poison the whole CUDA
context. Validate the compact, per-step metadata on the host before the
first attention layer launches a kernel. The result is cached in `meta`;
the cache object is intentionally shared by the shallow metadata copies used
by halo/A2A, so this costs one small device-to-host copy per step and shape,
not once per layer.

**Parameters:**

Per-step varlen metadata as produced by
:func:`precompute_blockdiag_varlen_meta` (`cu_q`/`cu_k` are 1D
int32 device tensors of per-document cumulative offsets).

Local query length `L` (rows this rank attends with).

Key length `S` visible to this rank (full sequence for
all-gather, needed range for halo/A2A).

Device the `cu_q`/`cu_k` tensors must live on.

**Returns:** `str | None`

A human-readable reason string when the metadata is unsafe, else `None`.

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._varlen_seg_for_rank(
    dids: torch.Tensor,
    row_offset: int,
    local_len: int,
    dev
) -> dict | None
```

Per-rank block-diagonal varlen segmentation for ONE packed sequence.

Factoring the segmentation out of the per-layer hot path lets
:func:`precompute_blockdiag_varlen_meta` run it once per step so the
attention layers do zero GPU->CPU host syncs.

**Parameters:**

Full (padded) per-position document id vector `[S]` (0 == pad).

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

`L`, the local query chunk length.

Device for the produced `cu_q`/`cu_k` tensors.

**Returns:** `dict | None`

The cu\_seqlens / slice metadata dict for this rank's local query chunk

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels.precompute_blockdiag_varlen_meta(
    doc_ids: torch.Tensor,
    row_offset: int,
    local_len: int,
    device
) -> dict
```

Precompute this rank's varlen segmentation once per step.

The block-diagonal varlen cu\_seqlens depend only on `(doc_ids, row_offset,
local_len)` -- all step-constant -- yet rebuilding them inline on every
attention layer (forward + AC recompute) costs 3-4 host syncs per rebuild.
:func:`~nemo_automodel.components.distributed.blockdiag_cp.batch.make_cp_blockdiag_batch_and_ctx`
calls this once and stashes the result in the CP state; `_cp_blockdiag_varlen(meta=...)`
then runs with zero `.item()` syncs.

Assumes one packed sequence per rank (B == 1, the CP contract).

**Parameters:**

Per-position document ids `[B, S]` (row 0 used) or `[S]`
(0 == padding), covering the full padded sequence.

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

`L`, the local query chunk length.

Device for the produced `cu_q`/`cu_k` tensors.

**Returns:** `dict`

A dict consumed directly by `_cp_blockdiag_varlen`; `{"n_real": 0}`

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_DETERMINISTIC = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_ENGAGED = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_LONG_SEGMENT_WARNED = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_WARNED = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_TE_DROPOUT_WARNED = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_VARLEN_SHAPE_LOGGED = False
```

```python
nemo_automodel.components.distributed.blockdiag_cp.kernels._TE_DPA_CACHE = {}
```

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