nemo_automodel.components.models.deepseek_v41.rope_triton

Fused Triton rotary embedding for DeepSeek V4.1, forward and backward in one kernel each.

_apply_rope in nemo_automodel.components.models.deepseek_v41.attention rotates the last 2 * rotary_pairs channels of a [batch, sequence, channels] or [batch, sequence, heads, channels] tensor as adjacent complex pairs and copies the remaining channels. Eagerly that is a complex multiply plus torch.cat in the forward and several zero-filled slice gradients in the backward, so every call streams the full tensor through memory more than once. This module writes the copied and the rotated channels in a single pass, and its backward rotates the incoming gradient by the conjugate in one pass as well. angles never receives a gradient here; inputs whose angles require one, inputs that are not on CUDA, and dtypes outside BF16/FP16/FP32 take the eager path.

Numerics follow the eager path by construction wherever PyTorch allows it:

• The frequencies come from the same torch.polar(torch.ones_like(angles), angles) call the eager path runs, so the cos/sin bits are identical rather than relying on libdevice agreeing with the CUDA math library PyTorch was built with. They are computed once per forward call and saved for the backward, which is the same footprint as the eager path’s saved frequencies.
• The complex product (x0 + i x1) * (c + i s) is evaluated in FP32 as r0 = fma(x0, c, -(x1 * s)) and r1 = fma(x0, s, x1 * c). c10::complex<float> multiplies as a * c - b * d and a * d + b * c, and nvcc contracts each expression into one FMA with the first product fused; the kernel spells that contraction out so the result does not depend on Triton’s FP-fusion default. The conjugate used by inverse=True and by every backward only negates s, which is exact and leaves the rounding pattern unchanged. If a GPU parity run ever reports one-ulp differences, the two tl.fma lines in _rope_kernel are the single place to adjust.
• The FP32 result is rounded to values.dtype with round-to-nearest-even, exactly as Tensor.to rounds.

Module Contents

Classes

Functions

Data

HAVE_TRITON

_ELEMENTS_PER_PROGRAM

_FUSED_DTYPES

API

class nemo_automodel.components.models.deepseek_v41.rope_triton._FusedRope()

nemo_automodel.components.models.deepseek_v41.rope_triton._FusedRope.backward(
    ctx: object,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, None, None]

nemo_automodel.components.models.deepseek_v41.rope_triton._FusedRope.forward(
    ctx: object,
    values: torch.Tensor,
    frequencies: torch.Tensor,
    inverse: bool
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v41.rope_triton._apply_rope_torch(
    values: torch.Tensor,
    angles: torch.Tensor,
    inverse: bool = False
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v41.rope_triton._launch(
    values: torch.Tensor,
    frequencies: torch.Tensor,
    negate_sin: bool
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v41.rope_triton._rope_kernel(
    values_ptr,
    frequencies_ptr,
    output_ptr,
    rows,
    sequence,
    heads,
    channels,
    pairs,
    pair_offset,
    stride_batch,
    stride_sequence,
    stride_head,
    NEGATE_SIN: tl.constexpr,
    BLOCK_R: tl.constexpr,
    BLOCK_K: tl.constexpr
)

nemo_automodel.components.models.deepseek_v41.rope_triton._uses_fused_kernel(
    values: torch.Tensor,
    angles: torch.Tensor
) -> bool

nemo_automodel.components.models.deepseek_v41.rope_triton.apply_rope_triton(
    values: torch.Tensor,
    angles: torch.Tensor,
    inverse: bool = False
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v41.rope_triton.HAVE_TRITON: bool = bool(_HAVE_TRITON and _HAVE_TRITON_LANGUAGE)

nemo_automodel.components.models.deepseek_v41.rope_triton._ELEMENTS_PER_PROGRAM = 4096

nemo_automodel.components.models.deepseek_v41.rope_triton._FUSED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)