flashinfer.norm.layernorm_quant

flashinfer.norm.layernorm_quant(out: Tensor, input: Tensor, gemma: Tensor, beta: Tensor, scale: float | Tensor, eps: float = 1e-06, enable_pdl: bool | None = None) → None

Layer normalization + fp8 quantization.

out[i] = (((input[i] - E[input]) / sqrt(Var[input] + eps)) * gemma[i] + beta[i]) / scale

Parameters:

  • out (torch.Tensor) – The output tensor, shape (batch_size, hidden_size). Need to be contiguous. The output is quantized to the dtype of this tensor, which must be float8_e4m3fn or float8_e5m2.

  • input (torch.Tensor) – Input tensor, shape (batch_size, hidden_size). Need to be bfloat16 and contiguous.

  • gemma (torch.Tensor) – Gemma tensor, shape (hidden_size,). Need to be float32.

  • beta (torch.Tensor) – Beta tensor, shape (hidden_size,). Need to be float32.

  • scale (Union[float, torch.Tensor]) – Scale factor for quantization; a tensor scale must have shape (1,). The normalized output is divided by this scale before the fp8 cast.

  • eps (float) – Epsilon for numerical stability.

  • enable_pdl (bool) – Whether to enable programmatic dependent launch. Only honored by the CuTe-DSL backend; the CUDA fallback ignores it.