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:

: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.
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`	SDPA replacement for cp1 packed runs: block-diagonal varlen from doc_ids.
`attach_cp1_packed_varlen_hooks`	Scope the cp1 packed varlen SDPA patch to the model’s attention forwards.
`cp1_packed_varlen_backend`	The configured cp1 packed varlen backend (‘te’/‘flash’), or None if disabled (dense).
`disable_cp1_packed_varlen`	Disarm stale cp1 state before starting a new outer model forward.
`enable_cp1_packed_varlen`	Arm cp1 packed block-diagonal varlen for the rest of this step.

Data

_PACKED_STATE

API

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:

query

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

key

Keys [B, Hkv, S, D].

value

Values [B, Hkv, S, D].

attn_mask

Defaults to None

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

dropout_p

Defaults to 0.0

Dropout probability.

is_causal

Defaults to False

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

scale

Defaults to None

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

enable_gqa

Defaults to False

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

**kwargs

Defaults to {}

Ignored; accepted for SDPA signature compatibility.

Returns:

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

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:

model

torch.nn.Module

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

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

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.

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:

doc_ids

torch.Tensor

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

backend

str

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

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