> This page is for version 0.7.0 · 26.10.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.gpt_oss.rope_utils

## Module Contents

### Classes

| Name                                                                                      | Description |
| ----------------------------------------------------------------------------------------- | ----------- |
| [`RotaryEmbedding`](#nemo_automodel-components-models-gpt_oss-rope_utils-RotaryEmbedding) | -           |

### Functions

| Name                                                                                                          | Description                                                    |
| ------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------- |
| [`_apply_rotary_emb_qk`](#nemo_automodel-components-models-gpt_oss-rope_utils-_apply_rotary_emb_qk)           | Rotate query/key together, sharing their dtype-rounded tables. |
| [`apply_rotary_emb`](#nemo_automodel-components-models-gpt_oss-rope_utils-apply_rotary_emb)                   | Apply rotary embeddings to input tensor.                       |
| [`apply_rotary_emb_qk`](#nemo_automodel-components-models-gpt_oss-rope_utils-apply_rotary_emb_qk)             | Apply rotary embeddings to query and key tensors.              |
| [`position_ids_to_freqs_cis`](#nemo_automodel-components-models-gpt_oss-rope_utils-position_ids_to_freqs_cis) | -                                                              |

### Data

[`_compiled_apply_rotary_emb_qk`](#nemo_automodel-components-models-gpt_oss-rope_utils-_compiled_apply_rotary_emb_qk)

### API

```python
class nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding(
    head_dim: int,
    base: int,
    dtype: torch.dtype,
    initial_context_length: int = 4096,
    scaling_factor: float = 1.0,
    ntk_alpha: float = 1.0,
    ntk_beta: float = 32.0,
    partial_rotary_factor: float = 1.0,
    device: torch.device | None = None
)
```

**Bases:** `Module`

**`rotary_dim`** `= int(head_dim * partial_rotary_factor)`

---

```python
nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding._compute_concentration_and_inv_freq() -> torch.Tensor
```

See YaRN paper: [https://arxiv.org/abs/2309.00071](https://arxiv.org/abs/2309.00071)

Uses rotary_dim instead of head_dim to support partial rotary embeddings.

```python
nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding._compute_cos_sin(
    num_tokens: int
)
```

```python
nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding.forward(
    query: torch.Tensor,
    key: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.gpt_oss.rope_utils._apply_rotary_emb_qk(
    q: torch.Tensor,
    k: torch.Tensor,
    freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Rotate query/key together, sharing their dtype-rounded tables.

**Parameters:**

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

Tensor of shape \[..., query_heads, head_dim], with arbitrary leading
dimensions and strides.

---

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

Tensor of shape \[..., key_heads, head_dim], sharing the leading
dimensions of `q` and accepting arbitrary strides.

---

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

Tensor of shape \[..., rotary_dim], with leading dimensions
broadcastable to `q` and `k`. The last dimension concatenates
cosine and sine tables; `rotary_dim` is even and at most
`head_dim`.

---

**Returns:** `torch.Tensor`

Independently allocated rotated query and key tensors, preserving

```python
nemo_automodel.components.models.gpt_oss.rope_utils.apply_rotary_emb(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor
```

Apply rotary embeddings to input tensor.

If cos/sin have fewer dimensions than x (due to partial_rotary_factor \< 1.0),
only the first rotary_dim dimensions of x are rotated, and the rest are passed through.

**Parameters:**

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

Input tensor (..., head_dim)

---

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

Cosine tensor (..., rotary_dim // 2)

---

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

Sine tensor (..., rotary_dim // 2)

---

```python
nemo_automodel.components.models.gpt_oss.rope_utils.apply_rotary_emb_qk(
    q: torch.Tensor,
    k: torch.Tensor,
    freqs_cis: torch.Tensor,
    format: str = 'bshd',
    rope_fusion: bool = True,
    cu_seqlens: torch.Tensor | None = None,
    concentration: float | None = None,
    cp_size: int = 1,
    cp_rank: int = 0
) -> tuple[torch.Tensor, torch.Tensor]
```

Apply rotary embeddings to query and key tensors.

**Parameters:**

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

Query tensor of shape \[batch, sequence, query_heads, head_dim] for
BSHD or \[tokens, query_heads, head_dim] for packed THD. Leading
dimensions contain the local context-parallel shard, if used.

---

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

Key tensor with the same layout as `q`, with `key_heads` in
place of `query_heads`. Arbitrary strides are supported.

---

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

For non-fused RoPE, tensor of shape \[batch, sequence,
rotary_dim] or \[tokens, rotary_dim], containing concatenated
cosine/sine tables with concentration applied. Leading dimensions
broadcast to the local tokens in `q` and `k`. For TE fused
RoPE, tensor of shape \[sequence, 1, 1, rotary_dim] containing
duplicated angles for the global sequence. `rotary_dim` is even
and at most `head_dim`.

---

**`format`** `str` — default: 'bshd'

QKV format ("bshd" or "thd").

---

**`rope_fusion`** `bool` — default: True

If True, use TE fused rope. If False, use non-fused rope.

---

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

Optional tensor of shape \[num_sequences + 1] containing
cumulative sequence lengths for variable-length TE fused RoPE.

---

**`concentration`** `float | None` — default: None

Optional scale applied after TE fused RoPE.

---

**`cp_size`** `int` — default: 1

Context parallelism size.

---

**`cp_rank`** `int` — default: 0

Context parallelism rank.

---

**Returns:** `torch.Tensor`

Tuple of rotated query and key tensors, preserving their respective

```python
nemo_automodel.components.models.gpt_oss.rope_utils.position_ids_to_freqs_cis(
    rotary_emb: nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding,
    position_ids: torch.Tensor,
    qkv_format: str = 'bshd',
    for_fused_rope: bool = True,
    cp_size: int = 1
) -> torch.Tensor
```

```python
nemo_automodel.components.models.gpt_oss.rope_utils._compiled_apply_rotary_emb_qk = torch.compile(_apply_rotary_emb_qk, dynamic=True, options={'emulate_precision_ca...
```