> This page is for version 0.4.0 · 26.04.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.common.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 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

| Name                                                                            | Description                                                                      |
| ------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- |
| [`_FusedRope`](#nemo_automodel-components-models-common-rope_triton-_FusedRope) | Rotate with the fused kernel and rotate the gradient by the conjugate frequency. |

### Functions

| Name                                                                                                                            | Description                                                                                 |
| ------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- |
| [`_apply_rope_torch`](#nemo_automodel-components-models-common-rope_triton-_apply_rope_torch)                                   | Eager rotation, mirroring `attention._apply_rope` so the fallback needs no circular import. |
| [`_imaginary_fma_uses_second_product`](#nemo_automodel-components-models-common-rope_triton-_imaginary_fma_uses_second_product) | Identify eager complex multiplication's contraction for this CUDA device.                   |
| [`_launch`](#nemo_automodel-components-models-common-rope_triton-_launch)                                                       | Run the fused kernel over `values` and return a new contiguous tensor.                      |
| [`_rope_kernel`](#nemo_automodel-components-models-common-rope_triton-_rope_kernel)                                             | Copy the leading channels and rotate the trailing pairs of `BLOCK_R` rows.                  |
| [`_uses_fused_kernel`](#nemo_automodel-components-models-common-rope_triton-_uses_fused_kernel)                                 | Check the shared kernel's device, layout, dtype, and autograd contract.                     |
| [`apply_complex_rope`](#nemo_automodel-components-models-common-rope_triton-apply_complex_rope)                                 | Rotate trailing channel pairs using a model's cached complex frequencies.                   |
| [`apply_rope_triton`](#nemo_automodel-components-models-common-rope_triton-apply_rope_triton)                                   | Rotate the final channels of `values` with the fused Triton kernel.                         |

### Data

[`HAVE_TRITON`](#nemo_automodel-components-models-common-rope_triton-HAVE_TRITON)

[`_ELEMENTS_PER_PROGRAM`](#nemo_automodel-components-models-common-rope_triton-_ELEMENTS_PER_PROGRAM)

[`_FUSED_DTYPES`](#nemo_automodel-components-models-common-rope_triton-_FUSED_DTYPES)

### API

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

**Bases:** `Function`

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

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

staticmethod

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

staticmethod

```python
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`** `bool` — default: False

Conjugate the rotation.

---

**Returns:** `torch.Tensor`

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

```python
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.

```python
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`.

```python
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`.

---

```python
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.

```python
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`** `bool` — default: False

Whether to conjugate the rotation.

---

**Returns:** `torch.Tensor`

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

```python
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`** `bool` — default: 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`.

```python
nemo_automodel.components.models.common.rope_triton.HAVE_TRITON: bool = bool(_HAVE_TRITON and _HAVE_TRITON_LANGUAGE)
```

```python
nemo_automodel.components.models.common.rope_triton._ELEMENTS_PER_PROGRAM = 4096
```

```python
nemo_automodel.components.models.common.rope_triton._FUSED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
```