nemo_automodel.components.models.step5.sparse

CSA proxy compression and learned region selection for Step-5.

Module Contents

Classes

Name	Description
`CausalDocumentLayout`	Native varlen layout for the valid prefix of one single-batch shard.
`RegionMetadata`	Layer-independent CSA layout for one contiguous CP query shard.
`RegionSelection`	Sparse selection tensors, all local to one CP rank’s queries.
`Step5IndexerLayerNorm`	Indexer variance E[x²] - E[x]², evaluated inside the module’s FSDP hooks.
`Step5IndexerRMSNorm`	Indexer RMSNorm with source-required activation-dtype parameter rounding.
`Step5IndexerWeightProjection`	Keep FP32 source weights while matching the reference BF16 projection.
`Step5SparseIndexer`	Checkpoint-backed, discrete CSA selector; LM loss does not train top-k indices.

Functions

Name	Description
`_gather_region_payload`	Gather existing token values without performing floating-point arithmetic.
`_pack_regions`	-
`_score_regions`	-
`_summarize_region_payload`	Preserve eager FP32 arithmetic and both required quantization roundings.
`_summary_ranges`	Merge each batch’s needed historical intervals without inspecting GPU data.
`gather_region_payload`	Gather only the historical regions needed by this rank’s sparse queries.
`indexer_score_tile`	Fuse weighted-ReLU scoring into one bounded FP32 output tile.
`pack_region_bitmap`	Pack selected regions directly without int64 scatter workspaces.
`region_metadata`	Prepare document/region metadata once per model forward.

Data

_compiled_gather_region_payload

API

class nemo_automodel.components.models.step5.sparse.CausalDocumentLayout(
    key_start: int,
    key_stop: int,
    cu_seqlens_q: torch.Tensor,
    cu_seqlens_k: torch.Tensor,
    max_seqlen_q: int,
    max_seqlen_k: int,
    valid_queries: int,
    padding_key: int | None = None,
    leading_padding: int = 0
)

Dataclass

Native varlen layout for the valid prefix of one single-batch shard.

cu_seqlens_k

Tensor

Int32 cumulative key counts for those same documents.

cu_seqlens_q

Tensor

Int32 cumulative query counts for the intersected documents.

key_start

int

Physical start of the first queried document’s keys.

key_stop

int

Exclusive last queried token; future keys must be omitted.

leading_padding

int = 0

Native SWA dummy keys before the received halo, used to keep document-relative 96-token K tile origins aligned across CP. They are outside every valid query window.

max_seqlen_k

int

Largest causal key prefix for one queried document.

max_seqlen_q

int

Largest local document query segment.

padding_key

int | None = None

Last document’s first token, if a trailing padding-query suffix must retain the original singleton-key attention semantics.

valid_queries

int

Query prefix evaluated by the native varlen kernel. Zero identifies a fully padded batch row with a proven padding_key.

class nemo_automodel.components.models.step5.sparse.RegionMetadata(
    token_regions: torch.Tensor,
    documents: torch.Tensor,
    region_documents: torch.Tensor,
    current: torch.Tensor,
    query_documents: torch.Tensor,
    dense_region_ids: torch.Tensor,
    sparse_queries: torch.Tensor,
    score_tiles: tuple[tuple[int, int, int, int, int], ...],
    region_count: int,
    needs_projection: bool,
    causal_layout: nemo_automodel.components.models.step5.sparse.CausalDocumentLayout | None = None,
    window_layout: nemo_automodel.components.models.step5.sparse.CausalDocumentLayout | None = None,
    summary_ranges: tuple[tuple[int, int, int], ...] = (),
    score_range_indices: tuple[int, ...] = (),
    region_sources: torch.Tensor | None = None,
    region_source_valid: torch.Tensor | None = None,
    all_padding_key: int | None = None,
    causal_batch_layouts: tuple[nemo_automodel.components.models.step5.sparse.CausalDocumentLayout, ...] | None = None,
    forward_full_block_prefix: int = 0,
    backward_full_block_prefix: int = 0,
    canonical_region_size: int = 0
)

Dataclass

Layer-independent CSA layout for one contiguous CP query shard.

all_padding_key

int | None = None

Singleton global key for a fully padded, fully selected local CSA shard, or None when the shortcut cannot be established.

backward_full_block_prefix

int = 0

The same proof as forward_full_block_prefix, independently enabled for backward mask skipping. Zero is fallback.

canonical_region_size

