> This page is for version Latest.
> 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.layers

Trainable Step-5 decoder layers and context-parallel attention.

## Module Contents

### Classes

| Name                                                                                    | Description                                                               |
| --------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`Step5Attention`](#nemo_automodel-components-models-step5-layers-Step5Attention)       | GQA with windowed or CSA-selected keys and contiguous query-shard CP.     |
| [`Step5DecoderLayer`](#nemo_automodel-components-models-step5-layers-Step5DecoderLayer) | Pre-norm decoder with FP32 residual accumulation and optional shared MoE. |
| [`Step5MLP`](#nemo_automodel-components-models-step5-layers-Step5MLP)                   | SwiGLU with separate gate/up projections and FP32 asymmetric clamping.    |
| [`Step5RMSNorm`](#nemo_automodel-components-models-step5-layers-Step5RMSNorm)           | Fuse the main-stream norm while retaining FP32 additive scale parameters. |

### Functions

| Name                                                                                          | Description                                                               |
| --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`_activation_rms_norm`](#nemo_automodel-components-models-step5-layers-_activation_rms_norm) | -                                                                         |
| [`_grouped_rms_norm`](#nemo_automodel-components-models-step5-layers-_grouped_rms_norm)       | Normalize grouped query rows without copying interleaved QKV gaps.        |
| [`_residual_rms_norm`](#nemo_automodel-components-models-step5-layers-_residual_rms_norm)     | -                                                                         |
| [`_rms_norm`](#nemo_automodel-components-models-step5-layers-_rms_norm)                       | Normalize \[..., width] in FP32 before rounding once to the output dtype. |
| [`_rms_norm_with_rstd`](#nemo_automodel-components-models-step5-layers-_rms_norm_with_rstd)   | Normalize in FP32 and retain the row statistic needed by backward.        |
| [`clamped_swiglu`](#nemo_automodel-components-models-step5-layers-clamped_swiglu)             | Fuse FP32 SwiGLU and its derivative without FP32 activation temporaries.  |
| [`fused_attention`](#nemo_automodel-components-models-step5-layers-fused_attention)           | Run the shared compiled FlexAttention kernel without a dense token mask.  |

### Data

[`region_replay`](#nemo_automodel-components-models-step5-layers-region_replay)

### API

```python
class nemo_automodel.components.models.step5.layers.Step5Attention(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)
```

**Bases:** `Module`

GQA with windowed or CSA-selected keys and contiguous query-shard CP.

**`cp_group`** `ProcessGroup | None = None`

---

**`g_proj`**

---

**`head_dim`** `= config.head_dim`

---

**`heads_per_group`** `= self.num_heads // self.num_kv_heads`

---

**`indexer`**

---

**`k_norm`**

---

**`num_heads`** `= config.num_attention_heads`

---

**`num_kv_heads`** `= config.num_attention_groups`

---

**`o_proj`**

---

**`q_norm`**

---

**`qkv_proj`**

---

**`rotary_emb`** `= Step5RotaryEmbedding(config, layer_idx, span=span)`

---

**`window`**

---

```python
nemo_automodel.components.models.step5.layers.Step5Attention._select_regions(
    hidden: torch.Tensor,
    positions: torch.Tensor,
    key_positions: torch.Tensor,
    key_padding: torch.Tensor,
    query_start: int,
    sparse_metadata: nemo_automodel.components.models.step5.sparse.RegionMetadata | None
) -> nemo_automodel.components.models.step5.sparse.RegionSelection
```

Select once or replay integer routes during activation recomputation.

**Parameters:**

**`hidden`** `torch.Tensor`

Local normalized states \[batch, local_queries, hidden].

---

**`positions`** `torch.Tensor`

Int64 local document positions \[batch, local_queries].

---

**`key_positions`** `torch.Tensor`

Int64 global document positions \[batch, global_keys].

---

**`key_padding`** `torch.Tensor`

Boolean global padding flags \[batch, global_keys].

---

**`query_start`** `int`

Physical start of the local query shard.

---

**`sparse_metadata`** `RegionMetadata | None`

Layer-independent global/local region layout.

---

**Returns:** `RegionSelection`

Integer-only region selection; no projected K/V or activations are

```python
nemo_automodel.components.models.step5.layers.Step5Attention.forward(
    hidden: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    global_positions: torch.Tensor,
    global_padding: torch.Tensor,
    sparse_metadata: nemo_automodel.components.models.step5.sparse.RegionMetadata | None = None
) -> torch.Tensor
```

Attend local queries to valid global keys, preserving document boundaries.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[batch, local_sequence, hidden].

---

**`positions`** `torch.Tensor`

Tensor of shape \[batch, local_sequence], document-local positions.

---

**`padding_mask`** `torch.Tensor`

Boolean tensor of shape \[batch, local_sequence], true for padding.

---

**`global_positions`** `torch.Tensor`

Int tensor \[batch, global_sequence], gathered once per model call.

---

**`global_padding`** `torch.Tensor`

Boolean tensor \[batch, global_sequence], true for padding.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, local_sequence, hidden].

```python
nemo_automodel.components.models.step5.layers.Step5Attention.setup_cp_attention(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None
```

Attach the CP group through AutoModel's model-owned attention hook.

```python
class nemo_automodel.components.models.step5.layers.Step5DecoderLayer(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    dtype: torch.dtype
)
```

**Bases:** `Module`

Pre-norm decoder with FP32 residual accumulation and optional shared MoE.

**`input_layernorm`**

---

**`moe`**

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

**`share_expert`**

---

```python
nemo_automodel.components.models.step5.layers.Step5DecoderLayer.forward(
    hidden: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    global_positions: torch.Tensor,
    global_padding: torch.Tensor,
    compute_dtype: torch.dtype,
    sparse_metadata: nemo_automodel.components.models.step5.sparse.RegionMetadata | None = None
) -> torch.Tensor
```

Apply attention and feed-forward residual branches.

**Parameters:**

**`hidden`** `torch.Tensor`

FP32 tensor of shape \[batch, local_sequence, hidden].

---

**`positions`** `torch.Tensor`

Tensor of shape \[batch, local_sequence].

---

**`padding_mask`** `torch.Tensor`

Boolean tensor of shape \[batch, local_sequence], true for padding.

---

**`global_positions`** `torch.Tensor`

Int tensor \[batch, global_sequence], gathered once per model call.

---

**`global_padding`** `torch.Tensor`

Boolean tensor \[batch, global_sequence], true for padding.

---

**`compute_dtype`** `torch.dtype`

Activation dtype from the completed embedding projection.

---

**`sparse_metadata`** `RegionMetadata | None` — default: None

Shared global region layout and local query metadata.

---

**Returns:** `torch.Tensor`

FP32 tensor of shape \[batch, local_sequence, hidden].

```python
nemo_automodel.components.models.step5.layers.Step5DecoderLayer.nemo_checkpoint_context_fn(
    context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
) -> collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
```

Compose the native AC context with one integer-route log per call.

**Parameters:**

**`context_fn`** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None`

Existing plain/selective checkpoint context factory.

---

**Returns:** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None`

A factory with a distinct recorder for each live pipeline microbatch,

```python
class nemo_automodel.components.models.step5.layers.Step5MLP(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    width: int,
    limit: float,
    dtype: torch.dtype
)
```

**Bases:** `Module`

SwiGLU with separate gate/up projections and FP32 asymmetric clamping.

**`down_proj`**

---

**`gate_proj`**

---

**`up_proj`**

---

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

Evaluate the dense or shared feed-forward block.

**Parameters:**

**`hidden`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden] in the input dtype.

```python
class nemo_automodel.components.models.step5.layers.Step5RMSNorm(
    hidden_size: int,
    eps: float = 1e-05,
    compile_norm: bool = False
)
```

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

Fuse the main-stream norm while retaining FP32 additive scale parameters.

Indexer normalization deliberately uses its separate implementation: its
checkpoint requires rounding the scale to the activation dtype first.

**Parameters:**

**`hidden_size`** `int`

Normalized feature width.

---

**`eps`** `float` — default: 1e-05

Additive variance epsilon.

---

**`compile_norm`** `bool` — default: False

Enable the compiled CUDA core. CPU always uses eager math.

---

```python
nemo_automodel.components.models.step5.layers.Step5RMSNorm._compiled_forward(
    hidden: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor
```

Run the compiled CUDA norm on \[..., hidden_size] and restore the input shape.

```python
nemo_automodel.components.models.step5.layers.Step5RMSNorm.forward(
    hidden: torch.Tensor,
    output_dtype: torch.dtype | None = None
) -> torch.Tensor
```

Normalize \[..., hidden_size] before the optional projection-input cast.

**Parameters:**

**`hidden`** `torch.Tensor`

Activations \[..., hidden_size] with arbitrary leading axes.
Grouped queries use \[batch, sequence, groups, heads_per_group,
hidden_size] and may be a strided slice of interleaved QKV.
Decoder residual inputs remain FP32 through normalization.

---

**`output_dtype`** `torch.dtype | None` — default: None

Final output dtype; defaults to the input dtype.

---

**Returns:** `torch.Tensor`

Normalized activations with the input shape. The learned FP32 scale

```python
nemo_automodel.components.models.step5.layers._activation_rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.step5.layers._grouped_rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]
```

Normalize grouped query rows without copying interleaved QKV gaps.

**Parameters:**

**`hidden`** `torch.Tensor`

Query tensor \[token_groups, heads_per_group, channels], where
token_groups flattens batch, sequence and KV group. The outer stride
includes the interleaved K/V rows, and channels has unit stride.

---

**`weight`** `torch.Tensor`

FP32 additive scale \[channels].

---

**`eps`** `float`

Additive variance epsilon.

---

**`output_dtype`** `torch.dtype`

Final normalized activation dtype.

---

**Returns:** `torch.Tensor`

Normalized query \[token_groups, heads_per_group, channels] and FP32

```python
nemo_automodel.components.models.step5.layers._residual_rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> torch.Tensor
```

```python
nemo_automodel.components.models.step5.layers._rms_norm(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> torch.Tensor
```

Normalize \[..., width] in FP32 before rounding once to the output dtype.

```python
nemo_automodel.components.models.step5.layers._rms_norm_with_rstd(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    output_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]
```

Normalize in FP32 and retain the row statistic needed by backward.

```python
nemo_automodel.components.models.step5.layers.clamped_swiglu(
    gate: torch.Tensor,
    up: torch.Tensor,
    limit: float
) -> torch.Tensor
```

Fuse FP32 SwiGLU and its derivative without FP32 activation temporaries.

**Parameters:**

**`gate`** `torch.Tensor`

Projection tensor \[batch, sequence, intermediate].

---

**`up`** `torch.Tensor`

Projection tensor with the same shape and dtype.

---

**`limit`** `float`

Positive clamp bound; non-positive values disable clamping.

---

**Returns:** `torch.Tensor`

Tensor with the same shape and dtype as gate.

```python
nemo_automodel.components.models.step5.layers.fused_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
) -> torch.Tensor
```

Run the shared compiled FlexAttention kernel without a dense token mask.

**Parameters:**

**`query`** `torch.Tensor`

CUDA tensor \[batch, local_queries, query_heads, head_dim].

---

**`key`** `torch.Tensor`

CUDA tensor \[batch, global_keys, kv_heads, head_dim].

---

**`value`** `torch.Tensor`

CUDA tensor with the same shape and dtype as key.

---

**`positions`** `torch.Tensor`

Document-local positions \[batch, local_queries].

---

**`key_positions`** `torch.Tensor`

Document-local positions \[batch, global_keys].

---

**`key_padding`** `torch.Tensor`

Boolean tensor \[batch, global_keys], true for padding.

---

**`documents`** `torch.Tensor`

Document ids \[batch, global_keys].

---

**`query_start`** `int`

Global offset of the contiguous query shard.

---

**`window`** `int | None`

Sliding window size, or None for CSA layers.

---

**`selection`** `RegionSelection | None`

Compact sorted CSA region ids, or None for sliding attention.

---

**Returns:** `torch.Tensor`

Attention output \[batch, local_queries, query_heads, head_dim].

```python
nemo_automodel.components.models.step5.layers.region_replay: RecomputeReplay[RegionSelection] = RecomputeReplay('Step5 CSA route')
```