> This page is for version 0.5.0 · 26.06. > 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.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. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.models.step5.fa4 Step-5 attention through FA4, with compact region membership metadata. ## Module Contents ### Classes | Name | Description | | ------------------------------------------------------------------------------------------------ | ------------------------------------------------------------- | | [`_BackwardBlockPartition`](#nemo_automodel-components-models-step5-fa4-_BackwardBlockPartition) | Independent int32 lists in the original provider/head layout. | ### Functions | Name | Description | | -------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- | | [`_all_padding_attention`](#nemo_automodel-components-models-step5-fa4-_all_padding_attention) | Copy singleton-key attention while retaining explicit zero Q/K gradients. | | [`_block_lists`](#nemo_automodel-components-models-step5-fa4-_block_lists) | Build block lists without allocating a dense block grid. | | [`_causal_document_attention`](#nemo_automodel-components-models-step5-fa4-_causal_document_attention) | Run the existing B1 causal prefix or proven singleton-padding operation. | | [`_compact_selected_tiles`](#nemo_automodel-components-models-step5-fa4-_compact_selected_tiles) | Compact a boolean block grid into FA4 partial-block counts and indices. | | [`_load_fa4`](#nemo_automodel-components-models-step5-fa4-_load_fa4) | Load FA4 and cache a mask, optionally specializing a proven document. | | [`_load_fa4_varlen`](#nemo_automodel-components-models-step5-fa4-_load_fa4_varlen) | Load native FA4 causal varlen attention without a custom mask. | | [`_provider_heads`](#nemo_automodel-components-models-step5-fa4-_provider_heads) | Copy provider-specific block metadata into FA4's query-head axis. | | [`_region_tiles`](#nemo_automodel-components-models-step5-fa4-_region_tiles) | Union selected regions for bounded chunks of 128-query tiles. | | [`_selected_block_lists`](#nemo_automodel-components-models-step5-fa4-_selected_block_lists) | Create route-derived forward/reverse FA4 lists with bounded scratch space. | | [`_split_backward_full_blocks`](#nemo_automodel-components-models-step5-fa4-_split_backward_full_blocks) | Classify complete causal rectangles using SM90 D192 backward tiles. | | [`_split_forward_full_blocks`](#nemo_automodel-components-models-step5-fa4-_split_forward_full_blocks) | Separate a low contiguous full-block prefix without changing K order. | | [`fa4_attention`](#nemo_automodel-components-models-step5-fa4-fa4_attention) | Run FA4 forward/backward with compact CSA metadata on SM90 or SM100. | ### API ```python class nemo_automodel.components.models.step5.fa4._BackwardBlockPartition() ``` **Bases:** `NamedTuple` Independent int32 lists in the original provider/head layout. **`full_counts`** `Tensor` Tensor with the partial_counts layout. --- **`full_indices`** `Tensor` Tensor with the partial_indices layout. --- **`partial_counts`** `Tensor` Tensor \[provider_batches, heads, key_tiles]. --- **`partial_indices`** `Tensor` Tensor \[provider_batches, heads, key_tiles, query_tiles]. --- ```python nemo_automodel.components.models.step5.fa4._all_padding_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, padding_key: int ) -> torch.Tensor ``` Copy singleton-key attention while retaining explicit zero Q/K gradients. **Parameters:** **`query`** `torch.Tensor` Tensor \[batch, local_queries, query_heads, head_dim]. --- **`key`** `torch.Tensor` Tensor \[batch, global_keys, kv_heads, head_dim]. --- **`value`** `torch.Tensor` Tensor shaped like key, potentially a noncontiguous KV-unbind view. --- **`padding_key`** `int` The only allowed global key, proven by shared metadata. query_heads must be divisible by kv_heads. --- **Returns:** `torch.Tensor` Materialized tensor \[batch, local_queries, query_heads, head_dim]. ```python nemo_automodel.components.models.step5.fa4._block_lists( queries: int, keys: int, query_start: int, window: int | None, block_q: int, block_k: int, reverse: bool, device: torch.device ) -> tuple[torch.Tensor, torch.Tensor] ``` Build block lists without allocating a dense block grid. **Parameters:** **`queries`** `int` Local query length. --- **`keys`** `int` Global key length. --- **`query_start`** `int` Physical global offset of the query shard. --- **`window`** `int | None` Sliding window size, or None for all causal blocks. --- **`block_q`** `int` Query block size. --- **`block_k`** `int` Key block size. --- **`reverse`** `bool` Return key-to-query lists used by backward. --- **`device`** `torch.device` CUDA metadata device. --- **Returns:** `torch.Tensor` Int32 counts \[1, 1, block_rows] and indices \[1, 1, block_rows, width]. ```python nemo_automodel.components.models.step5.fa4._causal_document_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, layout: nemo_automodel.components.models.step5.sparse.CausalDocumentLayout ) -> torch.Tensor ``` Run the existing B1 causal prefix or proven singleton-padding operation. **Parameters:** **`query`** `torch.Tensor` BF16 tensor \[1, local_queries, query_heads, head_dim]. --- **`key`** `torch.Tensor` BF16 tensor \[1, global_keys, kv_heads, head_dim], possibly strided. --- **`value`** `torch.Tensor` Tensor with the same shape/dtype as key; global storage is shared. --- **`layout`** `CausalDocumentLayout` Proven prefix metadata for this one batch row. A zero query count denotes a fully padded row with a valid singleton padding key. --- **Returns:** `torch.Tensor` Tensor \[1, local_queries, query_heads, head_dim], preserving B1 arithmetic. ```python nemo_automodel.components.models.step5.fa4._compact_selected_tiles( present: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor] ``` Compact a boolean block grid into FA4 partial-block counts and indices. ```python nemo_automodel.components.models.step5.fa4._load_fa4( window: int | None, bitmap: bool, slots: int, provider_groups: int, heads_per_group: int = 0, canonical_region_size: int = 0 ) -> tuple[collections.abc.Callable, type, collections.abc.Callable] ``` Load FA4 and cache a mask, optionally specializing a proven document. canonical_region_size>0 requires one unpadded B1 CP1 document with Q=K, zero query_start and positions exactly 0..K-1. Its token_regions are physical_key // canonical_region_size; learned membership is unchanged. ```python nemo_automodel.components.models.step5.fa4._load_fa4_varlen() -> collections.abc.Callable ``` Load native FA4 causal varlen attention without a custom mask. ```python nemo_automodel.components.models.step5.fa4._provider_heads( tensor: torch.Tensor, batch: int, heads_per_group: int ) -> torch.Tensor ``` Copy provider-specific block metadata into FA4's query-head axis. **Parameters:** **`tensor`** `torch.Tensor` Read-only int32 tensor \[batch \* groups, 1, query_tiles] for counts, or \[batch \* groups, 1, query_tiles, key_tiles] for indices. Backward metadata swaps query_tiles and key_tiles. Each provider owns one KV group, with a common number of query heads per group. --- **`batch`** `int` Original attention batch size before provider folding. --- **`heads_per_group`** `int` Number of consecutive query heads sharing a KV group. --- **Returns:** `torch.Tensor` Independent contiguous int32 metadata \[batch, query_heads, ...], ```python nemo_automodel.components.models.step5.fa4._region_tiles( ids: torch.Tensor, token_regions: torch.Tensor, regions: int, fwd_k: int, bwd_k: int ) -> tuple[torch.Tensor, torch.Tensor] ``` Union selected regions for bounded chunks of 128-query tiles. **Parameters:** **`ids`** `torch.Tensor` Sorted region IDs \[batch, groups, queries, slots]. --- **`token_regions`** `torch.Tensor` Monotonic region IDs \[batch, keys]. --- **`regions`** `int` Number of global regions; also the unused-route sentinel. --- **`fwd_k`** `int` Forward key tile width. --- **`bwd_k`** `int` Backward key tile width. --- **Returns:** `torch.Tensor` Forward and backward tile-presence grids \[batch, groups, query_tiles, ```python nemo_automodel.components.models.step5.fa4._selected_block_lists( selection: nemo_automodel.components.models.step5.sparse.RegionSelection, regions: int, query_start: int, fwd_k: int, bwd_k: int, fwd_q: int = 128 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] ``` Create route-derived forward/reverse FA4 lists with bounded scratch space. **Parameters:** **`selection`** `RegionSelection` Sorted region IDs \[batch, groups, queries, slots] and token-region IDs \[batch, keys] with monotonic key-region order. --- **`regions`** `int` Number of global regions. --- **`query_start`** `int` Global token offset of the local query shard. --- **`fwd_k`** `int` Key tokens represented by each forward metadata block. --- **`bwd_k`** `int` Key tokens represented by each backward metadata block. --- **`fwd_q`** `int` — default: 128 Query tokens per forward metadata block, either 128 or 256. --- **Returns:** `torch.Tensor` Forward counts \[batch \* groups, 1, ceil(queries / fwd_q)] and ```python nemo_automodel.components.models.step5.fa4._split_backward_full_blocks( counts: torch.Tensor, indices: torch.Tensor, prefix_length: int ) -> nemo_automodel.components.models.step5.fa4._BackwardBlockPartition ``` Classify complete causal rectangles using SM90 D192 backward tiles. **Parameters:** **`counts`** `torch.Tensor` Int32 counts \[provider_batches, heads, key_tiles]. --- **`indices`** `torch.Tensor` Ascending unique nonnegative query-block IDs of shape \[provider_batches, heads, key_tiles, query_tiles]. Entries beyond counts are ignored. Provider-folded tensors have heads=1; native layout has provider_batches=batch and heads=query_heads. --- **`prefix_length`** `int` Proven all-selected query prefix in one unpadded B1 CP1 document whose positions equal physical token indices. --- **Returns:** `_BackwardBlockPartition` Independent partial/full counts and index lists with the input layouts. ```python nemo_automodel.components.models.step5.fa4._split_forward_full_blocks( counts: torch.Tensor, indices: torch.Tensor, prefix_length: int ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] ``` Separate a low contiguous full-block prefix without changing K order. **Parameters:** **`counts`** `torch.Tensor` Int32 selected K-block counts \[provider_batches, heads, query_tiles]. Provider-folded layout uses provider_batches=batch\*groups and heads=1; native head layout uses provider_batches=batch and heads=query_heads. --- **`indices`** `torch.Tensor` Sorted unique nonnegative K-block indices \[provider_batches, heads, query_tiles, key_tiles], with ignored trailing slots. Both tensors use SM90 D192 tiles: 128 queries and 112 keys. --- **`prefix_length`** `int` Proven all-selected initial query count in a single unpadded CP1 document with physical positions starting at zero. --- **Returns:** `torch.Tensor` Partial-block counts, shifted partial-block indices, and full-block ```python nemo_automodel.components.models.step5.fa4.fa4_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, positions: torch.Tensor, key_positions: torch.Tensor, key_padding: torch.Tensor, documents: torch.Tensor, query_start: int, window: int | None, selection: nemo_automodel.components.models.step5.sparse.RegionSelection | None, window_layout: nemo_automodel.components.models.step5.sparse.CausalDocumentLayout | None = None ) -> torch.Tensor ``` Run FA4 forward/backward with compact CSA metadata on SM90 or SM100. **Parameters:** **`query`** `torch.Tensor` BF16 CUDA tensor \[batch, local_queries, query_heads, head_dim]. --- **`key`** `torch.Tensor` BF16 CUDA tensor \[batch, global_keys, kv_heads, head_dim]. --- **`value`** `torch.Tensor` Same shape/dtype/device as key. --- **`positions`** `torch.Tensor` Int tensor \[batch, local_queries], document-local positions. --- **`key_positions`** `torch.Tensor` Int tensor \[batch, global_keys], document-local positions. --- **`key_padding`** `torch.Tensor` Boolean tensor \[batch, global_keys], true for padding. --- **`documents`** `torch.Tensor` Int tensor \[batch, global_keys], document IDs. --- **`query_start`** `int` Physical global offset of the contiguous query shard. --- **`window`** `int | None` Sliding window size, or None for CSA. --- **`selection`** `RegionSelection | None` Compact sorted region metadata, or None for sliding attention. --- **`window_layout`** `CausalDocumentLayout | None` — default: None Opt-in valid SWA document segments in global token coordinates. The existing halo may contain local alignment padding. --- **Returns:** `torch.Tensor` BF16 tensor \[batch, local_queries, query_heads, head_dim].