Skip to content

vllm.model_executor.layers.fused_moe.silu_mul_token_quant

SiLU-and-mul fused with dynamic per-token FP8 quantization.

Replaces silu_and_mul + scaled_fp8_quant(use_per_token_if_dynamic=True) in the Triton MoE experts when the second GEMM takes per-token-quantized activations (for example PTPC-FP8). Under expert parallelism the first GEMM only writes the rows whose expert is local. This kernel reads the same expert_map[topk_ids] mask and skips the other rows instead of activating and quantizing them. The second GEMM and the finalize never read those rows.

Functions:

silu_mul_token_quant(x, topk_ids=None, expert_map=None, num_experts=None, out=None)

SiLU-and-mul of x ([rows, 2H], gate then up), quantized to FP8 per row.

Returns the FP8 activation [rows, H] and float32 scales [rows, 1]. With topk_ids ([rows] once flattened) and expert_map, rows routed to a remote expert, or to an id outside [0, num_experts), are skipped and their outputs are left untouched. num_experts defaults to expert_map.numel(). out optionally supplies the two output tensors.

Source code in vllm/model_executor/layers/fused_moe/silu_mul_token_quant.py 
 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
110
111
112
113
114
115
116
117
118
119
120
121
122
def silu_mul_token_quant(
    x: torch.Tensor,
    topk_ids: torch.Tensor | None = None,
    expert_map: torch.Tensor | None = None,
    num_experts: int | None = None,
    out: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """SiLU-and-mul of `x` (`[rows, 2H]`, gate then up), quantized to FP8 per row.

    Returns the FP8 activation `[rows, H]` and float32 scales `[rows, 1]`. With
    `topk_ids` (`[rows]` once flattened) and `expert_map`, rows routed to a remote
    expert, or to an id outside `[0, num_experts)`, are skipped and their outputs are
    left untouched. `num_experts` defaults to `expert_map.numel()`. `out` optionally
    supplies the two output tensors.
    """
    assert x.dim() == 2 and x.stride(-1) == 1 and x.size(1) % 2 == 0
    rows, H = x.size(0), x.size(1) // 2
    _, fp8_max = get_fp8_min_max()
    if out is None:
        q = torch.empty((rows, H), dtype=current_platform.fp8_dtype(), device=x.device)
        scale = torch.empty((rows, 1), dtype=torch.float32, device=x.device)
    else:
        q, scale = out
        assert q.shape == (rows, H) and q.stride(-1) == 1
        assert q.dtype == current_platform.fp8_dtype()
        assert scale.shape == (rows, 1) and scale.dtype == torch.float32
        assert scale.is_contiguous()
    if rows == 0:
        return q, scale
    has_mask = False
    ids = emap = x
    if topk_ids is not None and expert_map is not None:
        has_mask = True
        ids, emap = topk_ids.reshape(-1), expert_map
        if num_experts is None:
            num_experts = expert_map.numel()
        assert num_experts <= expert_map.numel()
    _silu_mul_token_quant_kernel[(rows,)](
        x,
        q,
        scale,
        ids,
        emap,
        num_experts or 0,
        H,
        x.stride(0),
        q.stride(0),
        fp8_max,
        _FP8_MIN_SCALING_FACTOR,
        HAS_MASK=has_mask,
        BLOCK=triton.next_power_of_2(H),
        num_warps=4 if H <= 1024 else 8,
    )
    return q, scale