Skip to content

vllm.ir.ops

Modules:

Functions:

  • fused_add_rms_norm –

    Fused add and weighted root-mean-square layer normalization

  • gelu_and_mul_sparse –

    Apply Gaussian sparsification, GELU, and gated multiplication.

  • 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 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
@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

gelu_and_mul_sparse(x, std_multiplier, approximate='none')

Apply Gaussian sparsification, GELU, and gated multiplication.

Source code in vllm/ir/ops/activation.py 
11
12
13
14
15
16
17
18
19
20
21
22
@register_op
def gelu_and_mul_sparse(
    x: Tensor, std_multiplier: float, approximate: str = "none"
) -> Tensor:
    """Apply Gaussian sparsification, GELU, and gated multiplication."""
    d = x.shape[-1] // 2
    gate = x[..., :d]
    # Statistics intentionally remain local to each tensor-parallel shard.
    mean = torch.mean(gate, dim=-1, keepdim=True)
    std = torch.std(gate, dim=-1, keepdim=True, unbiased=False)
    sparse_gate = F.relu(gate - (mean + std * std_multiplier))
    return F.gelu(sparse_gate, approximate=approximate) * x[..., d:]

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)