#012

fused_rmsnorm_quant

fp8_e4m3 aiter · both · aiter.ops.rmsnorm · importance 0.3%

Reference Implementation

reference.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):

    def __init__(self, eps: float=1e-05) -> None:
        super().__init__()
        self.eps = eps

    def forward(self, x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        residual_out = x + residual
        normed = F.rms_norm(residual_out, (residual_out.shape[-1],), weight, self.eps)
        normed_f = normed.to(torch.float32)
        finfo = torch.finfo(torch.float8_e4m3fnuz)
        amax = normed_f.abs().amax(dim=-1, keepdim=True)
        scale = amax / finfo.max
        scale = torch.where(scale == 0, torch.ones_like(scale), scale)
        quantized = (normed_f / scale).to(torch.float8_e4m3fnuz)
        return (quantized, scale.to(torch.float32), residual_out)

Shapes

T_SOL hardware:

#	token_count	hidden_size	T_SOL(XPU-A)	T_Prod	S
0	117	2048	0.32 us	7.90 us	4.1%
1	837	2048	2.27 us	11.10 us	20.5%
2	1769	2048	4.79 us	15.90 us	30.1%
3	3260	2048	8.82 us	23.40 us	37.7%
4	11399	2048	30.84 us	65.10 us	47.4%
5	26080	2048	70.56 us	140.90 us	50.1%
6	650	4096	3.52 us	13.50 us	26.1%
7	1013	4096	5.48 us	16.60 us	33.0%
8	2058	4096	11.14 us	27.40 us	40.7%
9	4083	4096	22.09 us	47.00 us	47.0%
10	8208	4096	44.41 us	86.50 us	51.3%

Input Generation

input.py

import torch


def _make_inputs(token_count: int, hidden_size: int) -> dict[str, torch.Tensor]:
    x = torch.randn(token_count, hidden_size, dtype=torch.bfloat16, device='cuda')
    residual = torch.randn(token_count, hidden_size, dtype=torch.bfloat16, device='cuda')
    weight = torch.randn(hidden_size, dtype=torch.bfloat16, device='cuda')
    return {'x': x, 'residual': residual, 'weight': weight}