nemo_automodel.components.models.deepseek_v4.kernels.qnorm_rope

Training Q RMSNorm plus partial interleaved RoPE in token-major storage.

The two BF16 rounding boundaries are intentional: normalize before rotating, and rotate the upstream gradient before applying RMSNorm backward.

Module Contents

Classes

Name	Description
`_QNormRoPE`	Autograd implementation with FP32 saved inverse RMS.

Functions

Name	Description
`_qnorm_rope`	Normalize and rotate M query rows per program into token-major storage.
`_qnorm_rope_backward`	Rotate the upstream gradient back, then apply the RMSNorm input gradient.
`qnorm_rope_triton`	Normalize Q and rotate its trailing interleaved pairs without a transpose copy.

API

class nemo_automodel.components.models.deepseek_v4.kernels.qnorm_rope._QNormRoPE()

Bases: Function

Autograd implementation with FP32 saved inverse RMS.

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

staticmethod

Apply conjugate rotation then the RMSNorm input gradient.

Parameters:

ctx

torch.autograd.function.FunctionCtx

Saved input, frequencies, and inverse RMS.

grad_output

torch.Tensor

Upstream gradient of shape [batch, heads, sequence, channels], any strides.

Returns: torch.Tensor

Input gradient of shape [batch, heads, sequence, channels] in x’s dtype

nemo_automodel.components.models.deepseek_v4.kernels.qnorm_rope._QNormRoPE.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    rotary: int,
    eps: float
) -> torch.Tensor

staticmethod

Normalize and rotate Q while saving its FP32 inverse RMS.

Parameters:

ctx

torch.autograd.function.FunctionCtx

Autograd context.

torch.Tensor

Query tensor of shape [batch, heads, sequence, channels], any strides.

cos

torch.Tensor

Fixed FP32 cosines of shape [batch, sequence, rotary_channels]; the first half holds the unique pair frequencies.

sin

torch.Tensor

Fixed FP32 sines with cos’s shape and frequency order.

rotary

int

Number of trailing channels to rotate.

eps

float

Normalization epsilon.

Returns: torch.Tensor

Tensor of shape [batch, heads, sequence, channels] in x’s dtype, backed

nemo_automodel.components.models.deepseek_v4.kernels.qnorm_rope._qnorm_rope(
    x,
    cos,
    sin,
    out,
    rstd,
    rows,
    S,
    H: tl.constexpr,
    D: tl.constexpr,
    R: tl.constexpr,
    XB,
    XH,
    XS,
    XD,
    CB,
    CS,
    CD,
    SB,
    SS,
    SD,
    EPS: tl.constexpr,
    M: tl.constexpr,
    K: tl.constexpr
)

Normalize and rotate M query rows per program into token-major storage.

Parameters:

Input with logical shape [batch, heads, sequence, channels] and strides (XB, XH, XS, XD); D is channels and R rotary_channels.

cos

Cosines of shape [batch, sequence, rotary_channels] with strides (CB, CS, CD); the first half holds the unique pair frequencies.

sin

Sines with cos’s layout and strides (SB, SS, SD).

out

Contiguous output storage of shape [batch, sequence, heads, channels].

rstd

FP32 inverse RMS of shape [batch * sequence * heads], one per row in the output’s token-major order.

rows

batch * sequence * heads.

nemo_automodel.components.models.deepseek_v4.kernels.qnorm_rope._qnorm_rope_backward(
    x,
    grad,
    cos,
    sin,
    rstd,
    out,
    rows,
    S,
    H: tl.constexpr,
    D: tl.constexpr,
    R: tl.constexpr,
    XB,
    XH,
    XS,
    XD,
    GB,
    GH,
    GS,
    GD,
    CB,
    CS,
    CD,
    SB,
    SS,
    SD,
    M: tl.constexpr,
    K: tl.constexpr
)

Rotate the upstream gradient back, then apply the RMSNorm input gradient.

Parameters:

Saved forward input with logical shape [batch, heads, sequence, channels] and strides (XB, XH, XS, XD).

grad

Upstream gradient with x’s logical shape and strides (GB, GH, GS, GD).

cos

Cosines of shape [batch, sequence, rotary_channels] with strides (CB, CS, CD).

sin

Sines with cos’s layout and strides (SB, SS, SD).

rstd

FP32 inverse RMS of shape [batch * sequence * heads] from the forward.

out

Contiguous input-gradient storage of shape [batch, sequence, heads, channels].

rows

batch * sequence * heads.

nemo_automodel.components.models.deepseek_v4.kernels.qnorm_rope.qnorm_rope_triton(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    rotary: int,
    eps: float
) -> torch.Tensor

Normalize Q and rotate its trailing interleaved pairs without a transpose copy.

Parameters:

torch.Tensor

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

cos

torch.Tensor

Fixed cosine table of shape [batch, sequence, rotary_channels].

sin

torch.Tensor

Fixed sine table with cos’s shape.

rotary

int

Positive even number of trailing channels to rotate.

eps

float

Positive RMSNorm epsilon.

Returns: torch.Tensor

Normalized and rotated tensor with x’s logical shape and token-major storage.