nemo_automodel.components.attention.fa4

FlashAttention 4 (flash_attn.cute) helpers for models that support BackendConfig.attn = "fa4".

import_fa4_modules imports the optional FA4/CuTe modules with one clear error when they are missing.
causal_fa4_attention runs plain causal attention.
document_causal_fa4_attention runs causal, per-document attention of a contiguous query shard against the global keys (packed sequences, and context parallelism with contiguous shards). It is expressed as FA4 varlen segments, so no mask is materialized or evaluated per element and unequal QK/V head dims (e.g. MLA’s 192/128) run without padding:
- Every document run of a row (padding runs included) is split at the end of the local query shard.
- A piece that overlaps the shard contributes its overlapping queries and all of its keys; FA4’s causal mask is aligned to the bottom-right corner of each segment, which is exactly “query at global position p sees keys of its document up to p”.
- Pieces without local queries (other ranks’ later tokens, earlier documents) become zero-query segments, so the key segments tile the flattened [batch * key_sequence] tensor and no gather is needed.

Optional FA4/CuTe dependencies are imported on first use.

Module Contents

Functions

Name	Description
`_load_fa4`	Import the FA4 dense and varlen entry points once.
`_varlen_layout`	Build (and cache for the step) FA4 varlen `cu_seqlens` and max lengths.
`_varlen_segments`	Return per-segment query and key lengths describing the document-causal pattern.
`causal_fa4_attention`	Run plain causal attention with FA4 when queries and keys cover the same single-document sequence.
`document_causal_fa4_attention`	Run causal, per-document attention of local queries against global keys with FA4 varlen.
`import_fa4_modules`	Import optional FlashAttention 4 / CuTe DSL modules, raising one clear error if any is missing.

Data

_LAYOUT_CACHE

_LAYOUT_GENERATION

_PAD_DOC_ID

API

nemo_automodel.components.attention.fa4._load_fa4() -> tuple[collections.abc.Callable, collections.abc.Callable]

Import the FA4 dense and varlen entry points once.

nemo_automodel.components.attention.fa4._varlen_layout(
    q_doc_ids: torch.Tensor,
    kv_doc_ids: torch.Tensor,
    q_global_start: int
) -> tuple[torch.Tensor, torch.Tensor, int, int]

Build (and cache for the step) FA4 varlen cu_seqlens and max lengths.

Parameters:

q_doc_ids

torch.Tensor

Integer tensor [batch, query_sequence] of 1-based document ids (0 = padding) for the local queries; equal to kv_doc_ids[:, q_global_start : q_global_start + query_sequence].

kv_doc_ids

torch.Tensor

Integer tensor [batch, key_sequence] of document ids for all keys, on the attention device. Its storage identifies the step: a new tensor invalidates the cache.

q_global_start

int

Global sequence offset of the first local query.

Returns: torch.Tensor

(cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k): two int32 tensors of shape [segments + 1] on

nemo_automodel.components.attention.fa4._varlen_segments(
    kv_doc_ids: torch.Tensor,
    q_length: int,
    q_global_start: int
) -> tuple[list[int], list[int]]

Return per-segment query and key lengths describing the document-causal pattern.

Parameters:

kv_doc_ids

torch.Tensor

CPU tensor of shape [batch, key_sequence] with 1-based document ids, 0 for padding.

q_length

int

Local query length.

q_global_start

int

Global sequence offset of the first local query.

Returns: tuple[list[int], list[int]]

Query lengths and key lengths of every segment, in flattened row-major order.

nemo_automodel.components.attention.fa4.causal_fa4_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: float
) -> torch.Tensor

Run plain causal attention with FA4 when queries and keys cover the same single-document sequence.

Parameters:

query

torch.Tensor

Tensor of shape [batch, heads, sequence, qk_head_dim].

key

torch.Tensor

Tensor of shape [batch, heads, sequence, qk_head_dim].

value

torch.Tensor

Tensor of shape [batch, heads, sequence, v_head_dim].

scale

float

Softmax scale applied to the query-key product.

Returns: torch.Tensor

Tensor of shape [batch, heads, sequence, v_head_dim].

nemo_automodel.components.attention.fa4.document_causal_fa4_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    q_doc_ids: torch.Tensor,
    kv_doc_ids: torch.Tensor,
    q_global_start: int,
    scale: float
) -> torch.Tensor

Run causal, per-document attention of local queries against global keys with FA4 varlen.

Parameters:

query

torch.Tensor

Tensor of shape [batch, heads, query_sequence, qk_head_dim].

key

torch.Tensor

Tensor of shape [batch, heads, key_sequence, qk_head_dim].

value

torch.Tensor

Tensor of shape [batch, heads, key_sequence, v_head_dim].

q_doc_ids

torch.Tensor

Tensor of shape [batch, query_sequence] with 1-based document ids.

kv_doc_ids

torch.Tensor

Tensor of shape [batch, key_sequence] with 1-based document ids.

q_global_start

int

Global sequence offset of the first query token.

scale

float

Softmax scale applied to the query-key product.

Returns: torch.Tensor

Tensor of shape [batch, heads, query_sequence, v_head_dim]; padding queries are zero.

nemo_automodel.components.attention.fa4.import_fa4_modules(
    names: str = ()
) -> tuple[types.ModuleType, ...]

Import optional FlashAttention 4 / CuTe DSL modules, raising one clear error if any is missing.

Parameters:

*names

strDefaults to ()

Fully qualified module names, e.g. "flash_attn.cute.interface" or "cutlass.cute".

Returns: tuple[ModuleType, ...]

The imported modules, in the order requested.

nemo_automodel.components.attention.fa4._LAYOUT_CACHE: dict[tuple, Any] = {}

nemo_automodel.components.attention.fa4._LAYOUT_GENERATION: list[Any] = [None, None]

nemo_automodel.components.attention.fa4._PAD_DOC_ID = 0