triton-layer-norm

[!CAUTION] Starting from September 13, 2026, we will be removing the "model" type repositories of kernels (e.g., kernels-community/flash-attn3). Make sure you're using a latest version of kernels. If you face any disruption, please report them here: https://github.com/huggingface/kernels/issues/new.

Triton layer normalization kernels.

This kernel implements layers normalization using Triton. This kernel is from the flash-attention project.

Functions

Function layer_norm

(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, residual: Optional[torch.Tensor] = None, x1: Optional[torch.Tensor] = None, weight1: Optional[torch.Tensor] = None, bias1: Optional[torch.Tensor] = None, eps: float = 1e-06, dropout_p: float = 0.0, rowscale=None, prenorm: bool = False, residual_in_fp32: bool = False, is_rms_norm: bool = False, return_dropout_mask: bool = False, out: Optional[torch.Tensor] = None, residual_out: Optional[torch.Tensor] = None)

Apply layer normalization to the input tensor with Triton acceleration.

Parameters

• x (torch.Tensor) -- Input tensor to normalize.
• weight (torch.Tensor) -- Scale parameter for normalization.
• bias (torch.Tensor) -- Shift parameter for normalization.
• residual (torch.Tensor, optional) -- Optional residual tensor to add to the input before normalization.
• x1 (torch.Tensor, optional) -- Optional second input tensor to combine with x. When provided, the function first adds x1 to x and then applies normalization.
• weight1 (torch.Tensor, optional) -- Scale parameter for the second normalization.
• bias1 (torch.Tensor, optional) -- Shift parameter for the second normalization.
• eps (float, optional, defaults to 1e-6) -- Small constant added for numerical stability in normalization.
• dropout_p (float, optional, defaults to 0.0) -- Dropout probability. If greater than 0, applies dropout to the input before normalization and residual addition.
• rowscale (torch.Tensor, optional) -- Optional scaling factor applied to each row of the input tensor. Not compatible with the use of x1.
• prenorm (bool, optional, defaults to False) -- If True, returns both the normalized output and the unnormalized input+residual.
• residual_in_fp32 (bool, optional, defaults to False) -- If True, performs the residual connection in FP32 precision.
• is_rms_norm (bool, optional, defaults to False) -- If True, uses RMS normalization instead of layer normalization.
• return_dropout_mask (bool, optional, defaults to False) -- If True, returns the dropout mask used for the computation.
• out (torch.Tensor, optional) -- Output tensor for the normalized result. If None, a new tensor is allocated.
• residual_out (torch.Tensor, optional) -- Output tensor for the residual result when using prenorm. If None, a new tensor is allocated when needed.

Returns

Type: torch.Tensor or tuple of torch.Tensor

• The normalized input.
• The second normalization of the input if weight1 is provided.
• The residual tensor if prenorm is set.
• The dropout mask if return_dropout_mask is set.
• The dropout mask for x1 if x1 is provided and return_dropout_mask is set.

Layers

Class LlamaRMSNorm

No documentation available.

Methods

Method forward

(self, hidden_states: torch.Tensor) -> torch.Tensor

No documentation available.