nemo_automodel.components.models.deepseek_v4.kernels.rope

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

Module Contents

Classes

Name	Description
`_PartialRoPE`	-

Functions

Name	Description
`_launch`	Run a rotation without materializing a heads-to-tokens transpose.
`_rotate`	Rotate [batch, heads, sequence, channels] into token-major storage.
`apply_partial_rope_triton`	Fuse partial interleaved RoPE and its conjugate backward.

API

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

Bases: Function

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

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.

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.

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:

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

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:

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.

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:

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,