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

cp\_size==1 packed-sequence varlen SDPA integration.

Same varlen kernel as the CP block-diagonal path, but for `cp_size == 1` (no
sequence sharding). The whole packed sequence lives on one rank, so it
degenerates to `_cp_blockdiag_varlen` with `row_offset=0` (q == k == full
sequence; square; `cu_q == cu_k`). This gives packed-sequence block-diagonal
attention to models whose softmax attention only dispatches to sdpa.

Integration mirrors :func:`nemo_automodel.components.distributed.cp_utils.attach_cp_sdpa_hooks`:

1. :func:`attach_cp1_packed_varlen_hooks` registers forward pre/post hooks on every
   `self_attn` module (the checkpoint-wrapped INNER module, so the hooks also fire
   during activation-checkpointing recompute). The pre-hook swaps
   `F.scaled_dot_product_attention` for :func:`_packed_varlen_sdpa`; the post-hook
   (`always_call=True`) restores stock SDPA. The patch is therefore live only
   while a hooked attention forward is running -- never process-wide.
2. The model arms the per-forward state with :func:`enable_cp1_packed_varlen`
   (doc\_ids + backend) at the start of its forward and clears stale state before
   the NEXT forward with :func:`disable_cp1_packed_varlen`.

IMPORTANT -- state must survive activation-checkpointing RECOMPUTE in backward:
the state is set per-forward and NOT reset at the end of the step, so the AC
worker thread re-running a layer's forward during backward reads the SAME doc\_ids
and reproduces the exact varlen output. (A per-step reset caused the forward to
use varlen but the recompute to fall back to dense -> AC "saved vs recomputed
metadata" shape mismatch.) `_ThreadSharedVar` makes the state visible in the
autograd worker thread. Because the outer model clears stale state before each
new forward, vision/unpacked attention cannot inherit it -- and the shape check
in :func:`_packed_varlen_sdpa` passes any non-matching call straight through.

## Module Contents

### Functions

| Name                                                                                                                          | Description                                                                           |
| ----------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`_packed_varlen_sdpa`](#nemo_automodel-components-distributed-blockdiag_cp-packed-_packed_varlen_sdpa)                       | SDPA replacement for cp1 packed runs: block-diagonal varlen from doc\_ids.            |
| [`attach_cp1_packed_varlen_hooks`](#nemo_automodel-components-distributed-blockdiag_cp-packed-attach_cp1_packed_varlen_hooks) | Scope the cp1 packed varlen SDPA patch to the model's attention forwards.             |
| [`cp1_packed_varlen_backend`](#nemo_automodel-components-distributed-blockdiag_cp-packed-cp1_packed_varlen_backend)           | The configured cp1 packed varlen backend ('te'/'flash'), or None if disabled (dense). |
| [`disable_cp1_packed_varlen`](#nemo_automodel-components-distributed-blockdiag_cp-packed-disable_cp1_packed_varlen)           | Disarm stale cp1 state before starting a new outer model forward.                     |
| [`enable_cp1_packed_varlen`](#nemo_automodel-components-distributed-blockdiag_cp-packed-enable_cp1_packed_varlen)             | Arm cp1 packed block-diagonal varlen for the rest of this step.                       |

### Data

[`_PACKED_STATE`](#nemo_automodel-components-distributed-blockdiag_cp-packed-_PACKED_STATE)

### API

```python
nemo_automodel.components.distributed.blockdiag_cp.packed._packed_varlen_sdpa(
    query,
    key,
    value,
    attn_mask = None,
    dropout_p = 0.0,
    is_causal = False,
    scale = None,
    enable_gqa = False,
    kwargs = {}
)
```

SDPA replacement for cp1 packed runs: block-diagonal varlen from doc\_ids.

Only intervenes for a full packed-sequence call whose batch and sequence
dims match the armed `doc_ids`; any other SDPA call (or state unset)
passes straight through to stock SDPA.

**Parameters:**

Queries `[B, Hq, S, D]` over the full (unsharded) packed
sequence (`B` = batch, `Hq` = query heads, `S` = sequence
length, `D` = head dim).

Keys `[B, Hkv, S, D]`.

Values `[B, Hkv, S, D]`.

Forwarded to stock SDPA on pass-through; ignored on the
varlen path (masking is rebuilt from `doc_ids`).

Dropout probability.

Forwarded on pass-through; the varlen path is always
per-document causal.

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

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

Ignored; accepted for SDPA signature compatibility.

**Returns:**

Attention output `[B, Hq, S, D]`.

```python
nemo_automodel.components.distributed.blockdiag_cp.packed.attach_cp1_packed_varlen_hooks(
    model: torch.nn.Module
) -> None
```

Scope the cp1 packed varlen SDPA patch to the model's attention forwards.

Registers a forward pre-hook / post-hook pair on every `self_attn` module
that installs :func:`_packed_varlen_sdpa` as `F.scaled_dot_product_attention`
for the duration of that module's forward and restores stock SDPA afterwards
(`always_call=True`, so a raising forward cannot leak the patch). Hooks are
attached to the checkpoint-wrapped INNER module because CheckpointWrapper's
recompute bypasses `__call__` on the wrapper -- this is what keeps the
varlen path active during activation-checkpointing recompute in backward.

Same bounded-patch pattern as
:func:`nemo_automodel.components.distributed.cp_utils.attach_cp_sdpa_hooks`.
Outside these hooks, `F.scaled_dot_product_attention` is untouched.

**Parameters:**

The model whose `self_attn` submodules route softmax attention
through `F.scaled_dot_product_attention`.

```python
nemo_automodel.components.distributed.blockdiag_cp.packed.cp1_packed_varlen_backend() -> str | None
```

The configured cp1 packed varlen backend ('te'/'flash'), or None if disabled (dense).

```python
nemo_automodel.components.distributed.blockdiag_cp.packed.disable_cp1_packed_varlen() -> None
```

Disarm stale cp1 state before starting a new outer model forward.

Decoder-layer activation-checkpoint recomputation happens before the next
outer forward, so the preceding step's state remains available for backward
and is cleared before vision or any other attention in the next batch runs.

```python
nemo_automodel.components.distributed.blockdiag_cp.packed.enable_cp1_packed_varlen(
    doc_ids: torch.Tensor,
    backend: str
) -> None
```

Arm cp1 packed block-diagonal varlen for the rest of this step.

Sets the per-forward doc\_ids/backend state read by the SDPA patch that
:func:`attach_cp1_packed_varlen_hooks` scopes to the attention forwards. The
state remains armed through backward's activation-checkpoint recomputation;
the next outer model forward clears it via :func:`disable_cp1_packed_varlen`.

**Parameters:**

Per-position document ids `[1, S]` or `[S]` (0 == padding)
over the full packed sequence.

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

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