This document is relevant for: Trn2, Trn3

RMSNorm Residual Prefill Kernel API Reference#

Pre-norm residual add followed by RMSNorm, for prefill-shaped inputs (T >= 128).

When residual is given the add runs in the DMA engines via dma_compute, so the sum is produced during the load; it is also spilled to HBM because it is the next block’s residual. Without it the kernel is a plain RMSNorm and returns only out.

Background#

The rmsnorm_residual_prefill kernel computes a pre-norm residual add followed by RMSNorm for prefill-shaped inputs (T >= 128). The residual add runs on the DMA engines during the load.

API Reference#

Source code for this kernel API can be found at: rmsnorm_residual_prefill.py

rmsnorm_residual_prefill#

nkilib.experimental.subkernels.rmsnorm_residual_prefill(hidden: nl.NkiTensor, gamma: nl.NkiTensor, residual: Optional[nl.NkiTensor] = None, eps: float = 1e-06, hidden_actual: Optional[int] = None) → List[nl.NkiTensor]#

Pre-norm residual add followed by RMSNorm, for prefill-shaped inputs (T >= 128).

Parameters:

  • hidden (nl.NkiTensor) – [T, H] or [B, S, H], Hidden states on HBM

  • gamma (nl.NkiTensor) – [1, H], Gamma tensor used in normalization on HBM

  • residual (Optional[nl.NkiTensor]) – [T, H] or [B, S, H], Residual to add before normalizing, on HBM. Omit for a plain RMSNorm

  • eps (float) – Epsilon to maintain numerical stability. Default is 1e-6

  • hidden_actual (Optional[int]) – True hidden width when H is padded, used as the mean divisor. Defaults to H

Returns:

[T, H], Normalized output

Return type:

nl.NkiTensor

Returns:

[T, H], hidden + residual, only when residual is given

Return type:

nl.NkiTensor

Dimensions:

  • T: Token count (batch x sequence, flattened)

This document is relevant for: Trn2, Trn3