int = 0

Region width when the full-block proof also establishes canonical global token coordinates; zero disables the optional arithmetic mask specialization.

causal_batch_layouts

tuple[CausalDocumentLayout, ...] | None = None

Prefix or all-padding layouts when both B2 rows satisfy native causal attention semantics.

causal_layout

CausalDocumentLayout | None = None

Optional native varlen metadata for this query shard.

current

Tensor

Int64 current region IDs [batch, local_queries].

dense_region_ids

Tensor

Int32 ascending full-prefix IDs [batch, local_queries, min(topk, regions) + 1]; unused entries contain region_count.

documents

Tensor

Int64 global document IDs [batch, global_keys].

forward_full_block_prefix

int = 0

Initial all-selected query count for opt-in forward full blocks, proven only for unpadded B1 CP1 inputs with positions exactly equal to global token indices. Zero is fallback.

needs_projection

bool

Whether any query on any CP rank needs learned scores. All ranks derive it from the same global layout before collectives.

query_documents

Tensor

Int64 document IDs [batch, local_queries].

region_count

int

Number of global regions, also the unused-route sentinel.

region_documents

Tensor

Int64 document IDs [batch, regions].

region_source_valid

Tensor | None = None

Boolean shape matching region_sources, excluding padding tokens as well as missing slots.

region_sources

Tensor | None = None

Optional int32 source token indices [batch, regions, region_size], shared across layers. Missing slots contain -1.

score_range_indices

tuple[int, ...] = ()

Index into summary_ranges for each score tile.

score_tiles

tuple[tuple[int, int, int, int, int], ...]

Host tuples (batch, query_start, query_stop, first_region, past_last_region), bounded to 128 local queries each. Only tiles containing a sparse query are included.

sparse_queries

Tensor

Boolean [batch, local_queries], true only when more than topk historical regions in the same document are visible.

summary_ranges

tuple[tuple[int, int, int], ...] = ()

Merged (batch, first_region, past_last_region) intervals needed by score tiles. Empty on prefix-only or irregular layouts.

token_regions

Tensor

Int64 global region IDs [batch, global_keys].

window_layout

CausalDocumentLayout | None = None

Optional native SWA layout, including at most window - 1 preceding keys. Independent of whether CSA selection is sparse.

class nemo_automodel.components.models.step5.sparse.RegionSelection(
    region_ids: torch.Tensor,
    token_regions: torch.Tensor,
    query_documents: torch.Tensor,
    key_documents: torch.Tensor,
    region_count: int,
    all_regions_selected: bool = False,
    causal_layout: nemo_automodel.components.models.step5.sparse.CausalDocumentLayout | None = None,
    all_padding_key: int | None = None,
    causal_batch_layouts: tuple[nemo_automodel.components.models.step5.sparse.CausalDocumentLayout, ...] | None = None,
    forward_full_block_prefix: int = 0,
    backward_full_block_prefix: int = 0,
    canonical_region_size: int = 0
)

Dataclass

Sparse selection tensors, all local to one CP rank’s queries.

all_padding_key

int | None = None

The only allowed global key for a fully padded shard, after proving standard document positions and trailing padding.

all_regions_selected

bool = False

Every query in this local shard selects all causal regions in its own document; membership adds no restriction.

backward_full_block_prefix

int = 0

The same proof as forward_full_block_prefix, independently enabled for backward mask skipping. Zero is fallback.

canonical_region_size

int = 0

Region width when the full-block proof also establishes canonical global token coordinates; zero disables the optional arithmetic mask specialization.

causal_batch_layouts

tuple[CausalDocumentLayout, ...] | None = None

Proven prefix or all-padding layout for both rows of a fully selected B2 shard; None if either row needs the fallback.

causal_layout

CausalDocumentLayout | None = None

Optional compact varlen layout for fully selected, single-batch shards with contiguous document positions.

forward_full_block_prefix

int = 0

Number of initial queries proven to select every causal key in one unpadded global B1 document, with CP1 and canonical positions starting at zero. Zero disables the proof.

key_documents

Tensor

Int64 tensor of shape [batch, global_keys].

query_documents

Tensor

Int64 tensor of shape [batch, local_queries].

region_count

int

Number of global regions; also the unused-route sentinel.

region_ids

Tensor

Sorted int32 tensor [batch, groups, local_queries, topk + 1]. Unused slots contain a sentinel larger than any valid region.

token_regions

Tensor

