#009

fused_moe

bf16 vllm · both · vllm.model_executor.layers.fused_moe.fused_moe · importance 11.4%

Reference Implementation

reference.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):

    def __init__(self, num_experts: int, intermediate_size: int, top_k: int) -> None:
        super().__init__()
        self.num_experts = num_experts
        self.intermediate_size = intermediate_size
        self.top_k = top_k

    def forward(self, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor) -> torch.Tensor:
        token_count = hidden_states.shape[0]
        hidden_size = hidden_states.shape[1]
        intermediate_size = self.intermediate_size
        output = torch.zeros(token_count, hidden_size, device=hidden_states.device, dtype=torch.float32)
        hidden_fp32 = hidden_states.float()
        for expert_index in range(self.num_experts):
            mask = topk_ids == expert_index
            if not mask.any():
                continue
            weight_for_expert = (topk_weights * mask.to(topk_weights.dtype)).sum(dim=1)
            token_idx = weight_for_expert.nonzero(as_tuple=True)[0]
            if token_idx.numel() == 0:
                continue
            x = hidden_fp32.index_select(0, token_idx)
            w1_e = w1[expert_index].float()
            w2_e = w2[expert_index].float()
            intermediate = F.linear(x, w1_e)
            gate = intermediate[:, :intermediate_size]
            up = intermediate[:, intermediate_size:]
            activated = F.silu(gate) * up
            expert_output = F.linear(activated, w2_e)
            output.index_add_(0, token_idx, expert_output * weight_for_expert.index_select(0, token_idx).unsqueeze(1))
        return output

Shapes

T_SOL hardware:

#	token_count	hidden_size	intermediate_size	num_experts	top_k	T_SOL(XPU-A)	T_Prod	S
0	135	2048	2048	8	2	38.30 us	232.70 us	16.5%
1	277	2048	2048	8	2	56.71 us	249.80 us	22.7%
2	668	2048	2048	8	2	136.79 us	477.10 us	28.7%
3	1023	2048	2048	8	2	208.22 us	539.00 us	38.6%
4	1979	2048	2048	8	2	403.02 us	772.50 us	52.2%
5	4195	2048	2048	8	2	851.82 us	1.60 ms	53.1%
6	7689	2048	2048	8	2	1.56 ms	2.68 ms	58.1%
7	15809	2048	2048	8	2	3.20 ms	5.37 ms	59.7%
8	1	2048	768	128	8	14.25 us	116.90 us	12.2%
9	2400	2048	768	128	8	757.05 us	2.16 ms	35.1%
10	4195	2048	768	128	8	1.32 ms	2.98 ms	44.4%
11	8192	2048	768	128	8	2.59 ms	4.99 ms	51.9%
12	15809	2048	768	128	8	4.99 ms	9.09 ms	54.9%
13	25632	2048	768	128	8	8.10 ms	14.48 ms	55.9%
14	45936	2048	768	128	8	14.52 ms	25.71 ms	56.5%
15	520	4096	384	128	8	230.33 us	595.80 us	38.7%
16	1019	4096	384	128	8	322.57 us	882.90 us	36.5%
17	2044	4096	384	128	8	646.04 us	1.61 ms	40.2%
18	3936	4096	384	128	8	1.24 ms	2.80 ms	44.5%
19	8192	4096	384	128	8	2.59 ms	5.69 ms	45.4%
20	4098	4096	4096	8	2	3.33 ms	5.33 ms	62.5%
21	8179	4096	4096	8	2	6.63 ms	10.10 ms	65.7%
22	15381	4096	4096	8	2	12.48 ms	18.64 ms	67.0%

Input Generation

input.py

import torch


def _make_inputs(
    token_count: int, hidden_size: int, intermediate_size: int, num_experts: int, top_k: int
) -> dict[str, torch.Tensor]:
    hidden_states = torch.randn(token_count, hidden_size, dtype=torch.bfloat16, device="cuda") * 0.1
    w1 = torch.randn(
        num_experts, 2 * intermediate_size, hidden_size, dtype=torch.bfloat16, device="cuda"
    ) * hidden_size ** (-0.5)
    w2 = torch.randn(
        num_experts, hidden_size, intermediate_size, dtype=torch.bfloat16, device="cuda"
    ) * intermediate_size ** (-0.5)
    # Realistic top-k routing: select top_k experts per token from router logits so that
    # each token's selected experts are unique (real MoE never routes a token to the same
    # expert twice). Using torch.randint here instead would produce duplicate experts per
    # token, an unphysical case that sorted-scatter MoE kernels handle differently from the
    # gather reference, spuriously failing accuracy checks.
    router_logits = torch.randn(token_count, num_experts, dtype=torch.float32, device="cuda")
    topk_weights, topk_ids = torch.topk(router_logits.softmax(dim=-1), top_k, dim=-1)
    topk_weights = (topk_weights / topk_weights.sum(dim=-1, keepdim=True)).to(torch.float32)
    topk_ids = topk_ids.to(torch.int32)
    return {
        "hidden_states": hidden_states,
        "w1": w1,
        "w2": w2,
        "topk_weights": topk_weights,
        "topk_ids": topk_ids,
    }