Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsCommonnemo_automodel.components.models.common.rope_triton

nemo_automodel.components.models.common.rope_triton

View as Markdown

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 is evaluated in FP32 with explicit FMAs. Eager CUDA complex multiplication contracts the imaginary component differently across PyTorch/CUDA builds and devices. A small eager multiplication selects the matching order once per device; only that choice is cached for the process lifetime. The same order is used for the conjugate in the backward. Unrecognized arithmetic defaults to the first-product order.
  • The FP32 result is rounded to values.dtype with round-to-nearest-even, exactly as Tensor.to rounds.

Module Contents

Classes

NameDescription
_FusedRopeRotate with the fused kernel and rotate the gradient by the conjugate frequency.

Functions

NameDescription
_apply_rope_torchEager rotation, mirroring attention._apply_rope so the fallback needs no circular import.
_imaginary_fma_uses_second_productIdentify eager complex multiplication’s contraction for this CUDA device.
_launchRun the fused kernel over values and return a new contiguous tensor.
_rope_kernelCopy the leading channels and rotate the trailing pairs of BLOCK_R rows.
_uses_fused_kernelCheck the shared kernel’s device, layout, dtype, and autograd contract.
apply_complex_ropeRotate trailing channel pairs using a model’s cached complex frequencies.
apply_rope_tritonRotate the final channels of values with the fused Triton kernel.

Data

HAVE_TRITON

_ELEMENTS_PER_PROGRAM

_FUSED_DTYPES

API

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

Bases: Function

Rotate with the fused kernel and rotate the gradient by the conjugate frequency.

nemo_automodel.components.models.common.rope_triton._FusedRope.backward(
    ctx: object,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, None, None]
staticmethod
nemo_automodel.components.models.common.rope_triton._FusedRope.forward(
    ctx: object,
    values: torch.Tensor,
    frequencies: torch.Tensor,
    inverse: bool
) -> torch.Tensor
staticmethod
nemo_automodel.components.models.common.rope_triton._apply_rope_torch(
    values: torch.Tensor,
    frequencies: torch.Tensor,
    inverse: bool = False
) -> torch.Tensor

Eager rotation, mirroring attention._apply_rope so the fallback needs no circular import.

Parameters:

values
torch.Tensor

Tensor of shape [batch, sequence, channels] or [batch, sequence, heads, channels].

frequencies
torch.Tensor

Complex tensor of shape [batch, sequence, rotary_pairs].

inverse
boolDefaults to False

Conjugate the rotation.

Returns: torch.Tensor

Tensor with the shape and dtype of values, in independent storage.

nemo_automodel.components.models.common.rope_triton._imaginary_fma_uses_second_product(
    device: torch.device
) -> bool

Identify eager complex multiplication’s contraction for this CUDA device.

The bounded, process-local cache retains only the arithmetic choice, not tensors. Calibration synchronizes on first use, before the first Triton kernel compilation. FP64 CPU products distinguish the two FP32 FMA orders without using a CUDA FMA implementation as the reference or changing the random-number generator state.

Parameters:

device
torch.device

CUDA device on which the eager reference will execute.

Returns: bool

Whether the eager reference matches the second-product contraction exactly.

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

Run the fused kernel over values and return a new contiguous tensor.

Parameters:

values
torch.Tensor

CUDA tensor of shape [batch, sequence, channels] or [batch, sequence, heads, channels]. Any leading strides are read in place; a strided channel dimension is made contiguous first.

frequencies
torch.Tensor

Contiguous FP32 tensor of shape [batch, sequence, pairs, 2].

negate_sin
bool

Rotate by the conjugate frequency.

Returns: torch.Tensor

Contiguous tensor with the shape and dtype of values.

nemo_automodel.components.models.common.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,
    IMAGINARY_SECOND_PRODUCT: tl.constexpr,
    BLOCK_R: tl.constexpr,
    BLOCK_K: tl.constexpr
)

Copy the leading channels and rotate the trailing pairs of BLOCK_R rows.