Int64 tensor of shape [batch, global_keys].

class nemo_automodel.components.models.step5.sparse.Step5IndexerLayerNorm()

Bases: LayerNorm

Indexer variance E[x²] - E[x]², evaluated inside the module’s FSDP hooks.

nemo_automodel.components.models.step5.sparse.Step5IndexerLayerNorm.forward(
    hidden: torch.Tensor
) -> torch.Tensor

Normalize […, proxy_dim] with BF16-rounded scale/bias and FP32 math.

class nemo_automodel.components.models.step5.sparse.Step5IndexerRMSNorm()

Bases: Step3p5RMSNorm

Indexer RMSNorm with source-required activation-dtype parameter rounding.

nemo_automodel.components.models.step5.sparse.Step5IndexerRMSNorm.forward(
    hidden: torch.Tensor
) -> torch.Tensor

Normalize […, proxy_dim] activations and preserve their dtype.

class nemo_automodel.components.models.step5.sparse.Step5IndexerWeightProjection()

Bases: Linear

Keep FP32 source weights while matching the reference BF16 projection.

nemo_automodel.components.models.step5.sparse.Step5IndexerWeightProjection.forward(
    hidden: torch.Tensor
) -> torch.Tensor

Project […, hidden_size] to signed FP32 indexer head weights.

class nemo_automodel.components.models.step5.sparse.Step5SparseIndexer(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    layer_idx: int,
    dtype: torch.dtype
)

Bases: Module

Checkpoint-backed, discrete CSA selector; LM loss does not train top-k indices.

config

= config.sparse_config

k_norm

k_proj

q_norm

= Step5IndexerRMSNorm(width, config.rms_norm_eps)

q_proj

rotary_emb

= Step5RotaryEmbedding(config, layer_idx, span=span)

w_proj

z_proj

nemo_automodel.components.models.step5.sparse.Step5SparseIndexer.forward(
    hidden: torch.Tensor,
    positions: torch.Tensor,
    keys: torch.Tensor | None,
    compression_logits: torch.Tensor | None,
    key_positions: torch.Tensor,
    key_padding: torch.Tensor,
    query_start: int = 0,
    metadata: nemo_automodel.components.models.step5.sparse.RegionMetadata | None = None
) -> nemo_automodel.components.models.step5.sparse.RegionSelection

Select historical CSA regions, scoring only genuinely sparse queries.

Parameters:

hidden

torch.Tensor

Tensor of shape [batch, local_queries, hidden].

positions

torch.Tensor

Int64 document-local positions [batch, local_queries].

keys

torch.Tensor | None

Projected global index keys [batch, global_keys, proxy_dim], or None when every global query sees at most topk historical regions.

compression_logits

torch.Tensor | None

Same shape as keys, or None on that full-prefix path.

key_positions

torch.Tensor

Int64 document-local positions [batch, global_keys].

key_padding

torch.Tensor

Boolean padding flags [batch, global_keys].

query_start

intDefaults to 0

Physical start of the contiguous local query shard.

metadata

RegionMetadata | NoneDefaults to None

Shared document/region layout for the current model call. Direct component callers may omit it.

Returns: RegionSelection

Ascending selected region IDs [batch, groups, local_queries, topk + 1]

