> This page is for version Nightly (default).
> 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.sparse

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

## Module Contents

### Classes

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

### Functions

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

### Data

[`_compiled_gather_region_payload`](#nemo_automodel-components-models-step5-sparse-_compiled_gather_region_payload)

### API

```python
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.

---

```python
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.

---

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

---

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

**Bases:** `LayerNorm`

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

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

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

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

**Bases:** [Step3p5RMSNorm](/nemo-automodel/nemo_automodel/components/models/step3p5/layers#nemo_automodel-components-models-step3p5-layers-Step3p5RMSNorm)

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

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

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

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

**Bases:** `Linear`

Keep FP32 source weights while matching the reference BF16 projection.

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

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

```python
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`**

---

```python
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`** `int` — default: 0

Physical start of the contiguous local query shard.

---

**`metadata`** `RegionMetadata | None` — default: 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]

```python
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.

```python
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.

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

```python
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
)
```

```python
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.

```python
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.

```python
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]

```python
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

```python
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)].

```python
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`** `bool` — default: False

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

---

**`csa_backward_full_blocks`** `bool` — default: False

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

---

**`window`** `int | None` — default: 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

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