nemo_automodel.components.distributed.blockdiag_cp.runtime

Block-diagonal SDPA routing and collective-safe fallback policy.

Module Contents

Functions

Name	Description
`_cp_group_spans_nodes`	True if the CP group spans more than one node.
`_needed_only_kernel_succeeded_on_all_ranks`	Make a rank-local varlen result safe to act on collectively.
`_needed_only_preflight`	Collectively decide whether every CP rank can run the varlen kernel.
`_resolve_cp_varlen_config`	Return `(attn_backend, kv_exchange)` for this step.
`_select_kv_exchange_path`	Decide this step’s KV-exchange path and WHY.
`cp_blockdiag_sdpa`	Block-diagonal context-parallel SDPA.

Data

_KV_EXCHANGE_PATH_LOGGED

_KV_XNODE_LOGGED

_ORIGINAL_SDPA

logger

API

nemo_automodel.components.distributed.blockdiag_cp.runtime._cp_group_spans_nodes(
    group
) -> bool

True if the CP group spans more than one node.

The needed-only exchanges (halo neighbor-p2p, a2a all-to-all-v) are verified when the CP group is node-local (cp_size <= GPUs/node, the common mesh layout where cp is the fast dim). Cross-node needed-only exchange is still fabric-sensitive, so the production default falls back to NCCL all-gather when the CP group spans nodes. Heuristic: cp_size > the local node’s world size (LOCAL_WORLD_SIZE / NPROC_PER_NODE / visible GPUs).

Set NEMO_CP_ALLOW_XNODE=1 to force the needed-only path cross-node for targeted validation on a specific cluster.

nemo_automodel.components.distributed.blockdiag_cp.runtime._needed_only_kernel_succeeded_on_all_ranks(
    out,
    group,
    device,
    diagnostic = None
) -> bool

Make a rank-local varlen result safe to act on collectively.

In particular, an all-padding rank returns a zero tensor without invoking FlashAttention/TE. It must nevertheless learn when a peer’s real-token kernel returned None before either rank advances to another collective.

nemo_automodel.components.distributed.blockdiag_cp.runtime._needed_only_preflight(
    state: dict,
    group,
    backend: str,
    query_dtype: torch.dtype,
    device: torch.device,
    query_len: int,
    key_len: int
) -> tuple[bool, str | None]

Collectively decide whether every CP rank can run the varlen kernel.

nemo_automodel.components.distributed.blockdiag_cp.runtime._resolve_cp_varlen_config(
    state: dict
) -> tuple[str, str]

Return (attn_backend, kv_exchange) for this step.

Prefers the per-step snapshot threaded into state by the batch/ctx builder (so :func:cp_blockdiag_sdpa is a function of (qkv, state) on the production path), and falls back to the runtime configuration owned by :mod:~nemo_automodel.components.distributed.blockdiag_cp.state for callers that set the step state manually (e.g. parity tests).

nemo_automodel.components.distributed.blockdiag_cp.runtime._select_kv_exchange_path(
    state,
    group,
    doc_ids,
    local_len,
    device,
    offset,
    query_dtype: torch.dtype | None = None,
    dropout_p: float = 0.0
)

Decide this step’s KV-exchange path and WHY.

Every downgrade to all-gather names its cause (mode, kernel, missing varlen meta, or cross-node topology) so a silent fall-through can’t hide a misconfiguration.

Parameters:

state

The per-step CP state dict (memoizes the plan and preflight).

group

The CP process group.

doc_ids

Replicated per-position document ids [B, S_full] (0 == pad).

local_len

Per-rank local sequence length L.

device

Device for plan tensors and the preflight all-reduce.

offset

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

query_dtype

torch.dtype | NoneDefaults to None

Query dtype for the kernel preflight (skip when None).

dropout_p

floatDefaults to 0.0

Attention dropout probability.

Returns:

(path, plan, reason) where path is "halo"/"a2a"/"allgather",

nemo_automodel.components.distributed.blockdiag_cp.runtime.cp_blockdiag_sdpa(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_mask: torch.Tensor | None = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: float | None = None,
    enable_gqa: bool = False,
    kwargs = {}
) -> torch.Tensor

Block-diagonal context-parallel SDPA.

Drop-in replacement for torch.nn.functional.scaled_dot_product_attention while the context returned by :func:~nemo_automodel.components.distributed.blockdiag_cp.batch.make_cp_blockdiag_batch_and_ctx is active; a plain pass-through to stock SDPA otherwise. K/V are exchanged across the CP group (all-gather, or a needed-only halo/a2a exchange) and one local attention runs the local queries against the delivered keys with a per-document causal mask. The passed attn_mask / is_causal are ignored on the CP path — masking is rebuilt from the document ids so packed sequences never attend across document boundaries.

Parameters:

query

torch.Tensor

This rank’s LOCAL query shard [B, Hq, L, D] (B = batch, Hq = query heads, L = local sequence length, D = head dim).

key

torch.Tensor

This rank’s LOCAL key shard [B, Hkv, L, D].

value

torch.Tensor

This rank’s LOCAL value shard [B, Hkv, L, D].

attn_mask

torch.Tensor | NoneDefaults to None

Ignored on the CP path (forwarded to stock SDPA otherwise).

dropout_p

floatDefaults to 0.0

Dropout probability.

is_causal

boolDefaults to False

Ignored on the CP path (forwarded to stock SDPA otherwise).

scale

float | NoneDefaults to None

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

enable_gqa

boolDefaults to False

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

**kwargs

Defaults to {}

Ignored; accepted for SDPA signature compatibility.

Returns: torch.Tensor

Attention output [B, Hq, L, D] for this rank’s local rows.

nemo_automodel.components.distributed.blockdiag_cp.runtime._KV_EXCHANGE_PATH_LOGGED = False

nemo_automodel.components.distributed.blockdiag_cp.runtime._KV_XNODE_LOGGED = False

nemo_automodel.components.distributed.blockdiag_cp.runtime._ORIGINAL_SDPA = torch.nn.functional.scaled_dot_product_attention

nemo_automodel.components.distributed.blockdiag_cp.runtime.logger = logging.getLogger(__name__)