nemo_automodel.components.models.step5.sparse.Step5SparseIndexer.project_keys(
    hidden: torch.Tensor,
    positions: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Compute normalized index keys and per-dimension compression logits.

Parameters:

hidden

torch.Tensor

Tensor of shape [batch, sequence, hidden].

positions

torch.Tensor

Tensor of shape [batch, sequence].

Returns: tuple[torch.Tensor, torch.Tensor]

Two tensors of shape [batch, sequence, proxy_dim], in the hidden dtype.

nemo_automodel.components.models.step5.sparse._gather_region_payload(
    keys: torch.Tensor,
    logits: torch.Tensor,
    sources: torch.Tensor,
    valid: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Gather existing token values without performing floating-point arithmetic.

nemo_automodel.components.models.step5.sparse._pack_regions(
    ids,
    bitmap,
    total,
    slots: tl.constexpr,
    regions,
    words,
    BLOCK: tl.constexpr
)

nemo_automodel.components.models.step5.sparse._score_regions(
    q,
    w,
    k,
    out,
    query_start,
    queries,
    regions,
    q_batch,
    q_token,
    q_group,
    q_head,
    w_batch,
    w_token,
    w_group,
    k_batch,
    k_region,
    out_batch,
    out_group,
    out_token,
    GROUPS: tl.constexpr,
    HEADS: tl.constexpr,
    DIM: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_Q: tl.constexpr,
    BLOCK_K: tl.constexpr
)

nemo_automodel.components.models.step5.sparse._summarize_region_payload(
    packed_keys: torch.Tensor,
    packed_z: torch.Tensor
) -> torch.Tensor

Preserve eager FP32 arithmetic and both required quantization roundings.

Inductor can remove the intermediate BF16 rounding when this arithmetic is compiled. Even emulate_precision_casts changes a few FP8 summaries through FP32 arithmetic, so only the integer gather is compiled.

nemo_automodel.components.models.step5.sparse._summary_ranges(
    score_tiles: list[tuple[int, int, int, int, int]]
) -> tuple[tuple[tuple[int, int, int], ...], tuple[int, ...]]

Merge each batch’s needed historical intervals without inspecting GPU data.

nemo_automodel.components.models.step5.sparse.gather_region_payload(
    keys: torch.Tensor,
    logits: torch.Tensor,
    sources: torch.Tensor,
    valid: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Gather only the historical regions needed by this rank’s sparse queries.

Parameters:

keys

torch.Tensor

Original indexer keys [batch, global_tokens, proxy_dim].

logits

torch.Tensor

Compression logits, shaped like keys.

sources

torch.Tensor

Int32 token indices [batch, required_regions, region_size]. Missing slots contain -1.

valid

torch.Tensor

Boolean shape matching sources, excluding padding and missing slots.

Returns: torch.Tensor

Key and logit tensors [batch, required_regions, region_size, proxy_dim]

nemo_automodel.components.models.step5.sparse.indexer_score_tile(
    query: torch.Tensor,
    weights: torch.Tensor,
    summary: torch.Tensor,
    start: int,
    stop: int
) -> torch.Tensor

Fuse weighted-ReLU scoring into one bounded FP32 output tile.

Parameters:

query

torch.Tensor

E4M3-rounded BF16 tensor [batch, queries, groups, heads, dim].

weights

torch.Tensor

FP32 signed head weights [batch, queries, groups, heads].

summary

torch.Tensor

E4M3-rounded BF16 region summaries [batch, regions, dim].

start

int

First query included in this tile.

stop

int

Exclusive last query included in this tile.

Returns: torch.Tensor

FP32 scores [batch, groups, stop - start, regions]. No gradients are

nemo_automodel.components.models.step5.sparse.pack_region_bitmap(
    ids: torch.Tensor,
    regions: int
) -> torch.Tensor

Pack selected regions directly without int64 scatter workspaces.

Parameters:

ids

torch.Tensor

Sorted CUDA int32 tensor [batch, groups, queries, slots]; ids at or above regions are padding. Repeated valid ids are allowed.

regions

int

Total number of key regions.

Returns: torch.Tensor

CUDA int32 bitmap [batch, groups, queries, ceil(regions / 32)].

nemo_automodel.components.models.step5.sparse.region_metadata(
    key_positions: torch.Tensor,
    key_padding: torch.Tensor,
    query_start: int,
    queries: int,
    region_size: int,
    topk: int,
    window: int | None = None,
    csa_forward_full_blocks: bool = False,
    csa_backward_full_blocks: bool = False
) -> nemo_automodel.components.models.step5.sparse.RegionMetadata

Prepare document/region metadata once per model forward.

Parameters:

key_positions

torch.Tensor

Int64 document-local positions [batch, global_keys].

key_padding

torch.Tensor

Boolean padding flags [batch, global_keys].

query_start

int

Physical start of this rank’s contiguous query shard.

queries

int

Number of local queries.

region_size

int

Tokens per region, reset at every document boundary.

topk

int

Maximum number of selected historical regions.

csa_forward_full_blocks

boolDefaults to False

Retain a single-document CP1 proof for FA4 forward mask skipping, using this call’s existing host metadata.

csa_backward_full_blocks

boolDefaults to False

Retain the same proof independently for backward mask skipping, without another host transfer.

window

int | NoneDefaults to None

Optional native SWA layout width. A globally trailing padded suffix keeps its custom-mask path. No additional host synchronization.

Returns: RegionMetadata

Shared immutable metadata. Only linear token/region storage and a bounded

nemo_automodel.components.models.step5.sparse._compiled_gather_region_payload = torch.compile(_gather_region_payload, dynamic=True)