Skip to content

vllm.ir.ops.activation

Functions:

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:]