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, scale=None)

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.

  • scale

    (Tensor | None, default: None ) –

    optional static per-tensor FP8 scale. When given, the result is quantized to float8_e4m3fn; ineligible inputs take the eager sequence followed by scaled_fp8_quant.

Returns:

  • Tensor –

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

Source code in vllm/kernels/triton/fused_sigmoid_gate_mul.py 
 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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def fused_sigmoid_gate_mul(
    x: torch.Tensor, gate: torch.Tensor, scale: torch.Tensor | None = None
) -> 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.
        scale: optional static per-tensor FP8 scale. When given, the result is
            quantized to float8_e4m3fn; ineligible inputs take the eager
            sequence followed by scaled_fp8_quant.

    Returns:
        A fresh (num_tokens, hidden) tensor with x's dtype (or FP8) 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
        # The FP8 output is hardcoded to e4m3fn; ROCm GEMMs take fnuz.
        or (scale is not None and not current_platform.is_cuda())
    ):
        out = x * torch.sigmoid(gate)
        if scale is None:
            return out
        # Strided gate view (num_heads == 1 per TP rank): quantize the eager
        # result, matching the unfused chain.
        from vllm import _custom_ops as ops

        return ops.scaled_fp8_quant(out, scale, group_shape=(-1, -1))[0]

    num_tokens, hidden = x.shape
    out = torch.empty_like(x, dtype=x.dtype if scale is None else torch.float8_e4m3fn)
    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,
        scale if scale is not None else x,
        hidden,
        x.stride(0),
        gate.stride(0),
        out.stride(0),
        BLOCK=block,
        QUANT_FP8=scale is not None,
        num_warps=num_warps,
    )
    return out