nemo_automodel.components.models.step5.fa4

Step-5 attention through FA4, with compact region membership metadata.

Module Contents

Classes

Name	Description
`_BackwardBlockPartition`	Independent int32 lists in the original provider/head layout.

Functions

Name	Description
`_all_padding_attention`	Copy singleton-key attention while retaining explicit zero Q/K gradients.
`_block_lists`	Build block lists without allocating a dense block grid.
`_causal_document_attention`	Run the existing B1 causal prefix or proven singleton-padding operation.
`_compact_selected_tiles`	Compact a boolean block grid into FA4 partial-block counts and indices.
`_load_fa4`	Load FA4 and cache a mask, optionally specializing a proven document.
`_load_fa4_varlen`	Load native FA4 causal varlen attention without a custom mask.
`_provider_heads`	Copy provider-specific block metadata into FA4’s query-head axis.
`_region_tiles`	Union selected regions for bounded chunks of 128-query tiles.
`_selected_block_lists`	Create route-derived forward/reverse FA4 lists with bounded scratch space.
`_split_backward_full_blocks`	Classify complete causal rectangles using SM90 D192 backward tiles.
`_split_forward_full_blocks`	Separate a low contiguous full-block prefix without changing K order.
`fa4_attention`	Run FA4 forward/backward with compact CSA metadata on SM90 or SM100.

API

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

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

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

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.

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.

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.

nemo_automodel.components.models.step5.fa4._load_fa4_varlen() -> collections.abc.Callable

Load native FA4 causal varlen attention without a custom mask.

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, …],

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,

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

intDefaults to 128

Query tokens per forward metadata block, either 128 or 256.

Returns: torch.Tensor

Forward counts [batch * groups, 1, ceil(queries / fwd_q)] and

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.

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

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 | NoneDefaults to 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].