> This page is for version 0.7.0 · 26.10. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > 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. # 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`](#nemo_automodel-components-attention-fa4-_load_fa4) | Import the FA4 dense and varlen entry points once. | | [`_varlen_layout`](#nemo_automodel-components-attention-fa4-_varlen_layout) | Build (and cache for the step) FA4 varlen `cu_seqlens` and max lengths. | | [`_varlen_segments`](#nemo_automodel-components-attention-fa4-_varlen_segments) | Return per-segment query and key lengths describing the document-causal pattern. | | [`causal_fa4_attention`](#nemo_automodel-components-attention-fa4-causal_fa4_attention) | Run plain causal attention with FA4 when queries and keys cover the same single-document sequence. | | [`document_causal_fa4_attention`](#nemo_automodel-components-attention-fa4-document_causal_fa4_attention) | Run causal, per-document attention of local queries against global keys with FA4 varlen. | | [`import_fa4_modules`](#nemo_automodel-components-attention-fa4-import_fa4_modules) | Import optional FlashAttention 4 / CuTe DSL modules, raising one clear error if any is missing. | ### Data [`_LAYOUT_CACHE`](#nemo_automodel-components-attention-fa4-_LAYOUT_CACHE) [`_LAYOUT_GENERATION`](#nemo_automodel-components-attention-fa4-_LAYOUT_GENERATION) [`_PAD_DOC_ID`](#nemo_automodel-components-attention-fa4-_PAD_DOC_ID) ### API ```python nemo_automodel.components.attention.fa4._load_fa4() -> tuple[collections.abc.Callable, collections.abc.Callable] ``` Import the FA4 dense and varlen entry points once. ```python 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 ```python 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. ```python 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]. ```python 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. ```python 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`** `str` — default: () Fully qualified module names, e.g. `"flash_attn.cute.interface"` or `"cutlass.cute"`. --- **Returns:** `tuple[ModuleType, ...]` The imported modules, in the order requested. ```python nemo_automodel.components.attention.fa4._LAYOUT_CACHE: dict[tuple, Any] = {} ``` ```python nemo_automodel.components.attention.fa4._LAYOUT_GENERATION: list[Any] = [None, None] ``` ```python nemo_automodel.components.attention.fa4._PAD_DOC_ID = 0 ```