Every row is read once as BLOCK_K adjacent channel pairs and written once. Pairs before the rotary block are passed through tl.where rather than multiplied by a unit frequency, so non-finite prefix values are copied bit for bit.

Parameters:

values_ptr

Input viewed as [batch, sequence, heads, channels] with a unit channel stride; heads is 1 for three-dimensional inputs.

frequencies_ptr

Contiguous FP32 tensor of shape [batch * sequence, pairs, 2] holding (cos, sin) per rotary pair.

output_ptr

Contiguous output of shape [rows, channels].

rows

batch * sequence * heads.

sequence

Sequence length.

heads

Head count, 1 for three-dimensional inputs.

channels

Channel count.

pairs

Rotary pair count; the last 2 * pairs channels rotate.

pair_offset

Channel pairs preceding the rotary block, rounded up so the pair grid covers channel 0 (it then starts at channel -1 for an odd prefix, masked off).

stride_batch

Input batch stride in elements.

stride_sequence

Input sequence stride in elements.

stride_head

Input head stride in elements; ignored when heads is 1.

NEGATE_SIN
tl.constexpr

Rotate by the conjugate frequency.

IMAGINARY_SECOND_PRODUCT
tl.constexpr

Fuse x1 * cos in the imaginary component to match the installed eager CUDA complex multiplication.

BLOCK_R
tl.constexpr

Rows handled by one program.

BLOCK_K
tl.constexpr

Power-of-two pair extent per row, at least pair_offset + pairs.

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

Check the shared kernel’s device, layout, dtype, and autograd contract.

Parameters:

values
torch.Tensor

Tensor of shape [batch, sequence, channels] or [batch, sequence, heads, channels].

frequencies
torch.Tensor

Complex tensor of shape [batch, sequence, rotary_pairs].

Returns: bool

Whether the inputs can use the Triton implementation.

nemo_automodel.components.models.common.rope_triton.apply_complex_rope(
    values: torch.Tensor,
    frequencies: torch.Tensor,
    inverse: bool = False
) -> torch.Tensor

Rotate trailing channel pairs using a model’s cached complex frequencies.

Parameters:

values
torch.Tensor

Tensor of shape [batch, sequence, channels] or [batch, sequence, heads, channels]. Leading dimensions may be strided.

frequencies
torch.Tensor

Complex tensor of shape [batch, sequence, rotary_pairs]. The rotary width must fit within channels. CUDA BF16/FP16/FP32 values and cached complex64 frequencies use Triton; other inputs use eager operations, preserving frequency gradients when required.

inverse
boolDefaults to False

Whether to conjugate the rotation.

Returns: torch.Tensor

New contiguous tensor with the shape and dtype of values. Its input

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

Rotate the final channels of values with the fused Triton kernel.

Drop-in for attention._apply_rope: the last 2 * angles.shape[-1] channels are rotated as adjacent pairs by exp(i * angles) (its conjugate when inverse) and the remaining channels are copied, all in one pass. The backward rotates the gradient by the conjugate in one pass; angles receives no gradient. Inputs that are not on CUDA, whose angles require a gradient, or whose dtype is outside BF16/FP16/FP32 use the eager path.

Parameters:

values
torch.Tensor

Tensor of shape [batch, sequence, channels] or [batch, sequence, heads, channels]; non-contiguous inputs are accepted.

angles
torch.Tensor

FP32 tensor of shape [batch, sequence, rotary_pairs].

inverse
boolDefaults to False

Conjugate the rotation, as for the attention output.

Returns: torch.Tensor

New contiguous tensor with the shape and dtype of values.

Raises:

  • ValueError: angles does not describe an even rotary width within values.
nemo_automodel.components.models.common.rope_triton.HAVE_TRITON: bool = bool(_HAVE_TRITON and _HAVE_TRITON_LANGUAGE)
 
nemo_automodel.components.models.common.rope_triton._ELEMENTS_PER_PROGRAM = 4096
 
nemo_automodel.components.models.common.rope_triton._FUSED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)