#015

l2_norm

bf16 vllm · both · vllm.model_executor.layers.fla.ops.l2norm · importance 0.5%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_float = x.float()
        norm_sq = (x_float * x_float).sum(dim=-1, keepdim=True)
        return (x_float * torch.rsqrt(norm_sq + self.eps)).to(x.dtype)

Shapes

T_SOL hardware:

#	num_tokens	hidden_size	T_SOL(XPU-A)	T_Prod	S
0	1	2048	0.00 us	7.60 us	0.0%
1	128	2048	0.20 us	7.70 us	2.6%
2	256	2048	0.40 us	7.90 us	5.1%
3	480	2048	0.74 us	9.00 us	8.2%
4	2304	2048	3.56 us	21.50 us	16.6%
5	4064	2048	6.28 us	30.00 us	20.9%
6	8736	2048	13.50 us	59.00 us	22.9%
7	16608	2048	25.67 us	117.20 us	21.9%
8	24480	2048	37.84 us	173.30 us	21.8%
9	48704	2048	75.28 us	315.90 us	23.8%
10	65792	2048	101.69 us	430.00 us	23.6%
11	98880	2048	152.83 us	635.00 us	24.1%
12	131072	2048	202.59 us	853.10 us	23.7%

Input Generation

input.py

import torch

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

def get_inputs() -> list[torch.Tensor]:
    return list(_make_inputs(num_tokens=1024, hidden_size=128).values())

def get_init_inputs() -> dict[str, object]:
    return {'hidden_size': 128}