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`	Rotary Position Embedding module for Llama and Qwen2 models.

Functions

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

Data

Qwen2RotaryEmbedding

RotaryEmbedding

__all__

API

class nemo_automodel.components.models.llama.rope_utils.LlamaRotaryEmbedding(
    config,
    device: typing.Optional[torch.device] = 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

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).

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).

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:

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

strDefaults to 'bshd'

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

cp_size

intDefaults to 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

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

Computes inverse frequencies for standard RoPE.

nemo_automodel.components.models.llama.rope_utils._compute_llama3_inv_freq(
    config,
    device: typing.Optional[torch.device] = 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

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)

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:

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.

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

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:

torch.Tensor

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

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 | NoneDefaults to None

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

cp_size

intDefaults to 1

Number of context-parallel ranks partitioning T.

cp_rank

intDefaults to 0

Rank within the context-parallel mesh.

Returns: torch.Tensor

Rotated tensors with the same local shapes and storage ownership as

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.

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

Rotates half the hidden dims of the input.

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

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

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