nemo_automodel.components.models.deepseek_v3.rope_utils

Module Contents

Functions

Name	Description
`apply_rotary_emb`	Rotate adjacent channel pairs using precomputed complex frequencies.
`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

Rotate adjacent channel pairs using precomputed complex frequencies.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, heads, channels] for bshd or [tokens, heads, channels] for thd. Shared keys may omit the head dimension when unsqueeze_dim inserts it.

freqs_cis

torch.Tensor

Complex frequency table with channels / 2 pairs per token, shaped [batch, sequence, channels / 2] or [tokens, channels / 2]. Cached CUDA complex64 frequencies use the fused kernel; trainable frequencies retain autograd.

qkv_format

strDefaults to 'bshd'

thd adds a temporary batch dimension for packed inputs.

unsqueeze_dim

int | NoneDefaults to None

Optional head dimension to insert after normalizing packed inputs to include a batch dimension.

Returns: torch.Tensor

Tensor with the input shape and dtype. Rotation uses FP32 complex

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
)