> This page is for version 0.7.0 · 26.10.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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_v4.kernels.rope

One-pass DSV4 partial interleaved RoPE with a conjugate backward.

## Module Contents

### Classes

| Name                                                                                      | Description |
| ----------------------------------------------------------------------------------------- | ----------- |
| [`_PartialRoPE`](#nemo_automodel-components-models-deepseek_v4-kernels-rope-_PartialRoPE) | -           |

### Functions

| Name                                                                                                                | Description                                                          |
| ------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [`_launch`](#nemo_automodel-components-models-deepseek_v4-kernels-rope-_launch)                                     | Run a rotation without materializing a heads-to-tokens transpose.    |
| [`_rotate`](#nemo_automodel-components-models-deepseek_v4-kernels-rope-_rotate)                                     | Rotate \[batch, heads, sequence, channels] into token-major storage. |
| [`apply_partial_rope_triton`](#nemo_automodel-components-models-deepseek_v4-kernels-rope-apply_partial_rope_triton) | Fuse partial interleaved RoPE and its conjugate backward.            |

### API

```python
class nemo_automodel.components.models.deepseek_v4.kernels.rope._PartialRoPE()
```

**Bases:** `Function`

```python
nemo_automodel.components.models.deepseek_v4.kernels.rope._PartialRoPE.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor, None, None, None]
```

staticmethod

Apply the conjugate rotation to the upstream gradient.

**Parameters:**

**`ctx`** `torch.autograd.function.FunctionCtx`

Saved frequency tables and rotary dimension.

---

**`grad_output`** `torch.Tensor`

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

---

**Returns:** `torch.Tensor`

Input gradient in the same logical layout, followed by None for

```python
nemo_automodel.components.models.deepseek_v4.kernels.rope._PartialRoPE.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    rotary: int
) -> torch.Tensor
```

staticmethod

Rotate trailing channels and retain only the frequency tables.

**Parameters:**

**`ctx`** `torch.autograd.function.FunctionCtx`

Autograd context.

---

**`x`** `torch.Tensor`

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

---

**`cos`** `torch.Tensor`

FP32 tensor of shape \[batch, sequence, rotary_channels].

---

**`sin`** `torch.Tensor`

FP32 tensor with the same shape as cos.

---

**`rotary`** `int`

Number of trailing channels to rotate.

---

**Returns:** `torch.Tensor`

Independent tensor with x's logical shape and dtype.

```python
nemo_automodel.components.models.deepseek_v4.kernels.rope._launch(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    rotary: int,
    inverse: bool
) -> torch.Tensor
```

Run a rotation without materializing a heads-to-tokens transpose.

**Parameters:**

**`x`** `torch.Tensor`

CUDA tensor of shape \[batch, heads, sequence, channels].

---

**`cos`** `torch.Tensor`

CUDA FP32 cosines of shape \[batch, sequence, rotary_channels].

---

**`sin`** `torch.Tensor`

CUDA FP32 sines of the same shape as cos.

---

**`rotary`** `int`

Even number of trailing channels to rotate.

---

**`inverse`** `bool`

Whether to negate sine for the backward rotation.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, heads, sequence, channels] with token-major

```python
nemo_automodel.components.models.deepseek_v4.kernels.rope._rotate(
    x,
    cos,
    sin,
    out,
    count,
    SEQUENCE,
    HEADS: tl.constexpr,
    DIM: tl.constexpr,
    ROTARY: tl.constexpr,
    XB,
    XH,
    XS,
    XD,
    CB,
    CS,
    CD,
    SB,
    SS,
    SD,
    INVERSE: tl.constexpr,
    BLOCK: tl.constexpr
)
```

Rotate \[batch, heads, sequence, channels] into token-major storage.

**Parameters:**

**`x`**

Input tensor with logical shape \[batch, heads, sequence, channels].

---

**`cos`**

Cosines of shape \[batch, sequence, rotary_channels], first half unique.

---

**`sin`**

Sines with the same logical layout as cos.

---

**`out`**

Output storage of shape \[batch, sequence, heads, channels].

---

**`count`**

Number of channel pairs in x.

---

```python
nemo_automodel.components.models.deepseek_v4.kernels.rope.apply_partial_rope_triton(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    rotary: int
) -> torch.Tensor
```

Fuse partial interleaved RoPE and its conjugate backward.

**Parameters:**

**`x`** `torch.Tensor`

CUDA FP32/FP16/BF16 tensor of shape \[batch, heads, sequence, channels].
The channel count must be even.

---

**`cos`** `torch.Tensor`

CUDA tensor of shape \[batch, sequence, rotary_channels]; the first
half contains the unique frequencies. Must not require gradients.

---

**`sin`** `torch.Tensor`

CUDA tensor with cos's shape and the same fixed-frequency contract.

---

**`rotary`** `int`

Positive even number of trailing channels, no greater than channels.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, heads, sequence, channels] in x's dtype,