> This page is for version Nightly (default).
> 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.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

| Name                                                                                  | Description                                                                      |
| ------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- |
| [`_FusedRope`](#nemo_automodel-components-models-deepseek_v41-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-deepseek_v41-rope_triton-_apply_rope_torch)   | Eager rotation, mirroring `attention._apply_rope` so the fallback needs no circular import. |
| [`_launch`](#nemo_automodel-components-models-deepseek_v41-rope_triton-_launch)                       | Run the fused kernel over `values` and return a new contiguous tensor.                      |
| [`_rope_kernel`](#nemo_automodel-components-models-deepseek_v41-rope_triton-_rope_kernel)             | Copy the leading channels and rotate the trailing pairs of `BLOCK_R` rows.                  |
| [`_uses_fused_kernel`](#nemo_automodel-components-models-deepseek_v41-rope_triton-_uses_fused_kernel) | -                                                                                           |
| [`apply_rope_triton`](#nemo_automodel-components-models-deepseek_v41-rope_triton-apply_rope_triton)   | Rotate the final channels of `values` with the fused Triton kernel.                         |

### Data

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

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

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

### API

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

**Bases:** `Function`

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

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

staticmethod

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

staticmethod

```python
nemo_automodel.components.models.deepseek_v41.rope_triton._apply_rope_torch(
    values: torch.Tensor,
    angles: 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].

---

**`angles`** `torch.Tensor`

FP32 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.deepseek_v41.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.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
)
```

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.

---

**`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.deepseek_v41.rope_triton._uses_fused_kernel(
    values: torch.Tensor,
    angles: torch.Tensor
) -> bool
```

```python
nemo_automodel.components.models.deepseek_v41.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:**

* `RuntimeError`: `values` is on CUDA but Triton is not importable.
* `ValueError`: `angles` does not describe an even rotary width within `values`.

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

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

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