> 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.llama.rope_utils

Rotary Position Embedding utilities for Llama and Qwen2 models.

This module provides RoPE implementation following HuggingFace's architecture.

Supports both:

* LlamaConfig: uses config.rope\_theta and config.rope\_scaling
* Qwen2Config: uses config.rope\_parameters\["rope\_theta"] and config.rope\_parameters

Note: gpt\_oss and deepseek\_v3 have their own specialized rope\_utils.py
with model-specific optimizations (YaRN, MLA, etc.).

## Module Contents

### Classes

| Name                                                                                              | Description                                                  |
| ------------------------------------------------------------------------------------------------- | ------------------------------------------------------------ |
| [`LlamaRotaryEmbedding`](#nemo_automodel-components-models-llama-rope_utils-LlamaRotaryEmbedding) | Rotary Position Embedding module for Llama and Qwen2 models. |

### Functions

| Name                                                                                                          | Description                                                                   |
| ------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`_compute_default_inv_freq`](#nemo_automodel-components-models-llama-rope_utils-_compute_default_inv_freq)   | Computes inverse frequencies for standard RoPE.                               |
| [`_compute_llama3_inv_freq`](#nemo_automodel-components-models-llama-rope_utils-_compute_llama3_inv_freq)     | Computes inverse frequencies for Llama3-style RoPE with smooth interpolation. |
| [`_get_rope_config`](#nemo_automodel-components-models-llama-rope_utils-_get_rope_config)                     | Extract rope parameters from config (handles both Llama and Qwen2 formats).   |
| [`apply_rotary_pos_emb`](#nemo_automodel-components-models-llama-rope_utils-apply_rotary_pos_emb)             | Applies Rotary Position Embedding to the query and key tensors.               |
| [`apply_rotary_pos_emb_fused`](#nemo_automodel-components-models-llama-rope_utils-apply_rotary_pos_emb_fused) | Applies RoPE using TE's fused kernel.                                         |
| [`apply_rotary_pos_emb_quack`](#nemo_automodel-components-models-llama-rope_utils-apply_rotary_pos_emb_quack) | Apply Llama-style RoPE with QuACK's fused rotary kernel.                      |
| [`rotate_half`](#nemo_automodel-components-models-llama-rope_utils-rotate_half)                               | Rotates half the hidden dims of the input.                                    |

### Data

[`Qwen2RotaryEmbedding`](#nemo_automodel-components-models-llama-rope_utils-Qwen2RotaryEmbedding)

[`RotaryEmbedding`](#nemo_automodel-components-models-llama-rope_utils-RotaryEmbedding)

[`__all__`](#nemo_automodel-components-models-llama-rope_utils-__all__)

### API

```python
class nemo_automodel.components.models.llama.rope_utils.LlamaRotaryEmbedding(
    config,
    device: torch.device | None = None,
    rope_fusion: bool = False
)
```

**Bases:** `Module`

Rotary Position Embedding module for Llama and Qwen2 models.

Returns (cos, sin) tuple for use with apply\_rotary\_pos\_emb.

**`dtype`**

---

**`inv_freq`** `Tensor`

---

**`max_seq_len_cached`** `= 0`

---

```python
nemo_automodel.components.models.llama.rope_utils.LlamaRotaryEmbedding._build_cache(
    seq_len: int,
    device: torch.device
) -> None
```

Build cos/sin cache in config dtype for positions \[0, seq\_len).

```python
nemo_automodel.components.models.llama.rope_utils.LlamaRotaryEmbedding._ensure_cache(
    seq_len: int,
    device: torch.device
) -> None
```

Build or grow the cos/sin cache so it covers positions `[0, seq_len)`.

```python
nemo_automodel.components.models.llama.rope_utils.LlamaRotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor,
    qkv_format: str = 'bshd',
    cp_size: int = 1
) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Return (cos, sin) for the given positions.

In the non-fused path `cos` / `sin` are gathered by the *values* in
`position_ids`, so non-contiguous positions receive the correct rotary
phase: EAGLE TTT depth offsets (`arange(seq_len) + step_idx`), packed
sequences, and context parallelism all pass
`position_ids != arange(seq_len)`. The earlier implementation returned
`cos_cache[:seq_len]`, which keyed only on the sequence *length* and
silently ignored the position values. For the common
`position_ids == arange(seq_len)` case the gather is numerically
identical to that slice.

The fused TE path (`rope_fusion=True`) consumes raw angles indexed by
sequence position and assumes contiguous `[0, seq_len)` positions, so
it keeps the legacy contiguous slice and does NOT honor non-contiguous
`position_ids`.

**Parameters:**

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

Hidden states `[B, S, H]` or packed `[T, H]`; used for
device placement. `H` is hidden size.

---

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

Position IDs `[B, S]` or packed local IDs `[T]`.

---

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

`"bshd"` for padded batches or `"thd"` for a
packed total-token layout.

---

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

Context-parallel size. For fused THD RoPE, the cached raw
frequency table spans the global token count `T * cp_size`.

---

**Returns:** `tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]`

Cosine and sine tensors `[B, S, D]` or `[T, D]`. When fused

```python
nemo_automodel.components.models.llama.rope_utils._compute_default_inv_freq(
    config,
    device: torch.device | None = None
) -> tuple[torch.Tensor, float]
```

Computes inverse frequencies for standard RoPE.

```python
nemo_automodel.components.models.llama.rope_utils._compute_llama3_inv_freq(
    config,
    device: torch.device | None = None
) -> tuple[torch.Tensor, float]
```

Computes inverse frequencies for Llama3-style RoPE with smooth interpolation.

Branch logic (matches HF \_compute\_llama3\_parameters):

* Long wavelength  (low freq,  wavelen > low\_freq\_wavelen)  → scale by factor
* Short wavelength (high freq, wavelen \< high\_freq\_wavelen) → unchanged
* Medium band → smooth interpolation

```python
nemo_automodel.components.models.llama.rope_utils._get_rope_config(
    config
) -> tuple[float, dict]
```

Extract rope parameters from config (handles both Llama and Qwen2 formats).

**Returns:** `tuple[float, dict]`

Tuple of (rope\_theta, rope\_scaling\_dict)

```python
nemo_automodel.components.models.llama.rope_utils.apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Applies Rotary Position Embedding to the query and key tensors.

**Parameters:**

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

Query tensor `[B, Hq, S, D]` in BSHD attention layout or
`[T, Hq, D]` in packed THD layout. `B` is batch, `S` is
sequence, `T` is total local tokens, `Hq` is query heads, and
`D` is head dimension.

---

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

Key tensor `[B, Hkv, S, D]` or `[T, Hkv, D]`, where `Hkv` is
key/value heads.

---

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

Cosine embeddings `[B, S, D]` or `[T, D]`.

---

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

Sine embeddings `[B, S, D]` or `[T, D]`.

---

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

Rotated (q, k) tensors

```python
nemo_automodel.components.models.llama.rope_utils.apply_rotary_pos_emb_fused(
    q: torch.Tensor,
    k: torch.Tensor,
    freqs_cis: torch.Tensor,
    cu_seqlens: torch.Tensor | None = None,
    cp_size: int = 1,
    cp_rank: int = 0
) -> tuple[torch.Tensor, torch.Tensor]
```

Applies RoPE using TE's fused kernel.

**Parameters:**

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

Query tensor `[B, Hq, S, D]` or packed `[T, Hq, D]`.

---

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

Key tensor `[B, Hkv, S, D]` or packed `[T, Hkv, D]`.

---

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

Global raw-angle table `[S, 1, 1, D]` in TE format.

---

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

Packed-document cumulative lengths `[N + 1]`. Required
for THD, where `N` is the number of packed documents.

---

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

Number of context-parallel ranks partitioning `T`.

---

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

Rank within the context-parallel mesh.

---

**Returns:** `torch.Tensor`

Rotated tensors with the same local shapes and storage ownership as

```python
nemo_automodel.components.models.llama.rope_utils.apply_rotary_pos_emb_quack(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    quack_apply_rotary_emb
) -> tuple[torch.Tensor, torch.Tensor]
```

Apply Llama-style RoPE with QuACK's fused rotary kernel.

QuACK consumes BSHD tensors and half-width cosine/sine tables, whereas the
HuggingFace contract uses BHSD tensors and duplicates the table across the
full head dimension. Batch size one uses QuACK's in-place path. Larger
batches are rotated independently so batch-specific position IDs remain
correct rather than assuming every sequence uses the same positions.

```python
nemo_automodel.components.models.llama.rope_utils.rotate_half(
    x: torch.Tensor
) -> torch.Tensor
```

Rotates half the hidden dims of the input.

```python
nemo_automodel.components.models.llama.rope_utils.Qwen2RotaryEmbedding = LlamaRotaryEmbedding
```

```python
nemo_automodel.components.models.llama.rope_utils.RotaryEmbedding = LlamaRotaryEmbedding
```

```python
nemo_automodel.components.models.llama.rope_utils.__all__ = ['RotaryEmbedding', 'LlamaRotaryEmbedding', 'Qwen2RotaryEmbedding', 'rotate_half...
```