flashinfer.norm.layernorm_quant

flashinfer.norm.layernorm_quant(out: Tensor, input: Tensor, gemma: Tensor, beta: Tensor, scale: float | Tensor, eps: float = 1e-06) → 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 (torch.Tensor) – Scale factor for quantization, shape (1,). The normalized output is divided by this scale before the fp8 cast.

• eps (float) – Epsilon for numerical stability.