372 lines
11 KiB
Python
372 lines
11 KiB
Python
import itertools
|
|
import math
|
|
import unittest
|
|
|
|
# TODO: use interface in cpu.py
|
|
import torch
|
|
|
|
from sglang.srt.layers.amx_utils import CPUQuantMethod
|
|
|
|
kernel = torch.ops.sgl_kernel
|
|
|
|
torch.manual_seed(128)
|
|
|
|
from utils import (
|
|
BLOCK_K,
|
|
BLOCK_N,
|
|
factor_for_scale,
|
|
fp8_max,
|
|
fp8_min,
|
|
native_fp8_fused_moe,
|
|
precision,
|
|
scaled_weight,
|
|
torch_naive_fused_moe,
|
|
torch_w8a8_per_column_fused_moe,
|
|
unpack_and_dequant_awq,
|
|
)
|
|
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
|
|
def fused_moe(a, w1, w2, score, topk, renormalize, prepack):
|
|
|
|
G = 1
|
|
topk_group = 1
|
|
|
|
B, D = a.shape
|
|
topk_weights = torch.empty(B, topk, dtype=torch.float32)
|
|
topk_ids = torch.empty(B, topk, dtype=torch.int32)
|
|
topk_weights, topk_ids = kernel.grouped_topk_cpu(
|
|
a, score, topk, renormalize, G, topk_group, 0, None, None
|
|
)
|
|
|
|
packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
|
|
packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
|
|
|
|
inplace = True
|
|
return kernel.fused_experts_cpu(
|
|
a,
|
|
packed_w1,
|
|
packed_w2,
|
|
topk_weights,
|
|
topk_ids,
|
|
inplace,
|
|
CPUQuantMethod.UNQUANT,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
prepack,
|
|
)
|
|
|
|
|
|
class TestFusedExperts(CustomTestCase):
|
|
M = [2, 114]
|
|
N = [32]
|
|
K = [32]
|
|
E = [4]
|
|
topk = [2]
|
|
renormalize = [False, True]
|
|
|
|
M_int8 = [1, 39]
|
|
N_int8 = [128]
|
|
K_int8 = [256]
|
|
E_int8 = [8]
|
|
topk_int8 = [3]
|
|
|
|
M_fp8 = [2, 121]
|
|
N_fp8 = [352, 512]
|
|
K_fp8 = [256, 320]
|
|
E_fp8 = [8]
|
|
topk_fp8 = [4]
|
|
|
|
M_int4 = [1, 6]
|
|
N_int4 = [512]
|
|
K_int4 = [256]
|
|
E_int4 = [8]
|
|
topk_int4 = [4]
|
|
|
|
def _bf16_moe(self, m, n, k, e, topk, renormalize):
|
|
dtype = torch.bfloat16
|
|
prepack = True
|
|
|
|
a = torch.randn((m, k), device="cpu", dtype=dtype) / 10
|
|
w1 = torch.randn((e, 2 * n, k), device="cpu", dtype=dtype) / 10
|
|
w2 = torch.randn((e, k, n), device="cpu", dtype=dtype) / 10
|
|
score = torch.randn((m, e), device="cpu", dtype=dtype)
|
|
|
|
torch_output = torch_naive_fused_moe(a, w1, w2, score, topk, renormalize)
|
|
fused_output = fused_moe(a, w1, w2, score, topk, renormalize, prepack)
|
|
|
|
atol = rtol = precision[torch_output.dtype]
|
|
torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
|
|
|
|
def test_bf16_moe(self):
|
|
for params in itertools.product(
|
|
self.M,
|
|
self.N,
|
|
self.K,
|
|
self.E,
|
|
self.topk,
|
|
self.renormalize,
|
|
):
|
|
with self.subTest(
|
|
m=params[0],
|
|
n=params[1],
|
|
k=params[2],
|
|
e=params[3],
|
|
topk=params[4],
|
|
renormalize=params[5],
|
|
):
|
|
self._bf16_moe(*params)
|
|
|
|
def _int8_moe(self, M, N, K, E, topk):
|
|
dtype = torch.bfloat16
|
|
prepack = True
|
|
|
|
# Initialize int8 quantization parameters
|
|
int8_factor_for_scale = 1e-2
|
|
int8_max = 127
|
|
int8_min = -128
|
|
|
|
# Input tensor
|
|
# M * K
|
|
a = torch.randn((M, K), dtype=dtype) / math.sqrt(K)
|
|
|
|
# Generate int8 weights
|
|
w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
|
|
w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
|
|
|
|
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
|
|
w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
|
|
|
|
# Generate scale for each column (per-column quantization)
|
|
w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * int8_factor_for_scale
|
|
w2_s = torch.rand(E, K, device=w2_fp32.device) * int8_factor_for_scale
|
|
|
|
# Calculate routing
|
|
score = torch.randn((M, E), dtype=dtype)
|
|
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
|
topk_weight, topk_ids = torch.topk(score, topk)
|
|
|
|
ref_out = torch_w8a8_per_column_fused_moe(
|
|
a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, topk
|
|
)
|
|
|
|
inplace = True
|
|
packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
|
|
packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
|
|
out = kernel.fused_experts_cpu(
|
|
a,
|
|
packed_w1,
|
|
packed_w2,
|
|
topk_weight,
|
|
topk_ids.to(torch.int32),
|
|
inplace,
|
|
CPUQuantMethod.INT8_W8A8,
|
|
w1_s,
|
|
w2_s,
|
|
None,
|
|
None,
|
|
None,
|
|
prepack,
|
|
)
|
|
|
|
atol = rtol = precision[ref_out.dtype]
|
|
# Increase the tolerance for large input shapes
|
|
if M > 35:
|
|
atol = rtol = 0.02
|
|
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
|
|
|
|
def test_int8_moe(self):
|
|
for params in itertools.product(
|
|
self.M_int8,
|
|
self.N_int8,
|
|
self.K_int8,
|
|
self.E_int8,
|
|
self.topk_int8,
|
|
):
|
|
with self.subTest(
|
|
M=params[0],
|
|
N=params[1],
|
|
K=params[2],
|
|
E=params[3],
|
|
topk=params[4],
|
|
):
|
|
self._int8_moe(*params)
|
|
|
|
def _fp8_moe(self, M, N, K, E, topk):
|
|
dtype = torch.bfloat16
|
|
|
|
a = torch.randn(M, K, dtype=dtype) / math.sqrt(K)
|
|
|
|
w1_fp32 = torch.randn(E, 2 * N, K)
|
|
w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
|
|
|
|
w2_fp32 = torch.randn(E, K, N)
|
|
w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
|
|
|
|
w1s = (
|
|
torch.randn(E, math.ceil(2 * N / BLOCK_N), math.ceil(K / BLOCK_K))
|
|
* factor_for_scale
|
|
)
|
|
w2s = (
|
|
torch.randn(E, math.ceil(K / BLOCK_N), math.ceil(N / BLOCK_K))
|
|
* factor_for_scale
|
|
)
|
|
|
|
w1_scaled = scaled_weight(w1, w1s)
|
|
w2_scaled = scaled_weight(w2, w2s)
|
|
|
|
score = torch.randn((M, E), dtype=dtype)
|
|
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
|
topk_weight, topk_ids = torch.topk(score, topk)
|
|
|
|
w1 = kernel.convert_weight_packed(w1)
|
|
w2 = kernel.convert_weight_packed(w2)
|
|
|
|
ref_out = native_fp8_fused_moe(
|
|
a, w1_scaled, w2_scaled, topk_weight, topk_ids, topk
|
|
)
|
|
out = kernel.fused_experts_cpu(
|
|
a,
|
|
w1,
|
|
w2,
|
|
topk_weight,
|
|
topk_ids.to(torch.int32),
|
|
False,
|
|
CPUQuantMethod.FP8_W8A16,
|
|
w1s,
|
|
w2s,
|
|
None,
|
|
None,
|
|
[BLOCK_N, BLOCK_K],
|
|
True,
|
|
)
|
|
|
|
atol = rtol = precision[dtype]
|
|
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
|
|
|
|
def test_fp8_moe(self):
|
|
for params in itertools.product(
|
|
self.M_fp8,
|
|
self.N_fp8,
|
|
self.K_fp8,
|
|
self.E_fp8,
|
|
self.topk_fp8,
|
|
):
|
|
with self.subTest(
|
|
M=params[0],
|
|
N=params[1],
|
|
K=params[2],
|
|
E=params[3],
|
|
topk=params[4],
|
|
):
|
|
self._fp8_moe(*params)
|
|
|
|
def _int4_moe(self, M, N, K, E, topk, group_size=128):
|
|
dtype = torch.bfloat16
|
|
|
|
a = torch.rand(M, K, dtype=dtype) / math.sqrt(K)
|
|
|
|
awq_w13_weight = torch.randint(-127, 128, (E, K, 2 * N // 8)).to(torch.int)
|
|
awq_w13_zero = torch.randint(0, 10, (E, K // group_size, 2 * N // 8)).to(
|
|
torch.int
|
|
)
|
|
awq_w13_scales = torch.rand(E, int(K // group_size), 2 * N).to(torch.bfloat16)
|
|
|
|
awq_w2_weight = torch.randint(-127, 128, (E, N, K // 8)).to(torch.int)
|
|
awq_w2_zero = torch.randint(0, 10, (E, N // group_size, K // 8)).to(torch.int)
|
|
awq_w2_scales = torch.rand(E, int(N // group_size), K).to(torch.bfloat16)
|
|
bf16_w13_weight = []
|
|
bf16_w2_weight = []
|
|
for i in range(E):
|
|
bf16_w13_weight_i, _ = unpack_and_dequant_awq(
|
|
awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i], 4, 128
|
|
)
|
|
bf16_w2_weight_i, _ = unpack_and_dequant_awq(
|
|
awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i], 4, 128
|
|
)
|
|
bf16_w13_weight.append(bf16_w13_weight_i)
|
|
bf16_w2_weight.append(bf16_w2_weight_i)
|
|
bf16_w13_weight = torch.stack(bf16_w13_weight).detach()
|
|
bf16_w2_weight = torch.stack(bf16_w2_weight).detach()
|
|
|
|
score = torch.rand((M, E), dtype=dtype)
|
|
|
|
ref_out = torch_naive_fused_moe(
|
|
a, bf16_w13_weight, bf16_w2_weight, score, topk, False
|
|
)
|
|
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.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.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,
|
|
awq_w13_weight_pack,
|
|
awq_w2_weight_pack,
|
|
topk_weight,
|
|
topk_ids.to(torch.int32),
|
|
False,
|
|
CPUQuantMethod.INT4_W4A8,
|
|
awq_w13_scales_pack,
|
|
awq_w2_scales_pack,
|
|
awq_w13_zero_pack,
|
|
awq_w2_zero_pack,
|
|
None,
|
|
True,
|
|
)
|
|
|
|
atol = rtol = precision[dtype]
|
|
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
|
|
|
|
def test_int4_moe(self):
|
|
for params in itertools.product(
|
|
self.M_int4,
|
|
self.N_int4,
|
|
self.K_int4,
|
|
self.E_int4,
|
|
self.topk_int4,
|
|
):
|
|
with self.subTest(
|
|
M=params[0],
|
|
N=params[1],
|
|
K=params[2],
|
|
E=params[3],
|
|
topk=params[4],
|
|
):
|
|
self._int4_moe(*params)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|