nemo_automodel.components.models.deepseek_v3.rope_utils

Module Contents

Functions

Name	Description
`apply_rotary_emb`	Applies rotary positional embeddings to the input tensor.
`apply_rotary_emb_half_split`	Apply half-split (`rotate_half`) rotary embeddings using the complex `freqs_cis` table.
`apply_rotary_emb_qk`	-
`freqs_cis_from_position_ids`	-
`precompute_freqs_cis`	Precomputes frequency-based complex exponential values for rotary positional embeddings.
`yarn_get_mscale`	-

API

nemo_automodel.components.models.deepseek_v3.rope_utils.apply_rotary_emb(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    qkv_format: str = 'bshd',
    unsqueeze_dim: int | None = None
) -> torch.Tensor

Applies rotary positional embeddings to the input tensor.

Parameters:

torch.Tensor

Input tensor with positional embeddings to be applied.

freqs_cis

torch.Tensor

Precomputed complex exponential values for positional embeddings.

Returns: torch.Tensor

torch.Tensor: Tensor with rotary embeddings applied.

nemo_automodel.components.models.deepseek_v3.rope_utils.apply_rotary_emb_half_split(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    qkv_format: str = 'bshd'
) -> torch.Tensor

Apply half-split (rotate_half) rotary embeddings using the complex freqs_cis table.

The DeepSeek-V3.2 DSA indexer rotates its rope slice with the non-interleaved layout (pair dim j with j + d/2), unlike the main MLA attention which pairs adjacent dims (apply_rotary_emb). freqs_cis is the same complex table used by the MLA (exp(i * theta_j * pos) for j in [0, d/2)); its real/imag parts are the cos/sin for frequency j, so the rotation angles match the interleaved path exactly. Compute the rotation in FP32, as in the official DeepSeek implementation and vLLM’s CUDA indexer RoPE, then cast the result back to the input dtype.

Parameters:

torch.Tensor

Rope slice, [B, S, H, d] / [B, S, d] (bshd) or [T, H, d] / [T, d] (thd).

freqs_cis

torch.Tensor

Complex RoPE table with trailing dim d/2 ([B, S, d/2] or [T, d/2]).

qkv_format

strDefaults to 'bshd'

"bshd" or "thd".

Returns: torch.Tensor

Rotated tensor with the same shape and dtype as x.

nemo_automodel.components.models.deepseek_v3.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,
    cp_size: int = 1,
    cp_rank: int = 0
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.models.deepseek_v3.rope_utils.freqs_cis_from_position_ids(
    position_ids: torch.Tensor,
    freqs: torch.Tensor,
    qkv_format: str = 'bshd',
    for_fused_rope: bool = False,
    cp_size: int = 1
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v3.rope_utils.precompute_freqs_cis(
    qk_rope_head_dim: int,
    max_seq_len: int,
    rope_theta: float,
    rope_scaling: dict[str, float | int] | None
) -> torch.Tensor

Precomputes frequency-based complex exponential values for rotary positional embeddings.

Parameters:

qk_rope_head_dim

int

Dimensionality of the rotary positional embeddings.

max_seq_len

int

Maximum sequence length.

original_seq_len

int

Original sequence length.

beta_fast

int

Fast beta value for the exponential computation.

beta_slow

int

Slow beta value for the exponential computation.

rope_theta

float

Base value for the exponential computation.

rope_factor

float

Factor value for the exponential computation.

Returns: torch.Tensor

torch.Tensor: Precomputed complex exponential values for positional embeddings.

nemo_automodel.components.models.deepseek_v3.rope_utils.yarn_get_mscale(
    scale = 1,
    mscale = 1
)