Skip to content

vllm.model_executor.hw_agnostic.layers.activation

Classes:

  • SiluAndMul –

    SwiGLU: x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.

Functions:

SiluAndMul

Bases: CustomOp

SwiGLU: x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.

Source code in vllm/model_executor/hw_agnostic/layers/activation.py 
 9
10
11
12
13
14
15
@CustomOp.register("silu_and_mul")
class SiluAndMul(CustomOp):
    """SwiGLU: ``x -> silu(x[:d]) * x[d:]`` where ``d = x.shape[-1] // 2``."""

    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        d = x.shape[-1] // 2
        return F.silu(x[..., :d]) * x[..., d:]

get_act_and_mul_fn(act_fn_name)

Build the hw-agnostic activation-and-mul op named act_fn_name.

Source code in vllm/model_executor/hw_agnostic/layers/activation.py 
25
26
27
def get_act_and_mul_fn(act_fn_name: str) -> CustomOp:
    """Build the hw-agnostic activation-and-mul op named `act_fn_name`."""
    return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name.lower()]()