> 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.beam.layers

Beam gated attention, root10 SSMax, rotary precision, and SwiGLU/router overrides.

## Module Contents

### Classes

| Name                                                                           | Description                                                                                   |
| ------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------- |
| [`BeamAttention`](#nemo_automodel-components-models-beam-layers-BeamAttention) | Gated attention with per-head q/k RMSNorm and full-width output gating.                       |
| [`BeamGate`](#nemo_automodel-components-models-beam-layers-BeamGate)           | Normalize top-1 routing like HF, including its epsilon and gradients.                         |
| [`BeamMLP`](#nemo_automodel-components-models-beam-layers-BeamMLP)             | Shared dense MLP with HF's fp32 SwiGLU boundary, preserving parameter names and MLP dispatch. |
| [`BeamMoE`](#nemo_automodel-components-models-beam-layers-BeamMoE)             | Shared MoE with BeamGate installed before loading or parallelization.                         |
| [`Block`](#nemo_automodel-components-models-beam-layers-Block)                 | One Beam decoder layer: sandwich norms + depth-scaled residual adds.                          |

### Functions

| Name                                                                                                       | Description                                                                              |
| ---------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------- |
| [`_cp_local_positions`](#nemo_automodel-components-models-beam-layers-_cp_local_positions)                 | Select the two load-balanced sequence chunks owned by this CP rank.                      |
| [`apply_rope_hf`](#nemo_automodel-components-models-beam-layers-apply_rope_hf)                             | Apply HF Beam's rotary expression and preserve its backward rounding.                    |
| [`depth_scaled_residual_coefs`](#nemo_automodel-components-models-beam-layers-depth_scaled_residual_coefs) | (sublayer_scale, residual_scale) for the depth-scaled residual add.                      |
| [`hf_rope_tables`](#nemo_automodel-components-models-beam-layers-hf_rope_tables)                           | Expand original FP32 rotary tables into HF's duplicated-half layout.                     |
| [`is_moe_layer`](#nemo_automodel-components-models-beam-layers-is_moe_layer)                               | Beam marks every layer `"dense"` or `"sparse"` in `config.mlp_layer_types`.              |
| [`layer_rope_parameters`](#nemo_automodel-components-models-beam-layers-layer_rope_parameters)             | Return the layer type's RoPE parameters, or None for NoPE+SSMax.                         |
| [`packed_sdpa_mask`](#nemo_automodel-components-models-beam-layers-packed_sdpa_mask)                       | Construct a block-diagonal causal mask for packed SDPA.                                  |
| [`reject_flex_attention_mask`](#nemo_automodel-components-models-beam-layers-reject_flex_attention_mask)   | Reject padding on Beam's unmasked-only flex backend without a device synchronization.    |
| [`rotate_half`](#nemo_automodel-components-models-beam-layers-rotate_half)                                 | Apply HF's half-vector rotation.                                                         |
| [`ssmax_query_scale`](#nemo_automodel-components-models-beam-layers-ssmax_query_scale)                     | Compute HF Beam's NoPE query scale in FP32.                                              |
| [`window_size_for_backend`](#nemo_automodel-components-models-beam-layers-window_size_for_backend)         | Translate the visible-token window W: SDPA uses W, TE/FA use inclusive offsets (W-1, 0). |

### Data

[`SSMAX_EXPONENT`](#nemo_automodel-components-models-beam-layers-SSMAX_EXPONENT)

[`SSMAX_REFERENCE_LENGTH`](#nemo_automodel-components-models-beam-layers-SSMAX_REFERENCE_LENGTH)

### API

```python
class nemo_automodel.components.models.beam.layers.BeamAttention(
    config: transformers.models.beam.BeamConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Gated attention with per-head q/k RMSNorm and full-width output gating.

Layer type selects full/sliding attention and RoPE or NoPE+SSMax.

**`attention_dropout`**

---

**`gate_bias`**

---

**`gate_proj`**

---

**`head_dim`**

---

**`k_norm`**

---

**`k_proj`**

---

**`layer_type`** `= config.layer_types[layer_idx]`

---

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

---

**`num_kv_heads`** `= config.num_key_value_heads`

---

**`o_proj`**

---

**`q_norm`**

---

**`q_proj`**

---

**`sliding_window`**

---

**`use_rope`**

---

**`use_ssmax`** `= not self.use_rope`

---

**`v_proj`**

---

**`window_size`**

---

```python
nemo_automodel.components.models.beam.layers.BeamAttention._attend(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attention_mask: torch.Tensor | None,
    qkv_format: str,
    attn_kwargs: dict[str, typing.Any]
) -> torch.Tensor
```

Run backend attention with the correct packed and sliding-window conventions.

**Parameters:**

**`q`** `torch.Tensor`

Tensor of shape \[batch, sequence, heads, head_dim] or packed
\[tokens, heads, head_dim], local under CP and in the compute dtype.

---

**`k`** `torch.Tensor`

Tensor with q's leading axes and shape \[..., kv_heads, head_dim].

---

**`v`** `torch.Tensor`

Tensor with the same shape/dtype/device as k.

---

**`attention_mask`** `torch.Tensor | None`

Optional \[batch, sequence] valid-token mask or backend-supported
\[batch, 1, sequence, sequence] bool/additive mask.

---

**`qkv_format`** `str`

bshd for batched input or thd for packed input.

---

**`attn_kwargs`** `dict[str, Any]`

Backend metadata. Integer cu_seqlens/cu_seqlens_padded tensors
have shape \[documents + 1]; max_seqlen and CP size/rank are scalars.

---

**Returns:** `torch.Tensor`

Attention output with q's leading axes and either \[heads, head_dim]

```python
nemo_automodel.components.models.beam.layers.BeamAttention._sdpa_dropout(
    attn_kwargs: dict[str, typing.Any]
) -> None
```

Set training-only dropout on the preprocessed SDPA arguments.

**Parameters:**

**`attn_kwargs`** `dict[str, Any]`

Mutable SDPA call mapping. Any attn_mask tensor has broadcast
layout \[batch, heads, query_sequence, key_sequence]; tensor values and
storage are unchanged. Only the scalar dropout_p entry is updated.

---

```python
nemo_automodel.components.models.beam.layers.BeamAttention._set_sliding_cp_transport(
    qkv_format: str
) -> None
```

Select TE's CP exchange for this sliding-window layer from the input layout.

Automodel's CP setup gives sliding-window layers TE's `all_gather` exchange. With packed (thd) input that
exchange fails before the first step in TE 2.18 (`cudaErrorIllegalInstruction`), so packed input uses
`a2a`, which matches `all_gather` numerically; batched (bshd) input keeps `all_gather`. `a2a` splits
heads across the CP group, so the CP size must divide both head counts. Full-attention layers are not touched.

**Parameters:**

**`qkv_format`** `str`

bshd for batched input or thd for packed input.

---

```python
nemo_automodel.components.models.beam.layers.BeamAttention.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor | None,
    rope_tables: tuple[torch.Tensor, torch.Tensor] | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
```

Apply gated attention with HF-compatible RoPE or NoPE query scaling.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden] or packed \[tokens, hidden],
in the compute dtype; sequence/tokens are local under CP.

---

**`freqs_cis`** `torch.Tensor | None`

Original FP32 concatenated half-width cosine/sine tensor of shape
\[batch, sequence, head_dim] or \[tokens, head_dim] for unfused RoPE.
Fused layout follows position_ids_to_freqs_cis; None for NoPE.

---

**`rope_tables`** `tuple[torch.Tensor, torch.Tensor] | None` — default: None

Optional original FP32 (cos, sin) pair, each of shape
\[batch, sequence, 1, head_dim] or \[tokens, 1, head_dim]. Shared with
other layers without mutation. When absent, derived from freqs_cis.

---

**`position_ids`** `torch.Tensor | None` — default: None

Integer tensor of shape \[batch, sequence] or \[tokens].
Normally local under CP; full-length positions can be selected using cp_size/cp_rank.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional \[batch, sequence] valid-token mask (1/True = attend),
or backend-supported \[batch, 1, sequence, sequence] bool/additive mask.

---

**`**attn_kwargs`** `Any` — default: \{}

Backend metadata: integer cu_seqlens/cu_seqlens_padded tensors
of shape \[documents + 1], scalar max_seqlen, qkv_format, cp_size/cp_rank.

---

**Returns:** `torch.Tensor`

Tensor with the same shape as x, in the output projection's compute dtype.

```python
nemo_automodel.components.models.beam.layers.BeamAttention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

Initialize projection and normalization parameters; buffer_device follows the model initializer API.

```python
class nemo_automodel.components.models.beam.layers.BeamGate()
```

**Bases:** [Gate](/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-Gate)

Normalize top-1 routing like HF, including its epsilon and gradients.

The override also serves the CUDA-graph routing core via Gate.\_route_scores.

```python
nemo_automodel.components.models.beam.layers.BeamGate._route_scores(
    scores: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
```

Route scores with HF's normalization for top-1 as well as larger top-k.

**Parameters:**

**`scores`** `torch.Tensor`

Tensor of shape \[tokens, experts], using the shared Gate score dtype.

---

**Returns:** `torch.Tensor`

Routing weights \[tokens, topk], integer expert indices \[tokens, topk], and

```python
class nemo_automodel.components.models.beam.layers.BeamMLP(
    dim: int,
    inter_dim: int,
    backend: str,
    dtype: torch.dtype = torch.bfloat16
)
```

**Bases:** [MLP](/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-MLP)

Shared dense MLP with HF's fp32 SwiGLU boundary, preserving parameter names and MLP dispatch.

```python
nemo_automodel.components.models.beam.layers.BeamMLP.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Compute SwiGLU in FP32 before the down projection.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[..., hidden], with arbitrary leading dimensions,
in the projection compute dtype.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., hidden] with the same leading dimensions. The FP32

```python
class nemo_automodel.components.models.beam.layers.BeamMoE(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** [MoE](/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-MoE)

Shared MoE with BeamGate installed before loading or parallelization.

Shared compiled SwiGLU preserves the FP32 activation boundary; Dynamo eager
fallback can lose it. Unlike HF, default expert routing weights apply before
the down projection and expert outputs accumulate in FP32.

```python
class nemo_automodel.components.models.beam.layers.Block(
    layer_idx: int,
    config: transformers.models.beam.BeamConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

One Beam decoder layer: sandwich norms + depth-scaled residual adds.

**`attn_scales`**

---

**`input_layernorm`** `= norm()`

---

**`layer_type`** `= config.layer_types[layer_idx]`

---

**`mlp`** `= BeamMoE(moe_config, backend)`

---

**`mlp_scales`**

---

**`post_attention_layernorm`** `= norm()`

---

**`post_feedforward_layernorm`** `= norm()`

---

**`pre_feedforward_layernorm`** `= norm()`

---

**`self_attn`** `= BeamAttention(config, layer_idx, backend)`

---

```python
nemo_automodel.components.models.beam.layers.Block._mlp(
    x: torch.Tensor,
    padding_mask: torch.Tensor | None
) -> torch.Tensor
```

Dispatch to the dense MLP or the routed experts.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden] or packed \[tokens, hidden],
local to this rank under CP, in the layer compute dtype.

---

**`padding_mask`** `torch.Tensor | None`

Optional bool tensor matching x's leading axes; True skips routing.

---

**Returns:** `torch.Tensor`

Tensor with the same shape as x.

```python
nemo_automodel.components.models.beam.layers.Block.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor | None,
    rope_tables: tuple[torch.Tensor, torch.Tensor] | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
```

Apply attention, feedforward computation and depth-scaled residuals.

**Parameters:**

**`x`** `torch.Tensor`

Tensor \[batch, sequence, hidden] or packed \[tokens, hidden], local under CP;
only x is cast to the FSDP compute dtype.

---

**`freqs_cis`** `torch.Tensor | None`

FP32 rotary tensor \[..., head_dim], aligned with x's leading axes;
concatenated half-width cos/sin when unfused, fused layout otherwise. None for NoPE.

---

**`rope_tables`** `tuple[torch.Tensor, torch.Tensor] | None` — default: None

Optional FP32 (cos, sin), each \[..., 1, head_dim], shared across
layers without mutation or casting; leading axes match x.

---

**`position_ids`** `torch.Tensor | None` — default: None

Integer tensor \[batch, sequence] or \[tokens], absolute within each document.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional \[batch, sequence] bool/0-1 mask; True/1 means valid.

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional bool tensor matching x's leading axes; True skips routing.

---

**`**attn_kwargs`** `Any` — default: \{}

Backend metadata: qkv_format (bshd/thd), integer cu_seqlens and
cu_seqlens_padded \[documents + 1], scalar max_seqlen, cp_size and cp_rank.

---

**Returns:** `torch.Tensor`

Tensor with x's shape, covering this rank's local positions.

```python
nemo_automodel.components.models.beam.layers.Block.init_weights(
    buffer_device: torch.device
) -> None
```

Initialize the layer's norms, attention, MLP and router buffers on buffer_device.

```python
nemo_automodel.components.models.beam.layers._cp_local_positions(
    position_ids: torch.Tensor,
    local_len: int,
    cp_size: int,
    cp_rank: int
) -> torch.Tensor
```

Select the two load-balanced sequence chunks owned by this CP rank.

**Parameters:**

**`position_ids`** `torch.Tensor`

Integer tensor of shape \[batch, global_sequence] or
\[global_sequence], or the corresponding already-local shape.

---

**`local_len`** `int`

Tokens per local sequence, equal to global_sequence / cp_size;
must be even when selecting the two head/tail chunks.

---

**`cp_size`** `int`

Number of context-parallel ranks.

---

**`cp_rank`** `int`

Zero-based context-parallel rank.

---

**Returns:** `torch.Tensor`

Position tensor of shape \[batch, local_len] or \[local_len]. Already-local

```python
nemo_automodel.components.models.beam.layers.apply_rope_hf(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor
```

Apply HF Beam's rotary expression and preserve its backward rounding.

Do not eagerly cast x to FP32: HF rounds each product's gradient to x.dtype.

**Parameters:**

**`x`** `torch.Tensor`

Query/key tensor of shape \[batch, sequence, heads, head_dim] or packed
\[tokens, heads, head_dim], in the activation dtype.

---

**`cos`** `torch.Tensor`

Original FP32 cosine tensor of shape \[batch, sequence, 1, head_dim]
or \[tokens, 1, head_dim], broadcastable to x and on the same device.

---

**`sin`** `torch.Tensor`

Original FP32 sine tensor with the same shape/device as cos.

---

**Returns:** `torch.Tensor`

New tensor with x's shape/dtype/device, computed as

```python
nemo_automodel.components.models.beam.layers.depth_scaled_residual_coefs(
    layer_idx: int,
    is_attention: bool
) -> tuple[float, float]
```

(sublayer_scale, residual_scale) for the depth-scaled residual add.

1-based residual-add index: `2L` for attention, `2L + 1` for the MLP; the sublayer output is
scaled by `1/sqrt(idx)` and the residual stream by `sqrt((idx - 1)/idx)`.

```python
nemo_automodel.components.models.beam.layers.hf_rope_tables(
    freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Expand original FP32 rotary tables into HF's duplicated-half layout.

**Parameters:**

**`freqs_cis`** `torch.Tensor`

FP32 tensor of shape \[..., head_dim], with arbitrary leading axes
corresponding to \[batch, sequence] or packed \[tokens]. The last axis
concatenates half-width cosine and sine values, in that order.

---

**Returns:** `torch.Tensor`

Pair (cos, sin), each a new FP32 tensor of shape \[..., 1, head_dim] on the

```python
nemo_automodel.components.models.beam.layers.is_moe_layer(
    config: transformers.models.beam.BeamConfig,
    layer_idx: int
) -> bool
```

Beam marks every layer `"dense"` or `"sparse"` in `config.mlp_layer_types`.

```python
nemo_automodel.components.models.beam.layers.layer_rope_parameters(
    config: transformers.models.beam.BeamConfig,
    layer_type: str
) -> dict[str, typing.Any] | None
```

Return the layer type's RoPE parameters, or None for NoPE+SSMax.

Reject flat pre-Beam configs and non-default rope_type rather than silently changing positional encoding.

```python
nemo_automodel.components.models.beam.layers.packed_sdpa_mask(
    cu_seqlens: torch.Tensor,
    total_tokens: int,
    sliding_window: int | None
) -> torch.Tensor
```

Construct a block-diagonal causal mask for packed SDPA.

**Parameters:**

**`cu_seqlens`** `torch.Tensor`

Integer tensor of shape \[documents + 1], beginning at zero and
ending at total_tokens, with strictly increasing document boundaries.

---

**`total_tokens`** `int`

Number of tokens in the concatenated packed buffer.

---

**`sliding_window`** `int | None`

Number of visible tokens, including the query, or None.

---

**Returns:** `torch.Tensor`

Bool tensor of shape \[1, 1, tokens, tokens] on cu_seqlens.device, where

```python
nemo_automodel.components.models.beam.layers.reject_flex_attention_mask(
    attention_mask: torch.Tensor | None
) -> None
```

Reject padding on Beam's unmasked-only flex backend without a device synchronization.

**Parameters:**

**`attention_mask`** `torch.Tensor | None`

Optional tensor, normally \[batch, sequence]; any non-None value
is rejected without inspecting its contents. Use TE or SDPA for padded input.

---

```python
nemo_automodel.components.models.beam.layers.rotate_half(
    x: torch.Tensor
) -> torch.Tensor
```

Apply HF's half-vector rotation.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[..., head_dim], with arbitrary leading axes and even head_dim.

---

**Returns:** `torch.Tensor`

New tensor of the same shape/dtype/device, concatenating -x's second half

```python
nemo_automodel.components.models.beam.layers.ssmax_query_scale(
    position_ids: torch.Tensor
) -> torch.Tensor
```

Compute HF Beam's NoPE query scale in FP32.

**Parameters:**

**`position_ids`** `torch.Tensor`

Integer tensor of shape \[batch, sequence] or packed \[tokens],
containing absolute per-document positions, local to this rank under CP.

---

**Returns:** `torch.Tensor`

FP32 tensor with the same shape/device as position_ids:

```python
nemo_automodel.components.models.beam.layers.window_size_for_backend(
    sliding_window: int | None,
    attn_impl: str
) -> tuple[int, int] | None
```

Translate the visible-token window W: SDPA uses W, TE/FA use inclusive offsets (W-1, 0).

Flex takes sliding_window=W directly.

```python
nemo_automodel.components.models.beam.layers.SSMAX_EXPONENT = 0.1
```

```python
nemo_automodel.components.models.beam.layers.SSMAX_REFERENCE_LENGTH = 2048
```