#027

rms_norm

bf16 vllm · both · vllm.csrc.layernorm_kernels · importance 2.8%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

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

    def forward(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
        x_f = x.float()
        variance = x_f.square().mean(dim=-1, keepdim=True)
        out = x_f * torch.rsqrt(variance + self.eps) * weight.float()
        return out.to(x.dtype)

Shapes

T_SOL hardware:

#	token_count	hidden_size	T_SOL(XPU-A)	T_Prod	S
0	2500	1024	1.93 us	9.90 us	19.5%
1	4968	1024	3.84 us	14.10 us	27.2%
2	10000	1024	7.73 us	22.80 us	33.9%
3	16200	1024	12.52 us	32.70 us	38.3%
4	22100	1024	17.08 us	42.10 us	40.6%
5	360	1152	0.31 us	7.80 us	4.0%
6	720	1152	0.63 us	8.10 us	7.8%
7	1200	1152	1.04 us	9.20 us	11.3%
8	2116	1152	1.84 us	12.30 us	15.0%
9	3844	1152	3.34 us	17.70 us	18.9%
10	8136	1152	7.07 us	31.10 us	22.7%
11	15000	1152	13.04 us	52.10 us	25.0%
12	24368	1152	21.19 us	81.20 us	26.1%
13	49152	1152	42.73 us	157.40 us	27.1%
14	65556	1152	57.00 us	208.00 us	27.4%
15	3	128	0.00 us	7.50 us	0.0%
16	1504	128	0.15 us	7.70 us	1.9%
17	2048	128	0.20 us	7.70 us	2.6%
18	4096	128	0.40 us	8.70 us	4.6%
19	8192	128	0.79 us	11.80 us	6.7%
20	16384	128	1.58 us	18.50 us	8.5%
21	24544	128	2.37 us	24.40 us	9.7%
22	49024	128	4.74 us	42.00 us	11.3%
23	65488	128	6.33 us	53.90 us	11.7%
24	95952	128	9.27 us	75.80 us	12.2%
25	131072	128	12.66 us	101.10 us	12.5%
26	247680	128	23.93 us	190.80 us	12.5%
27	505888	128	48.87 us	371.20 us	13.2%
28	7	1280	0.01 us	7.50 us	0.1%
29	24752	1280	23.91 us	82.10 us	29.1%
30	1736	1536	2.01 us	11.60 us	17.3%
31	5000	1536	5.80 us	22.10 us	26.2%
32	128	2048	0.20 us	7.70 us	2.6%
33	256	2048	0.40 us	7.80 us	5.1%
34	512	2048	0.79 us	8.20 us	9.6%
35	1029	2048	1.59 us	9.10 us	17.5%
36	2010	2048	3.11 us	12.60 us	24.7%
37	4174	2048	6.45 us	21.00 us	30.7%
38	8290	2048	12.81 us	34.90 us	36.7%
39	16530	2048	25.55 us	62.80 us	40.7%
40	20538	2048	31.75 us	80.50 us	39.4%
41	3	256	0.00 us	7.50 us	0.0%
42	1672	2560	3.23 us	14.20 us	22.7%
43	298	3584	0.81 us	7.90 us	10.3%
44	514	3584	1.39 us	8.70 us	16.0%
45	1024	3584	2.77 us	11.40 us	24.3%
46	3	512	0.00 us	7.40 us	0.0%
47	1	5120	0.01 us	7.50 us	0.1%
48	53	7168	0.29 us	7.70 us	3.8%
49	141	7168	0.77 us	7.90 us	9.7%
50	632	7168	3.42 us	12.90 us	26.5%
51	917	7168	4.96 us	16.40 us	30.2%
52	1688	7168	9.13 us	24.90 us	36.7%
53	3804	7168	20.58 us	46.50 us	44.3%
54	154	8192	0.96 us	8.00 us	12.0%
55	510	8192	3.16 us	12.40 us	25.5%

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') * 0.02
    weight = torch.randn(hidden_size, dtype=torch.bfloat16, device='cuda') * 0.02 + 1.0
    return {'x': x, 'weight': weight}