minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] expand the interface of shared_expert without scaling factor (#22933)

Browse Source 
merge since this is CPU only change on sgl-kernel.
This commit is contained in:
Ma Mingfei
2026-04-21 20:03:39 +08:00
committed by GitHub
parent 48daa831ea
commit 929e00eeab
9 changed files with 313 additions and 623 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/cpu/test_moe.py
+8 -27
View File
@@ -300,35 +300,16 @@ class TestFusedExperts(CustomTestCase):
)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
awq_w13_weight_pack = []
awq_w13_zero_pack = []
awq_w13_scales_pack = []
awq_w2_weight_pack = []
awq_w2_zero_pack = []
awq_w2_scales_pack = []
for i in range(E):
packed_weight_13_i, packed_zero_13_i, packed_scales_13_i = (
torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i]
)
awq_w13_weight_pack, awq_w13_zero_pack, awq_w13_scales_pack = (
torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
awq_w13_weight, awq_w13_zero, awq_w13_scales
)
awq_w13_weight_pack.append(packed_weight_13_i)
awq_w13_zero_pack.append(packed_zero_13_i)
awq_w13_scales_pack.append(packed_scales_13_i)
packed_weight_2_i, packed_zero_2_i, packed_scales_2_i = (
torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i]
)
)
awq_w2_weight_pack, awq_w2_zero_pack, awq_w2_scales_pack = (
torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
awq_w2_weight, awq_w2_zero, awq_w2_scales
)
awq_w2_weight_pack.append(packed_weight_2_i)
awq_w2_zero_pack.append(packed_zero_2_i)
awq_w2_scales_pack.append(packed_scales_2_i)
awq_w13_weight_pack = torch.stack(awq_w13_weight_pack).detach()
awq_w13_zero_pack = torch.stack(awq_w13_zero_pack).detach()
awq_w13_scales_pack = torch.stack(awq_w13_scales_pack).detach()
awq_w2_weight_pack = torch.stack(awq_w2_weight_pack).detach()
awq_w2_zero_pack = torch.stack(awq_w2_zero_pack).detach()
awq_w2_scales_pack = torch.stack(awq_w2_scales_pack).detach()
)
out = kernel.fused_experts_cpu(
a,
test/srt/cpu/test_shared_expert.py
+64 -50
View File
@@ -2,12 +2,10 @@ import itertools
import math
import unittest
# TODO: use interface in cpu.py
import torch
from utils import (
BLOCK_K,
BLOCK_N,
SiluAndMul,
factor_for_scale,
fp8_max,
fp8_min,
@@ -18,7 +16,6 @@ from utils import (
torch_w8a8_per_column_moe,
)
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.test_utils import CustomTestCase
torch.manual_seed(1234)
@@ -29,37 +26,41 @@ class TestSharedExpert(CustomTestCase):
N = [32, 32 * 4]
K = [32, 32 * 2]
routed_scaling_factor = [16]
apply_scaling_factor = [True, False]
M_fp8 = [2, 12]
N_fp8 = [512]
K_fp8 = [256]
def _bf16_shared_expert(self, m, n, k, routed_scaling_factor):
def _bf16_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
dtype = torch.bfloat16
prepack = True
hidden_states = torch.randn(m, k, dtype=dtype) / k
w1 = torch.randn(2 * n, k, dtype=dtype)
w2 = torch.randn(k, n, dtype=dtype)
fused_output = torch.randn(m, k, dtype=dtype) / k
fused_output = (
torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None
)
routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None
# fused moe mutates content in hs
hidden_states2 = hidden_states.clone()
# bfloat16
ref = torch_naive_moe(
hidden_states.float(),
w1.float(),
w2.float(),
fused_output.float(),
routed_scaling_factor,
).to(dtype=dtype)
res = torch.ops.sgl_kernel.shared_expert_cpu(
hidden_states,
w1,
w2,
fused_output,
routed_scaling_factor,
output_dtype=dtype,
)
out = torch.ops.sgl_kernel.shared_expert_cpu(
hidden_states2,
w1,
w2,
fused_output,
routed_scaling_factor,
True,
False,
False,
@@ -70,7 +71,7 @@ class TestSharedExpert(CustomTestCase):
)
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, res, atol=atol, rtol=rtol)
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
def test_bf16_shared_expert(self):
for params in itertools.product(
@@ -78,39 +79,43 @@ class TestSharedExpert(CustomTestCase):
self.N,
self.K,
self.routed_scaling_factor,
self.apply_scaling_factor,
):
with self.subTest(
m=params[0],
n=params[1],
k=params[2],
routed_scaling_factor=params[3],
apply_scaling_factor=params[4],
):
self._bf16_shared_expert(*params)
def _int8_shared_expert(self, m, n, k, routed_scaling_factor):
def _int8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
dtype = torch.bfloat16
prepack = True
hidden_states = torch.randn(m, k, dtype=dtype) / k
w1 = torch.randn(2 * n, k, dtype=dtype)
w2 = torch.randn(k, n, dtype=dtype)
fused_output = torch.randn(m, k, dtype=dtype) / k
fused_output = (
torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None
)
routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None
# fused moe mutates content in hs
hidden_states2 = hidden_states.clone()
w1_q, w1_s = per_token_quant_int8(w1)
w2_q, w2_s = per_token_quant_int8(w2)
ref2 = torch_w8a8_per_column_moe(
hidden_states2.float(),
ref = torch_w8a8_per_column_moe(
hidden_states,
w1_q,
w2_q,
w1_s,
w2_s,
fused_output.float(),
fused_output,
routed_scaling_factor,
).to(dtype=dtype)
res2 = torch.ops.sgl_kernel.shared_expert_cpu(
)
out = torch.ops.sgl_kernel.shared_expert_cpu(
hidden_states2,
w1_q,
w2_q,
@@ -125,8 +130,8 @@ class TestSharedExpert(CustomTestCase):
False,
)
atol = rtol = precision[ref2.dtype]
torch.testing.assert_close(ref2, res2, atol=atol, rtol=rtol)
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
def test_int8_shared_expert(self):
for params in itertools.product(
@@ -134,57 +139,64 @@ class TestSharedExpert(CustomTestCase):
self.N,
self.K,
self.routed_scaling_factor,
self.apply_scaling_factor,
):
with self.subTest(
m=params[0],
n=params[1],
k=params[2],
routed_scaling_factor=params[3],
apply_scaling_factor=params[4],
):
self._int8_shared_expert(*params)
def _fp8_shared_expert(self, M, N, K, routed_scaling_factor):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
def _fp8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
dtype = torch.bfloat16
prepack = True
a = torch.randn(M, K, dtype=dtype) / math.sqrt(K)
hidden_states = torch.randn(m, k, dtype=dtype) / math.sqrt(k)
w1_fp32 = torch.randn(1, 2 * N, K)
w1_fp32 = torch.randn(1, 2 * n, k)
w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w2_fp32 = torch.randn(1, K, N)
w2_fp32 = torch.randn(1, k, n)
w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w1s = torch.randn(1, 2 * N // BLOCK_N, K // BLOCK_K) * factor_for_scale
w2s = torch.randn(1, K // BLOCK_N, N // BLOCK_K) * factor_for_scale
w1s = torch.randn(1, 2 * n // BLOCK_N, k // BLOCK_K) * factor_for_scale
w2s = torch.randn(1, k // BLOCK_N, n // BLOCK_K) * factor_for_scale
w1_scaled = scaled_weight(w1, w1s).view(2 * N, K)
w2_scaled = scaled_weight(w2, w2s).view(K, N)
w1_scaled = scaled_weight(w1, w1s).view(2 * n, k)
w2_scaled = scaled_weight(w2, w2s).view(k, n)
# change back to 2D
w1, w2 = w1.squeeze(0), w2.squeeze(0)
w1s, w2s = w1s.squeeze(0), w2s.squeeze(0)
w1_scaled, w2_scaled = w1_scaled.squeeze(0), w2_scaled.squeeze(0)
fused_out = torch.randn(M, K, dtype=dtype) / math.sqrt(K)
a2 = a.clone()
fused_output = (
torch.randn(m, k, dtype=dtype) / math.sqrt(k)
if apply_scaling_factor
else None
)
routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None
hidden_states2 = hidden_states.clone()
# ref
ic0 = torch.matmul(a.float(), w1_scaled.transpose(0, 1))
ic1 = SiluAndMul(ic0)
shared_out = torch.matmul(ic1, w2_scaled.transpose(0, 1))
ref_out = shared_out + fused_out.float() * routed_scaling_factor
ref_out = ref_out.to(dtype=dtype)
# ref with bfloat16
ref = torch_naive_moe(
hidden_states,
w1_scaled,
w2_scaled,
fused_output,
routed_scaling_factor,
output_dtype=dtype,
)
w1 = torch.ops.sgl_kernel.convert_weight_packed(w1) # [2N, K]
w2 = torch.ops.sgl_kernel.convert_weight_packed(w2) # [K, N]
out = torch.ops.sgl_kernel.shared_expert_cpu(
a2,
hidden_states2,
w1,
w2,
fused_out,
fused_output,
routed_scaling_factor,
True,
False,
@@ -195,8 +207,8 @@ class TestSharedExpert(CustomTestCase):
True,
)
atol = rtol = precision[ref_out.dtype]
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
def test_fp8_shared_expert(self):
for params in itertools.product(
@@ -204,12 +216,14 @@ class TestSharedExpert(CustomTestCase):
self.N_fp8,
self.K_fp8,
self.routed_scaling_factor,
self.apply_scaling_factor,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
m=params[0],
n=params[1],
k=params[2],
routed_scaling_factor=params[3],
apply_scaling_factor=params[4],
):
self._fp8_shared_expert(*params)
test/srt/cpu/utils.py
+18 -4
View File
@@ -126,16 +126,28 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, bias, output_dtype=torch.bfloat16
return C.reshape(origin_C_shape).to(output_dtype)
def torch_naive_moe(a, w1, w2, b, routed_scaling_factor):
def torch_naive_moe(a, w1, w2, b, routed_scaling_factor, output_dtype=torch.bfloat16):
a = a.to(torch.float32)
w1 = w1.to(torch.float32)
w2 = w2.to(torch.float32)
b = b.to(torch.float32) if b is not None else None
ic1 = torch.matmul(a, w1.transpose(0, 1))
ic2 = SiluAndMul(ic1)
ic3 = torch.matmul(ic2, w2.transpose(0, 1))
return ic3 + b * routed_scaling_factor
out = ic3 if b is None else ic3 + b * routed_scaling_factor
return out.to(output_dtype)
def torch_w8a8_per_column_moe(a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor):
def torch_w8a8_per_column_moe(
a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor, output_dtype=torch.bfloat16
):
a = a.to(torch.float32)
b = b.to(torch.float32) if b is not None else None
# Perform per-token quantization
a_q, a_s = per_token_quant_int8(a)
@@ -150,7 +162,9 @@ def torch_w8a8_per_column_moe(a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_facto
a1_q, w2_q, a1_s, w2_s, bias=None, output_dtype=torch.float32
)
return ic3 + b * routed_scaling_factor
out = ic3 if b is None else ic3 + b * routed_scaling_factor
return out.to(output_dtype)
def scaled_weight(weight, scales):