Skip to content

vllm.ir.ops.layernorm

Functions:

  • fused_add_rms_norm –

    Fused add and weighted root-mean-square layer normalization

  • rms_norm –

    Weighted root-mean-square layer normalization

fused_add_rms_norm(x, x_residual, weight, epsilon, variance_size=None)

Fused add and weighted root-mean-square layer normalization

Source code in vllm/ir/ops/layernorm.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@register_op(allow_inplace=True)
def fused_add_rms_norm(
    x: Tensor,
    x_residual: Tensor,
    weight: Tensor | None,
    epsilon: float,
    variance_size: int | None = None,
) -> tuple[Tensor, Tensor]:
    """Fused add and weighted root-mean-square layer normalization"""
    orig_dtype = x.dtype
    x = x.to(torch.float32)
    x = x + x_residual.to(torch.float32)
    x_residual = x.to(orig_dtype)

    x_var = x if variance_size is None else x[..., :variance_size]
    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
    x = x * torch.rsqrt(variance + epsilon)
    if weight is not None:
        x = x.to(weight.dtype) * weight
    return x.to(orig_dtype), x_residual

rms_norm(x, weight, epsilon, variance_size=None)

Weighted root-mean-square layer normalization

Source code in vllm/ir/ops/layernorm.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
@register_op
def rms_norm(
    x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
) -> Tensor:
    """Weighted root-mean-square layer normalization"""
    orig_dtype = x.dtype
    x = x.to(torch.float32)
    x_var = x if variance_size is None else x[..., :variance_size]
    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
    x = x * torch.rsqrt(variance + epsilon)
    if weight is not None:
        x = x.to(weight.dtype) * weight
    return x.to(orig_dtype)