#016

layer_norm

bf16 vllm · both · vllm · importance 0.1%

Reference Implementation

reference.py

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

class Model(nn.Module):

    def __init__(self) -> None:
        super().__init__()

    def forward(self, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
        if x.ndim == 0:
            return x.clone()
        return F.layer_norm(x.float(), (x.shape[-1],), weight=weight.float(), bias=bias.float()).to(x.dtype)

Shapes

T_SOL hardware:

#	shape	T_SOL(XPU-A)	T_Prod	S
0	[276, 1, 1152]	0.24 us	8.90 us	2.7%
1	[600, 1, 1152]	0.52 us	9.90 us	5.3%
2	[1012, 1, 1152]	0.88 us	12.10 us	7.3%
3	[2024, 1, 1152]	1.76 us	15.80 us	11.1%
4	[4100, 1, 1152]	3.57 us	23.30 us	15.3%
5	[8184, 1, 1152]	7.12 us	38.00 us	18.7%
6	[15476, 1, 1152]	13.46 us	64.00 us	21.0%
7	[24952, 1, 1152]	21.69 us	97.50 us	22.2%
8	[40560, 1, 1152]	35.27 us	152.20 us	23.2%
9	[91192, 2048]	140.95 us	788.40 us	17.9%
10	[1, 2048]	0.00 us	7.60 us	0.0%
11	[1, 4096]	0.01 us	8.10 us	0.1%
12	[1, 5120]	0.01 us	8.90 us	0.1%

Input Generation

input.py

import torch

def _make_inputs(shape: list[int]) -> dict[str, torch.Tensor]:
    tensor_shape = tuple((int(dim) for dim in shape))
    if len(tensor_shape) == 0:
        tensor_shape = (1,)
    x = torch.randn(tensor_shape, dtype=torch.bfloat16, device='cuda')
    feature_dim = tensor_shape[-1]
    weight = torch.randn(feature_dim, dtype=torch.bfloat16, device='cuda')
    bias = torch.randn(feature_dim, dtype=torch.bfloat16, device='cuda')
    return {'x': x, 'weight': weight, 'bias': bias}