Skip to content

vllm.kernels.triton.fused_sigmoid_gate_mul

Fused sigmoid-gate-multiply Triton kernel: out = x * sigmoid(gate).

Functions:

fused_sigmoid_gate_mul(x, gate)

Compute x * sigmoid(gate) in a single Triton kernel.

Falls back to the eager x * torch.sigmoid(gate) sequence unless both inputs are same-shape, same-dtype, contiguous 2D CUDA tensors.

Parameters:

  • x

    (Tensor) –

    (num_tokens, hidden) tensor, e.g. attention output.

  • gate

    (Tensor) –

    (num_tokens, hidden) pre-sigmoid gate.

Returns:

  • Tensor –

    A fresh (num_tokens, hidden) tensor with x's dtype and device.

Source code in vllm/kernels/triton/fused_sigmoid_gate_mul.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def fused_sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
    """Compute ``x * sigmoid(gate)`` in a single Triton kernel.

    Falls back to the eager ``x * torch.sigmoid(gate)`` sequence unless both
    inputs are same-shape, same-dtype, contiguous 2D CUDA tensors.

    Args:
        x: (num_tokens, hidden) tensor, e.g. attention output.
        gate: (num_tokens, hidden) pre-sigmoid gate.

    Returns:
        A fresh (num_tokens, hidden) tensor with x's dtype and device.

    """
    if (
        x.ndim != 2
        or x.shape != gate.shape
        or x.dtype != gate.dtype
        or not x.is_contiguous()
        or not gate.is_contiguous()
        or not x.is_cuda
    ):
        return x * torch.sigmoid(gate)

    num_tokens, hidden = x.shape
    out = torch.empty_like(x)
    if num_tokens == 0 or hidden == 0:
        return out

    if hidden <= 1024:
        block = triton.next_power_of_2(hidden)
        num_warps = 4
    elif hidden <= 4096:
        block = 1024
        num_warps = 4
    else:
        block = 2048
        num_warps = 8

    grid = (num_tokens, triton.cdiv(hidden, block))
    _fused_sigmoid_gate_mul_kernel[grid](
        x,
        gate,
        out,
        hidden,
        x.stride(0),
        gate.stride(0),
        out.stride(0),
        BLOCK=block,
        num_warps=num_warps,
    )
    return out