#017

linear_sigmoid_mul

bf16 sglang · both · sglang.srt.models.qwen2_moe · importance 0.7%

Reference Implementation

reference.py

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

class Model(nn.Module):

    def __init__(self) -> None:
        super().__init__()

    def forward(self, hidden_states: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, post_mul_mat: torch.Tensor) -> torch.Tensor:
        gate = F.linear(hidden_states, weight, bias)
        out = torch.sigmoid(gate.float()) * post_mul_mat.float()
        return out.to(post_mul_mat.dtype)

Shapes

T_SOL hardware:

#	token_count	hidden_size	out_features	T_SOL(XPU-A)	T_Prod	S
0	8037	4096	4096	1.16 ms	1.60 ms	72.3%
1	6717	4096	4096	968.89 us	1.27 ms	76.5%
2	4001	4096	4096	577.12 us	2.06 ms	28.0%
3	15381	4096	4096	2.22 ms	7.21 ms	30.8%
4	6576	4096	4096	948.55 us	1.26 ms	75.3%
5	1	2048	4096	3.17 us	20.90 us	15.2%
6	90	2048	4096	6.49 us	32.20 us	20.2%
7	13	2048	4096	3.22 us	17.30 us	18.6%
8	2748	4096	4096	396.38 us	564.50 us	70.2%

Input Generation

input.py

import torch

def _make_inputs(token_count: int, hidden_size: int, out_features: int) -> dict[str, torch.Tensor]:
    hidden_states = torch.randn(token_count, hidden_size, dtype=torch.bfloat16, device='cuda') * 0.02
    weight = torch.randn(out_features, hidden_size, dtype=torch.bfloat16, device='cuda') * hidden_size ** (-0.5)
    bias = torch.randn(out_features, dtype=torch.bfloat16, device='cuda') * 0.02
    post_mul_mat = torch.randn(token_count, out_features, dtype=torch.bfloat16, device='cuda') * 0.02
    return {'hidden_states': hidden_states, 'weight': weight, 'bias': bias, 'post_mul_mat': post_mul_mat}