> This page is for version Latest. > 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.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`](#nemo_automodel-components-models-deepseek_v4-kernels-qnorm_rope-_QNormRoPE) | Autograd implementation with FP32 saved inverse RMS. | ### Functions | Name | Description | | --------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- | | [`_qnorm_rope`](#nemo_automodel-components-models-deepseek_v4-kernels-qnorm_rope-_qnorm_rope) | Normalize and rotate M query rows per program into token-major storage. | | [`_qnorm_rope_backward`](#nemo_automodel-components-models-deepseek_v4-kernels-qnorm_rope-_qnorm_rope_backward) | Rotate the upstream gradient back, then apply the RMSNorm input gradient. | | [`qnorm_rope_triton`](#nemo_automodel-components-models-deepseek_v4-kernels-qnorm_rope-qnorm_rope_triton) | Normalize Q and rotate its trailing interleaved pairs without a transpose copy. | ### API ```python class nemo_automodel.components.models.deepseek_v4.kernels.qnorm_rope._QNormRoPE() ``` **Bases:** `Function` Autograd implementation with FP32 saved inverse RMS. ```python 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 ```python 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. --- **`x`** `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 ```python 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:** **`x`** 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. --- ```python 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:** **`x`** 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. --- ```python 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:** **`x`** `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.