mirrors/vllm-project--vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
pre-commit / pre-run-check (push) Has been cancelled
Details
pre-commit / pre-commit (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:55:37 +08:00
commit 7ce4c8e27e
5900 changed files with 1668062 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_cpu_quant_fused_moe.py
+888
View File
@@ -0,0 +1,888 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for CPU quantized fused MoE kernels."""
import math
import sys
import pytest
import torch
import torch.nn.functional as F
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
if not current_platform.is_cpu():
pytest.skip("skipping CPU-only tests", allow_module_level=True)
import vllm._custom_ops as ops # noqa: E402
if not hasattr(torch.ops._C, "fused_experts_cpu"):
pytest.skip("fused_experts_cpu op not available", allow_module_level=True)
def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
d = x.shape[-1] // 2
return F.silu(x[..., :d]) * x[..., d:]
def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
"""VNNI-prepack expert weights via ``convert_weight_packed``."""
return torch.ops._C.convert_weight_packed(w)
# ===========================================================================
# FP8 W8A16 MoE
# ===========================================================================
BLOCK_SIZE = [128, 128] # [block_n, block_k]
_FP8_INFO = torch.finfo(torch.float8_e4m3fn)
FP8_SCALE = _FP8_INFO.max # 448.0
FACTOR_FOR_SCALE = 1e-3
def _block_dequant_weight(
weight: torch.Tensor,
scales: torch.Tensor,
block_size: list[int],
) -> torch.Tensor:
"""Block-dequantize FP8 weight [E, N, K] -> float [E, N, K]."""
E, N, K = weight.shape
block_n, block_k = block_size
pad_N = (block_n - N % block_n) % block_n
pad_K = (block_k - K % block_k) % block_k
if pad_N > 0 or pad_K > 0:
weight = F.pad(weight, (0, pad_K, 0, pad_N))
n_tiles = math.ceil(N / block_n)
k_tiles = math.ceil(K / block_k)
weight_block = (
weight.view(E, n_tiles, block_n, k_tiles, block_k)
.permute(0, 1, 3, 2, 4)
.float()
.contiguous()
)
weight_scaled = (
(weight_block * scales.view(E, n_tiles, k_tiles, 1, 1))
.permute(0, 1, 3, 2, 4)
.contiguous()
)
if pad_N > 0 or pad_K > 0:
weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
weight_scaled = weight_scaled[..., :N, :K].contiguous()
else:
weight_scaled = weight_scaled.view(E, N, K)
return weight_scaled
def ref_w8a16_block_fp8_moe(
a: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
w1_s: torch.Tensor,
w2_s: torch.Tensor,
topk_weight: torch.Tensor,
topk_ids: torch.Tensor,
block_size: list[int],
) -> torch.Tensor:
"""Reference FP8 W8A16 block-scaled fused MoE in pure torch."""
B, D = a.shape
topk = topk_ids.size(1)
w1_dq = _block_dequant_weight(w1, w1_s, block_size)
w2_dq = _block_dequant_weight(w2, w2_s, block_size)
a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
topk_weight_flat = topk_weight.view(-1)
topk_ids_flat = topk_ids.view(-1)
for i in range(w1_dq.shape[0]):
mask = topk_ids_flat == i
if mask.sum():
ic0 = torch.matmul(a_exp[mask], w1_dq[i].transpose(0, 1))
ic1 = _silu_and_mul(ic0)
out[mask] = torch.matmul(ic1, w2_dq[i].transpose(0, 1))
return (
(out.view(B, -1, w2_dq.shape[1]) * topk_weight_flat.view(B, -1, 1))
.sum(dim=1)
.to(a.dtype)
)
def _make_fp8_moe_weights(
E: int,
N: int,
K: int,
block_size: list[int],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Generate random FP8 MoE weights with random block scales."""
block_n, block_k = block_size
w1 = (
(torch.randn(E, 2 * N, K) * FP8_SCALE)
.clamp(min=-FP8_SCALE, max=FP8_SCALE)
.to(torch.float8_e4m3fn)
)
w2 = (
(torch.randn(E, K, N) * FP8_SCALE)
.clamp(min=-FP8_SCALE, max=FP8_SCALE)
.to(torch.float8_e4m3fn)
)
w1_s = (
torch.randn(E, math.ceil(2 * N / block_n), math.ceil(K / block_k))
* FACTOR_FOR_SCALE
)
w2_s = (
torch.randn(E, math.ceil(K / block_n), math.ceil(N / block_k))
* FACTOR_FOR_SCALE
)
return w1, w2, w1_s, w2_s
FP8_NUM_TOKENS = [1, 2, 64, 121]
FP8_MOE_CONFIGS = [
(256, 512, 8, 2),
(256, 512, 8, 4),
(512, 256, 8, 2),
(512, 256, 8, 4),
(512, 512, 8, 2),
(512, 512, 8, 4),
(768, 2048, 8, 2),
(768, 2048, 8, 4),
(768, 2048, 128, 8),
]
@pytest.mark.parametrize("M", FP8_NUM_TOKENS)
@pytest.mark.parametrize("N,K,E,topk", FP8_MOE_CONFIGS)
@pytest.mark.parametrize("seed", [0])
def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
"""Test fused_experts_cpu FP8 W8A16 against dequantised torch reference."""
set_random_seed(seed)
a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)
score = torch.randn(M, E, dtype=torch.bfloat16)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_ids = topk_ids.to(torch.int32)
ref_out = ref_w8a16_block_fp8_moe(
a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, BLOCK_SIZE
)
pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
# Test inplace=False against reference
out = ops.fused_experts_cpu(
a.clone(),
pw1,
pw2,
topk_weight,
topk_ids,
False,
ops.CPUQuantMethod.FP8_W8A16,
w1_s,
w2_s,
None,
None,
BLOCK_SIZE,
is_vnni=True,
)
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
# Test inplace=True produces identical output
out_inplace = ops.fused_experts_cpu(
a.clone(),
pw1,
pw2,
topk_weight,
topk_ids,
True,
ops.CPUQuantMethod.FP8_W8A16,
w1_s,
w2_s,
None,
None,
BLOCK_SIZE,
is_vnni=True,
)
torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
# ===========================================================================
# MXFP4 W4A16 MoE
# ===========================================================================
class MXFP4QuantizeUtil:
"""MXFP4 quantization utility."""
E2M1_max = 6.0
E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
block_size = 32
@classmethod
def quantize(cls, input: torch.Tensor) -> tuple:
"""Quantize BF16 tensor to MXFP4 packed uint8 format."""
def cast_fp4(x):
sign = torch.sign(x)
sign_bit = (2 - sign) // 2
ord_ = torch.sum(
(x.abs().unsqueeze(-1) - cls.E2M1_bounds.to(x.device)) > 0, dim=-1
)
fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8)
return fp4_val
def fuse_uint4_to_uint8(x):
left_side = x[..., 0::2]
right_side = x[..., 1::2]
new_data = right_side.clone() << 4
new_data[..., : left_side.shape[-1]] += left_side
return new_data
original_shape = input.shape
input = input.view(-1, cls.block_size)
input_amax = input.abs().max(dim=-1, keepdim=True).values
descale = input_amax / cls.E2M1_max
min_value = torch.tensor(-127.0, device=descale.device)
e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value))
input = (input / torch.exp2(e8m0_scale)).view(original_shape)
input_q = cast_fp4(input)
input_q = fuse_uint4_to_uint8(input_q)
e8m0_scale = (e8m0_scale + 127).to(torch.uint8)
return input_q, e8m0_scale
@classmethod
def dequantize(cls, quantized_data, dtype: torch.dtype, scale):
"""Dequantize MXFP4 packed tensor back to float."""
def unfuse_uint8_to_uint4(x):
left_side = x & 0x0F
right_side = (x >> 4) & 0x0F
shape = list(x.shape)
shape[-1] = shape[-1] * 2
result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
result[..., 0::2] = left_side
result[..., 1::2] = right_side
return result
e8m0_scale = scale
x_unfused = unfuse_uint8_to_uint4(quantized_data)
sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(torch.float32)
magnitude = (x_unfused & 0b0111).to(torch.long)
values = torch.tensor(cls.E2M1_values, device=quantized_data.device)
original_shape = magnitude.shape
x_float = values[magnitude.reshape(-1)].reshape(original_shape)
x_float = sign.float() * x_float
x_float = x_float.reshape(-1, cls.block_size)
scale_factor = torch.exp2(e8m0_scale.float() - 127)
scale_factor = scale_factor.reshape(-1, 1)
x_float = x_float * scale_factor
return x_float.reshape(original_shape).to(dtype)
def _swiglu(x: torch.Tensor, alpha: float, limit: float) -> torch.Tensor:
"""SwigLU activation used in GPT-OSS.
Input is interleaved: [gate_0, up_0, gate_1, up_1, ...] in last dim.
"""
gate = x[..., 0::2]
up = x[..., 1::2]
gate_clamped = torch.clamp(gate, max=limit)
up_clamped = torch.clamp(up, min=-limit, max=limit)
return gate_clamped * torch.sigmoid(alpha * gate_clamped) * (up_clamped + 1)
def ref_mxfp4_fused_moe(
a: torch.Tensor,
w1_dq: torch.Tensor,
w2_dq: torch.Tensor,
topk_weight: torch.Tensor,
topk_ids: torch.Tensor,
topk: int,
) -> torch.Tensor:
"""Reference MXFP4 fused MoE with SiLU activation."""
B, D = a.shape
a_f = a.float()
out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
topk_ids_flat = topk_ids.view(-1)
for i in range(w1_dq.shape[0]):
mask = topk_ids_flat == i
if mask.sum() == 0:
continue
token_indices = torch.where(mask)[0]
source_indices = token_indices // topk
ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
ic1 = _silu_and_mul(ic0)
out[mask] = torch.matmul(ic1, w2_dq[i].float().T)
return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
def ref_mxfp4_fused_moe_gptoss(
a: torch.Tensor,
w1_dq: torch.Tensor,
w2_dq: torch.Tensor,
w1_bias: torch.Tensor,
w2_bias: torch.Tensor,
topk_weight: torch.Tensor,
topk_ids: torch.Tensor,
alpha: float,
limit: float,
) -> torch.Tensor:
"""Reference MXFP4 fused MoE with SwigLU+bias (GPT-OSS style)."""
B, D = a.shape
topk = topk_ids.shape[1]
a_f = a.float()
E = w1_dq.shape[0]
out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
topk_ids_flat = topk_ids.view(-1)
for i in range(E):
mask = topk_ids_flat == i
if mask.sum() == 0:
continue
token_indices = torch.where(mask)[0]
source_indices = token_indices // topk
ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
ic0 = ic0 + w1_bias[i].float()
ic1 = _swiglu(ic0, alpha, limit)
ic2 = torch.matmul(ic1, w2_dq[i].float().T)
ic2 = ic2 + w2_bias[i].float()
out[mask] = ic2
return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
def _prepack_mxfp4_experts(
w: torch.Tensor, w_scale: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""VNNI-prepack MXFP4 weights and repack scales."""
packed_w = torch.ops._C.convert_weight_packed(w)
packed_s = torch.ops._C.convert_scale_packed(w_scale)
return packed_w, packed_s
MXFP4_NUM_TOKENS = [1, 2, 32, 121]
MXFP4_MOE_CONFIGS = [
(128, 128, 4, 2),
(256, 256, 8, 4),
(352, 256, 8, 4),
(512, 320, 8, 4),
]
@pytest.mark.parametrize("M", MXFP4_NUM_TOKENS)
@pytest.mark.parametrize("N,K,E,topk", MXFP4_MOE_CONFIGS)
@pytest.mark.parametrize("seed", [0])
def test_mxfp4_cpu_fused_moe(M, N, K, E, topk, seed):
"""Test fused_experts_mxfp4_cpu against dequantized torch reference."""
set_random_seed(seed)
dtype = torch.bfloat16
a = torch.randn(M, K, dtype=dtype) / 10
# Generate and quantize weights
w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
w1s = w1s.reshape(E, 2 * N, K // 32)
w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
w2s = w2s.reshape(E, K, N // 32)
w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
# 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)
topk_ids = topk_ids.to(torch.int32)
# Reference
ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, topk)
# Pack weights for kernel
pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
# Kernel
out = ops.fused_experts_cpu(
a.clone(),
pw1,
pw2,
topk_weight,
topk_ids,
False, # inplace
ops.CPUQuantMethod.MXFP4,
pw1s, # w1_scale
pw2s, # w2_scale
None, # w1_zero
None, # w2_zero
None, # block_size
)
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
@pytest.mark.parametrize("M", [1, 32])
@pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)])
@pytest.mark.parametrize("seed", [0])
def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
"""Test fused_experts_mxfp4_cpu with bias and SwigLU activation (GPT-OSS)."""
set_random_seed(seed)
dtype = torch.bfloat16
alpha = 1.702
limit = 7.0
a = torch.randn(M, K, dtype=dtype) / 10
# Generate and quantize weights
w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
w1s = w1s.reshape(E, 2 * N, K // 32)
w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
w1_b = torch.randn(E, 2 * N, dtype=torch.float32) / 10
w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
w2s = w2s.reshape(E, K, N // 32)
w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
w2_b = torch.randn(E, K, dtype=torch.float32) / 10
# 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)
topk_ids = topk_ids.to(torch.int32)
# Reference
ref_out = ref_mxfp4_fused_moe_gptoss(
a, w1dq, w2dq, w1_b, w2_b, topk_weight, topk_ids, alpha, limit
)
# Pack weights for kernel
pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
# Kernel
out = ops.fused_experts_cpu(
a.clone(),
pw1,
pw2,
topk_weight,
topk_ids,
False, # inplace
ops.CPUQuantMethod.MXFP4,
pw1s, # w1_scale
pw2s, # w2_scale
None, # w1_zero
None, # w2_zero
None, # block_size
w1_bias=w1_b,
w2_bias=w2_b,
alpha=alpha,
limit=limit,
)
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
# ===========================================================================
# INT4 W4A16 MoE
# ===========================================================================
def _pack_int4_gptq(w_int4: torch.Tensor) -> torch.Tensor:
"""Pack INT4 values [N, K] → [N, K//8] int32 along K dim (GPTQ format)."""
N, K = w_int4.shape
assert K % 8 == 0
w = w_int4.to(torch.int32)
w_packed = torch.zeros(N, K // 8, dtype=torch.int32)
for j in range(8):
w_packed |= (w[:, j::8] & 0xF) << (j * 4)
return w_packed
def _pack_int4_awq(w_int4: torch.Tensor) -> torch.Tensor:
"""Pack INT4 values [..., N] → [..., N//8] int32 along last dim (AWQ format)."""
# AWQ packing bitshifts: indices {0,4,1,5,2,6,3,7} * 4 bits each
_AWQ_BITSHIFTS = [0, 16, 4, 20, 8, 24, 12, 28]
N = w_int4.shape[-1]
assert N % 8 == 0
w = w_int4.to(torch.int32)
w_packed = torch.zeros(*w.shape[:-1], N // 8, dtype=torch.int32)
for j, shift in enumerate(_AWQ_BITSHIFTS):
w_packed |= (w[..., j::8] & 0xF) << shift
return w_packed
def _ref_int4_moe(
a: torch.Tensor,
w1_int4: torch.Tensor,
w2_int4: torch.Tensor,
w1_zeros: torch.Tensor | None,
w2_zeros: torch.Tensor | None,
w1_s: torch.Tensor,
w2_s: torch.Tensor,
topk_weight: torch.Tensor,
topk_ids: torch.Tensor,
group_size: int,
) -> torch.Tensor:
"""Reference INT4 W4A16 group-quantized fused MoE in pure torch."""
B = a.shape[0]
topk = topk_ids.size(1)
K_out = a.shape[1]
out = torch.zeros(B, topk, K_out, dtype=torch.float32)
for b in range(B):
for t in range(topk):
eid = topk_ids[b, t].item()
x = a[b : b + 1].float()
# Dequantize w1: [K, 2*N], groups along K (input dim)
K_dim = w1_int4.shape[1]
w1_dq = torch.zeros(K_dim, w1_int4.shape[2], dtype=torch.float32)
for g in range(w1_s.shape[1]):
k_start = g * group_size
k_end = min((g + 1) * group_size, K_dim)
zp = w1_zeros[eid, g, :].float() if w1_zeros is not None else 8.0
w1_dq[k_start:k_end, :] = (
w1_int4[eid, k_start:k_end, :].float() - zp
) * w1_s[eid, g, :].float()
ic = torch.matmul(x, w1_dq) # [1, K] @ [K, 2*N] → [1, 2*N]
ic = _silu_and_mul(ic) # [1, N]
# Dequantize w2: [N, K], groups along N (input dim)
N_dim = w2_int4.shape[1]
w2_dq = torch.zeros(N_dim, w2_int4.shape[2], dtype=torch.float32)
for g in range(w2_s.shape[1]):
n_start = g * group_size
n_end = min((g + 1) * group_size, N_dim)
zp = w2_zeros[eid, g, :].float() if w2_zeros is not None else 8.0
w2_dq[n_start:n_end, :] = (
w2_int4[eid, n_start:n_end, :].float() - zp
) * w2_s[eid, g, :].float()
oc = torch.matmul(ic, w2_dq) # [1, N] @ [N, K] → [1, K]
out[b, t] = oc.squeeze(0)
return (out * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
def _make_int4_moe_weights(E, N, K, group_size, quant_algo):
"""Create INT4 MoE weights in GPTQ or AWQ packed format.
Canonical layout (input × output):
w1_int4: [E, K, 2*N] w2_int4: [E, N, K]
GPTQ packed (pack transposed weight along input/K dim):
w1_packed: [E, K//8, 2*N] w2_packed: [E, N//8, K]
zeros: actual int4 zero points, same packing as weights
AWQ packed (pack along output/N dim):
w1_packed: [E, K, 2*N//8] w2_packed: [E, N, K//8]
zeros: actual int4 zero points, same packing as weights
Returns:
w1_int4, w2_int4,
w1_packed, w2_packed,
w1_zeros, w2_zeros,
w1_zeros_packed, w2_zeros_packed,
w1_s, w2_s
"""
w1_int4 = torch.randint(0, 16, (E, K, 2 * N), dtype=torch.int32)
w2_int4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32)
num_groups_w1 = K // group_size
num_groups_w2 = N // group_size
w1_s = (
torch.randn(E, num_groups_w1, 2 * N, dtype=torch.bfloat16) * 0.01
).abs() + 0.001
w2_s = (torch.randn(E, num_groups_w2, K, dtype=torch.bfloat16) * 0.01).abs() + 0.001
if quant_algo == ops.CPUQuantAlgo.GPTQ:
# Pack: canonical [E, K, 2*N] → transpose [E, 2*N, K] → GPTQ pack
# [E, 2*N, K//8] → transpose [E, K//8, 2*N]
w1_t = w1_int4.transpose(1, 2).contiguous() # [E, 2*N, K]
w1_packed = (
torch.stack([_pack_int4_gptq(w1_t[e]) for e in range(E)])
.transpose(1, 2)
.contiguous()
) # [E, K//8, 2*N]
w2_t = w2_int4.transpose(1, 2).contiguous() # [E, K, N]
w2_packed = (
torch.stack([_pack_int4_gptq(w2_t[e]) for e in range(E)])
.transpose(1, 2)
.contiguous()
) # [E, N//8, K]
w1_zeros = w2_zeros = None
w1_zeros_packed = torch.full(
(E, num_groups_w1, 2 * N // 8), 0x77777777, dtype=torch.int32
)
w2_zeros_packed = torch.full(
(E, num_groups_w2, K // 8), 0x77777777, dtype=torch.int32
)
else: # AWQ
# Asymmetric: actual zero points, packed along output dim.
w1_zeros = torch.randint(1, 15, (E, num_groups_w1, 2 * N), dtype=torch.int32)
w2_zeros = torch.randint(1, 15, (E, num_groups_w2, K), dtype=torch.int32)
w1_packed = torch.stack(
[_pack_int4_awq(w1_int4[e]) for e in range(E)]
) # [E, K, 2*N//8]
w2_packed = torch.stack(
[_pack_int4_awq(w2_int4[e]) for e in range(E)]
) # [E, N, K//8]
w1_zeros_packed = torch.stack(
[_pack_int4_awq(w1_zeros[e]) for e in range(E)]
) # [E, K//gs, 2*N//8]
w2_zeros_packed = torch.stack(
[_pack_int4_awq(w2_zeros[e]) for e in range(E)]
) # [E, N//gs, K//8]
return (
w1_int4,
w2_int4,
w1_packed,
w2_packed,
w1_zeros,
w2_zeros,
w1_zeros_packed,
w2_zeros_packed,
w1_s,
w2_s,
)
INT4_MOE_CONFIGS = [
# (N, K, E, topk, group_size)
(256, 512, 8, 2, 128),
(512, 256, 8, 2, 128),
(512, 512, 8, 4, 128),
(768, 2048, 8, 2, 128),
]
@pytest.mark.parametrize("M", [1, 2, 64, 121])
@pytest.mark.parametrize("N,K,E,topk,group_size", INT4_MOE_CONFIGS)
@pytest.mark.parametrize("quant_algo", [ops.CPUQuantAlgo.GPTQ, ops.CPUQuantAlgo.AWQ])
@pytest.mark.parametrize("seed", [0])
def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed):
"""Test fused_experts_cpu INT4 W4A16 for both GPTQ and AWQ quant formats."""
set_random_seed(seed)
a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
(
w1_int4,
w2_int4,
w1_packed,
w2_packed,
w1_zeros,
w2_zeros,
w1_zeros_packed,
w2_zeros_packed,
w1_s,
w2_s,
) = _make_int4_moe_weights(E, N, K, group_size, quant_algo)
score = torch.randn(M, E, dtype=torch.bfloat16)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_ids = topk_ids.to(torch.int32)
ref_out = _ref_int4_moe(
a,
w1_int4,
w2_int4,
w1_zeros,
w2_zeros,
w1_s,
w2_s,
topk_weight,
topk_ids,
group_size,
)
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
prepare_int4_moe_layer_for_cpu,
)
(blocked_w1, blocked_w2, blocked_s1, blocked_s2, blocked_z1, blocked_z2) = (
prepare_int4_moe_layer_for_cpu(
w1_packed,
w2_packed,
w1_s,
w2_s,
quant_algo=quant_algo,
w13_zeros=w1_zeros_packed,
w2_zeros=w2_zeros_packed,
)
)
out = ops.fused_experts_cpu(
a.clone(),
blocked_w1,
blocked_w2,
topk_weight,
topk_ids,
False, # inplace
ops.CPUQuantMethod.INT4_W4A8,
blocked_s1,
blocked_s2,
blocked_z1,
blocked_z2,
None, # block_size
None, # w1_bias
None, # w2_bias
None, # alpha
None, # limit
True, # is_vnni
)
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
# ===========================================================================
# INT8 W8A8 MoE
# ===========================================================================
def _quantize_per_channel(w):
"""Symmetric per-channel INT8 quantisation. w: [N, K] -> (int8, scale)."""
amax = w.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
scale = amax / 127.0
w_q = (w / scale).round().clamp(-128, 127).to(torch.int8)
return w_q, scale.float()
def _quantize_per_token(x):
"""Symmetric per-token INT8 quantisation. x: [M, K] -> (int8, scale)."""
amax = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
scale = amax / 127.0
x_q = (x / scale).round().clamp(-128, 127).to(torch.int8)
return x_q, scale.float()
def _ref_int8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids):
"""Reference INT8 W8A8 per-channel fused MoE in pure torch."""
B, D = a.shape
topk = topk_ids.size(1)
out = torch.zeros(B, topk, w2.shape[1], dtype=torch.float32)
for b in range(B):
for t in range(topk):
eid = topk_ids[b, t].item()
x = a[b : b + 1].float()
x_q, x_s = _quantize_per_token(x)
ic = torch.matmul(x_q.float(), w1[eid].float().t())
ic = ic * x_s * w1_s[eid].view(1, -1)
ic = _silu_and_mul(ic)
ic_q, ic_s = _quantize_per_token(ic)
oc = torch.matmul(ic_q.float(), w2[eid].float().t())
oc = oc * ic_s * w2_s[eid].view(1, -1)
out[b, t] = oc.squeeze(0)
result = (out * topk_weight.unsqueeze(-1)).sum(dim=1)
return result.to(a.dtype)
def _make_int8_moe_weights(E, N, K):
factor = 1e-2
w1_f = (torch.randn(E, 2 * N, K) - 0.5) * 2
w2_f = (torch.randn(E, K, N) - 0.5) * 2
w1_q_list, w1_s_list = [], []
w2_q_list, w2_s_list = [], []
for e in range(E):
q, s = _quantize_per_channel(w1_f[e])
w1_q_list.append(q)
w1_s_list.append(s)
q, s = _quantize_per_channel(w2_f[e])
w2_q_list.append(q)
w2_s_list.append(s)
return (
torch.stack(w1_q_list),
torch.stack(w2_q_list),
torch.stack(w1_s_list) * factor,
torch.stack(w2_s_list) * factor,
)
INT8_NUM_TOKENS = [1, 2, 64, 121]
INT8_MOE_CONFIGS = [
# (N, K, E, topk)
(256, 512, 8, 2),
(512, 256, 8, 2),
(512, 512, 8, 4),
(768, 2048, 8, 2),
]
@pytest.mark.parametrize("M", INT8_NUM_TOKENS)
@pytest.mark.parametrize("N,K,E,topk", INT8_MOE_CONFIGS)
@pytest.mark.parametrize("seed", [0])
@pytest.mark.parametrize("is_vnni", [False, True])
@pytest.mark.parametrize("inplace", [False, True])
def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace):
"""Test fused_experts_cpu INT8 W8A8 against torch reference."""
set_random_seed(seed)
a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
w1_q, w2_q, w1_s, w2_s = _make_int8_moe_weights(E, N, K)
score = torch.randn(M, E, dtype=torch.bfloat16)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_ids = topk_ids.to(torch.int32)
ref_out = _ref_int8_moe(a, w1_q, w2_q, w1_s, w2_s, topk_weight, topk_ids)
w1 = _prepack_experts(w1_q) if is_vnni else w1_q
w2 = _prepack_experts(w2_q) if is_vnni else w2_q
out = ops.fused_experts_cpu(
a.clone(),
w1,
w2,
topk_weight,
topk_ids,
inplace,
ops.CPUQuantMethod.INT8_W8A8,
w1_s,
w2_s,
None, # w1_zero
None, # w2_zero
None, # block_size
None, # w1_bias
None, # w2_bias
None, # alpha
None, # limit
is_vnni,
)
torch.testing.assert_close(
ref_out.bfloat16(),
out,
atol=2e-1,
rtol=2e-1,
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))