#022

moe_topk_gating_softmax

fp32 vllm · both · vllm._custom_ops · importance 2.1%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

    def __init__(self, top_k: int, num_expert_groups: int = 0, topk_group: int = 0) -> None:
        super().__init__()
        self.top_k = top_k
        self.num_expert_groups = num_expert_groups
        self.topk_group = topk_group

    def forward(self, gating_output: torch.Tensor) -> dict[str, torch.Tensor]:
        scores = torch.softmax(gating_output.float(), dim=-1)
        if self.num_expert_groups > 0 and self.topk_group > 0:
            n, ne = scores.shape
            group_size = ne // self.num_expert_groups
            grouped = scores.view(n, self.num_expert_groups, group_size)
            group_scores = grouped.amax(dim=-1)
            _, top_groups = torch.topk(group_scores, k=self.topk_group, dim=-1)
            mask = torch.zeros(n, self.num_expert_groups, dtype=torch.bool, device=scores.device)
            mask.scatter_(1, top_groups, True)
            mask = mask.unsqueeze(-1).expand(-1, -1, group_size).reshape(n, ne)
            scores = scores.masked_fill(~mask, 0.0)
        (topk_weights, topk_ids) = torch.topk(scores, k=self.top_k, dim=-1)
        return {'topk_weights': topk_weights, 'topk_ids': topk_ids.to(torch.int32)}

Shapes

T_SOL hardware:

#	token_count	num_experts	T_SOL(XPU-A)	T_Prod	S
0	1	128	0.00 us	11.40 us	0.0%
1	129	128	0.01 us	12.80 us	0.1%
2	256	128	0.03 us	13.00 us	0.2%
3	520	128	0.06 us	13.00 us	0.5%
4	1019	128	0.11 us	13.10 us	0.8%
5	2044	128	0.22 us	13.30 us	1.7%
6	4083	128	0.44 us	17.10 us	2.6%
7	8192	128	0.89 us	25.60 us	3.5%
8	8037	256	1.65 us	39.40 us	4.2%
9	13526	256	2.78 us	56.20 us	4.9%

Input Generation

input.py

import torch

def _make_inputs(token_count: int, num_experts: int) -> dict[str, torch.Tensor]:
    gating_output = torch.randn(token_count, num_experts, dtype=torch.float32, device='cuda')
    return {'gating_output': gating_output}