ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsDistributedBlockdiag Cp

nemo_automodel.components.distributed.blockdiag_cp.kernels

View as Markdown

Dense and varlen block-diagonal attention kernels.

Module Contents

Functions

NameDescription
_cp_blockdiag_maskPer-document causal attention mask for block-diagonal CP, shape [B, 1, L, S].
_cp_blockdiag_varlenBlock-diagonal CP attention via varlen (flash or TE) — no dense [B,1,L,S] mask.
_flash_varlen_with_long_prefix_guardRun FlashAttention without an asymmetric-prefix varlen launch.
_te_varlen_dpaCached TE DotProductAttention module for thd/varlen block-diagonal CP attention.
_varlen_backend_unavailable_reasonReturn a cheap, non-kernel-launching reason a varlen kernel cannot run.
_varlen_metadata_unavailable_reasonValidate every index consumed by a varlen CUDA kernel.
_varlen_seg_for_rankPer-rank block-diagonal varlen segmentation for ONE packed sequence.
precompute_blockdiag_varlen_metaPrecompute this rank’s varlen segmentation once per step.

Data

_CP_FLASH_DETERMINISTIC

_CP_FLASH_ENGAGED

_CP_FLASH_LONG_SEGMENT_WARNED

_CP_FLASH_WARNED

_CP_TE_DROPOUT_WARNED

_CP_VARLEN_SHAPE_LOGGED

_TE_DPA_CACHE

logger

API

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:

doc_ids
torch.Tensor

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

row_offset
int

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

local_len
int

L, the number of local query rows.

full_len
int

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

batch_size
int

B, used to expand a 1D doc_ids.

Returns: torch.Tensor

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

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:

query

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

key_full

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

value_full

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

doc_ids

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

row_offset

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

scale
Defaults to None

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

backend
Defaults to 'flash'

"flash" (flash_attn_varlen_func) or "te" (TransformerEngine DotProductAttention thd).

dropout_p
Defaults to 0.0

Attention dropout probability.

meta
Defaults to None

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

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:

q_packed
torch.Tensor

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

k_packed
torch.Tensor

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

v_packed
torch.Tensor

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

cu_q
torch.Tensor

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

cu_k
torch.Tensor

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

max_q
int

Maximum per-document query segment length.

max_k
int

Maximum per-document key segment length.

local_query_len
int

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

scale

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

dropout_p
float

Attention dropout probability.

meta
dict

Per-step metadata carrying first_q/first_k/max_tail.

Returns:

Packed attention output [n_real, Hq, D].

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.

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:

backend
str

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

dtype
torch.dtype

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

device
torch.device

Device the kernel would run on.

Returns: str | None

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

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:

meta
dict | None

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

query_len
int

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

key_len
int

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

device
torch.device

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.

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:

dids
torch.Tensor

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

row_offset
int

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

local_len
int

L, the local query chunk length.

dev

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

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:

doc_ids
torch.Tensor

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

row_offset
int

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

local_len
int

L, the local query chunk length.

device

Device for the produced cu_q/cu_k tensors.

Returns: dict

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

nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_DETERMINISTIC = False
 
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_ENGAGED = False
 
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_LONG_SEGMENT_WARNED = False
 
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_FLASH_WARNED = False
 
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_TE_DROPOUT_WARNED = False
 
nemo_automodel.components.distributed.blockdiag_cp.kernels._CP_VARLEN_SHAPE_LOGGED = False
 
nemo_automodel.components.distributed.blockdiag_cp.kernels._TE_DPA_CACHE = {}
 
nemo_automodel.components.distributed.blockdiag_cp.kernels.logger = logging.getLogger(__name__)