#021

moe_sum_reduce

bf16 sglang · both · sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels · importance 0.3%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

    def __init__(self, routed_scaling_factor: float=1.0) -> None:
        super().__init__()
        self.routed_scaling_factor = float(routed_scaling_factor)

    def forward(self, expert_outputs: torch.Tensor) -> torch.Tensor:
        out = expert_outputs.float().sum(dim=1) * self.routed_scaling_factor
        return out.to(expert_outputs.dtype)

Shapes

T_SOL hardware:

#	token_count	top_k	hidden_size	T_SOL(XPU-A)	T_Prod	S
0	11027	8	2048	76.70 us	362.40 us	21.2%
1	6030	8	2048	41.94 us	209.80 us	20.0%
2	5983	8	2048	41.61 us	213.30 us	19.5%
3	5043	8	2048	35.08 us	179.80 us	19.5%
4	4996	8	2048	34.75 us	178.40 us	19.5%
5	4363	8	2048	30.35 us	159.20 us	19.1%
6	4195	8	2048	29.18 us	154.00 us	18.9%

Input Generation

input.py

import torch

def _make_inputs(token_count: int, top_k: int, hidden_size: int) -> dict[str, torch.Tensor]:
    expert_outputs = torch.randn(token_count, top_k, hidden_size, dtype=torch.bfloat16, device='cuda') * 0.02
    return {'expert_outputs': expert_outputs}