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

[CPU] Add GPT-OSS model optimization for CPU (#16775)

Browse Source 
Co-authored-by: mingfeima <mingfei.ma@intel.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
This commit is contained in:
blzheng
2026-05-29 16:05:26 +08:00
committed by GitHub
co-authored by mingfeima jianan-gu
parent 5601b7139d
commit 3ecf2c76ad
35 changed files with 2000 additions and 530 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cpu/test_decode.py
+139 -19
View File
@@ -12,6 +12,91 @@ torch.manual_seed(1234)
class TestDecodeAttention(CustomTestCase):
def _scaled_dot_product_attention(self, Q, K, V, S, scaling, sliding_window):
# sliding_window <= 0 means no sliding window
# Q: [n_tokens_q, n_heads, q_mult, d_head]
# K: [n_tokens_kv, n_heads, d_head]
# V: [n_tokens_kv, n_heads, d_head]
n_tokens_q, n_heads, q_mult, d_head = Q.shape
n_tokens_kv = K.shape[0]
assert K.shape == (n_tokens_kv, n_heads, d_head)
assert V.shape == (n_tokens_kv, n_heads, d_head)
K = K[:, :, None, :].expand(-1, -1, q_mult, -1)
V = V[:, :, None, :].expand(-1, -1, q_mult, -1)
S = S.reshape(n_heads, q_mult, 1, 1).expand(-1, -1, n_tokens_q, -1)
if n_tokens_q == n_tokens_kv: # Prefill
mask = torch.triu(
Q.new_full((n_tokens_q, n_tokens_kv), -float("inf")), diagonal=1
)
else: # Decode
mask = Q.new_zeros((n_tokens_q, n_tokens_kv))
if sliding_window is not None and sliding_window > 0:
mask += torch.tril(
mask.new_full((n_tokens_q, n_tokens_kv), -float("inf")),
diagonal=n_tokens_kv - n_tokens_q - sliding_window,
)
QK = torch.einsum("qhmd,khmd->hmqk", Q, K)
QK *= scaling
QK += mask[None, None, :, :]
QK = torch.cat([QK, S], dim=-1)
W = torch.softmax(QK, dim=-1)
W = W[..., :-1]
attn = torch.einsum("hmqk,khmd->qhmd", W, V)
return attn.reshape(n_tokens_q, -1)
def _run_sdpa_forward_decode_sink(
self,
query: torch.Tensor,
output: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
num_kv_heads: int,
q_mult: int,
scaling=None,
sliding_window=None,
attention_sinks=None,
enable_gqa=False,
causal=False,
):
# [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
query = query.movedim(0, query.dim() - 2)
start_q, start_kv = 0, 0
for seq_idx in range(seq_lens.shape[0]):
# TODO: this loop process a sequence per iter, this is inefficient.
# Need optimize the performance later.
seq_len_q = 1
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + seq_len_q
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_query = per_req_query.permute(1, 0, 2).reshape(
seq_len_q, num_kv_heads, q_mult, per_req_query.shape[-1]
)
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_key = per_req_key.permute(1, 0, 2)
per_req_value = per_req_value.permute(1, 0, 2)
per_req_out = self._scaled_dot_product_attention(
per_req_query,
per_req_key,
per_req_value,
attention_sinks,
scaling=scaling,
sliding_window=sliding_window,
).reshape(seq_len_q, -1, per_req_value.shape[-1])
output[start_q:end_q, :, :] = per_req_out
start_q, start_kv = end_q, end_kv
return output
def _run_sdpa_forward_decode(
self,
query: torch.Tensor,
@@ -71,11 +156,11 @@ class TestDecodeAttention(CustomTestCase):
return output
def _test_grouped_decode_attention_once(
self, B, H_Q, H_KV, D, D_V, is_cross_attn, dtype, device
self, B, H_Q, H_KV, D, D_V, sliding_window, sink, is_cross_attn, dtype, device
):
# This represents the number of tokens already in the sequence
seq_len = 1024
encoder_len = 10
encoder_len = 0 if sink else 10
total_tokens = B * (seq_len + encoder_len)
sm_scale = 1.0 / (D**0.5)
logit_cap = 0.0
@@ -84,6 +169,7 @@ class TestDecodeAttention(CustomTestCase):
# q represents the new token being generated, one per batch
q = torch.randn(B, H_Q, D, dtype=dtype, device=device)
sinks = torch.rand(H_Q, dtype=dtype, device=device) * 10
# k_buffer and v_buffer represent all previous tokens
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
@@ -137,22 +223,41 @@ class TestDecodeAttention(CustomTestCase):
sm_scale,
logit_cap,
is_cross_attn,
sliding_window if sliding_window is not None else 0,
encoder_lens,
sinks if sink else None,
)
self._run_sdpa_forward_decode(
q,
o_grouped,
k_buffer,
v_buffer,
req_to_token,
b_req_idx,
b_seq_len,
scaling=sm_scale,
enable_gqa=enable_gqa,
encoder_lens=encoder_lens,
is_cross_attn=is_cross_attn,
)
if sink:
self._run_sdpa_forward_decode_sink(
q,
o_grouped,
k_buffer,
v_buffer,
req_to_token,
b_req_idx,
b_seq_len,
num_kv_heads=H_KV,
q_mult=H_Q // H_KV if enable_gqa else 1,
scaling=sm_scale,
sliding_window=sliding_window if sliding_window is not None else None,
attention_sinks=sinks,
enable_gqa=enable_gqa,
)
else:
self._run_sdpa_forward_decode(
q,
o_grouped,
k_buffer,
v_buffer,
req_to_token,
b_req_idx,
b_seq_len,
scaling=sm_scale,
enable_gqa=enable_gqa,
encoder_lens=encoder_lens,
is_cross_attn=is_cross_attn,
)
cos_sim = torch.nn.functional.cosine_similarity(
o.flatten(), o_grouped.flatten(), dim=0
)
@@ -176,11 +281,26 @@ class TestDecodeAttention(CustomTestCase):
for B, H_Q, H_KV, D, D_V in configs:
for dtype in [torch.bfloat16, torch.float16]:
for sink in [True, False]:
if D != D_V and sink:
continue
for sliding_window in [None, 10]:
if sliding_window is not None and not sink:
continue
self._test_grouped_decode_attention_once(
B,
H_Q,
H_KV,
D,
D_V,
sliding_window,
sink,
False,
dtype=dtype,
device=device,
)
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, False, dtype=dtype, device=device
)
self._test_grouped_decode_attention_once(
B, H_Q, H_KV, D, D_V, True, dtype=dtype, device=device
B, H_Q, H_KV, D, D_V, None, False, True, dtype=dtype, device=device
)
def test_grouped_decode_attention(self):
test/registered/cpu/test_extend.py
+156 -21
View File
@@ -12,6 +12,37 @@ torch.manual_seed(1234)
class TestExtendAttention(CustomTestCase):
def _scaled_dot_product_attention(self, Q, K, V, S, scaling, sliding_window):
# sliding_window <= 0 means no sliding window
# Q: [n_tokens_q, n_heads, q_mult, d_head]
# K: [n_tokens_kv, n_heads, d_head]
# V: [n_tokens_kv, n_heads, d_head]
n_tokens_q, n_heads, q_mult, d_head = Q.shape
n_tokens_kv = K.shape[0]
assert K.shape == (n_tokens_kv, n_heads, d_head)
assert V.shape == (n_tokens_kv, n_heads, d_head)
K = K[:, :, None, :].expand(-1, -1, q_mult, -1)
V = V[:, :, None, :].expand(-1, -1, q_mult, -1)
S = S.reshape(n_heads, q_mult, 1, 1).expand(-1, -1, n_tokens_q, -1)
if n_tokens_q == n_tokens_kv: # Prefill
mask = torch.triu(
Q.new_full((n_tokens_q, n_tokens_kv), -float("inf")), diagonal=1
)
else: # Decode
mask = Q.new_zeros((n_tokens_q, n_tokens_kv))
if sliding_window is not None and sliding_window > 0:
mask += torch.tril(
mask.new_full((n_tokens_q, n_tokens_kv), -float("inf")),
diagonal=n_tokens_kv - n_tokens_q - sliding_window,
)
QK = torch.einsum("qhmd,khmd->hmqk", Q, K)
QK *= scaling
QK += mask[None, None, :, :]
QK = torch.cat([QK, S], dim=-1)
W = torch.softmax(QK, dim=-1)
W = W[..., :-1]
attn = torch.einsum("hmqk,khmd->qhmd", W, V)
return attn.reshape(n_tokens_q, -1)
def _run_sdpa_forward_extend(
self,
@@ -86,6 +117,70 @@ class TestExtendAttention(CustomTestCase):
start_q, start_kv = end_q, end_kv
return output
def _run_sdpa_forward_extend_sink(
self,
query: torch.Tensor,
output: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
extend_prefix_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
num_kv_heads: int,
q_mult: int,
scaling=None,
sliding_window=None,
attention_sinks=None,
enable_gqa=False,
causal=False,
):
assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
assert seq_lens.shape[0] == extend_seq_lens.shape[0]
# [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
query = query.movedim(0, query.dim() - 2)
start_q, start_kv = 0, 0
for seq_idx in range(seq_lens.shape[0]):
# TODO: this loop process a sequence per iter, this is inefficient.
# Need optimize the performance later.
extend_seq_len_q = extend_seq_lens[seq_idx]
prefill_seq_len_q = extend_prefix_lens[seq_idx]
seq_len_kv = seq_lens[seq_idx]
end_q = start_q + extend_seq_len_q
end_kv = start_kv + seq_len_kv
per_req_query = query[:, start_q:end_q, :]
per_req_query_redudant = torch.empty(
(per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]),
dtype=per_req_query.dtype,
device=per_req_query.device,
)
per_req_query_redudant[:, prefill_seq_len_q:, :] = per_req_query
# get key and value from cache. per_req_tokens contains the kv cache
# index for each token in the sequence.
req_pool_idx = req_pool_indices[seq_idx]
per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_query_redudant = per_req_query_redudant.permute(1, 0, 2).reshape(
seq_len_kv, num_kv_heads, q_mult, per_req_query_redudant.shape[-1]
)
per_req_key = per_req_key.permute(1, 0, 2)
per_req_value = per_req_value.permute(1, 0, 2)
per_req_out_redudant = self._scaled_dot_product_attention(
per_req_query_redudant,
per_req_key,
per_req_value,
attention_sinks,
scaling=scaling,
sliding_window=sliding_window,
).reshape(seq_len_kv, -1, per_req_value.shape[-1])
output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :]
start_q, start_kv = end_q, end_kv
return output
def _test_extend_attention_once(
self,
B,
@@ -94,6 +189,8 @@ class TestExtendAttention(CustomTestCase):
H_KV,
D,
DV,
sliding_window=None,
has_sink=False,
mla=False,
is_cross_attn=False,
*,
@@ -108,9 +205,13 @@ class TestExtendAttention(CustomTestCase):
b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)
encoder_lens = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int64)
scale = 20
if mla:
b_seq_len_prefix.zero_()
encoder_lens.zero_()
if has_sink:
encoder_lens.zero_()
scale = 1
if b_seq_len_extend is None:
b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
@@ -143,6 +244,7 @@ class TestExtendAttention(CustomTestCase):
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype)
v_extend = torch.empty((extend_token_num, H_KV, DV), dtype=dtype)
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
sinks = torch.rand(H_Q, dtype=dtype)
for i in range(B):
extend_start_in_buffer = (
@@ -158,7 +260,7 @@ class TestExtendAttention(CustomTestCase):
extend_start_in_buffer:extend_end_in_buffer
]
q_extend[extend_start:extend_end] = (
torch.randn((b_seq_len_extend[i], H_Q, D), dtype=dtype) * 20
torch.randn((b_seq_len_extend[i], H_Q, D), dtype=dtype) * scale
)
# q_extend, k_extend, v_extend, k_buffer and v_buffer supports non-contiguous tensors
@@ -182,22 +284,40 @@ class TestExtendAttention(CustomTestCase):
enable_gqa = H_Q != H_KV
o_ref = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
self._run_sdpa_forward_extend(
q_extend,
o_ref,
k_buffer,
v_buffer,
req_to_tokens,
b_req_idx,
b_seq_len,
b_seq_len_prefix,
b_seq_len_extend,
scaling=sm_scale,
enable_gqa=enable_gqa,
causal=not is_cross_attn,
is_cross_attn=is_cross_attn,
encoder_lens=encoder_lens,
)
if has_sink:
self._run_sdpa_forward_extend_sink(
q_extend,
o_ref,
k_buffer,
v_buffer,
req_to_tokens,
b_req_idx,
b_seq_len,
b_seq_len_prefix,
b_seq_len_extend,
H_KV,
H_Q // H_KV if enable_gqa else 1,
scaling=sm_scale,
sliding_window=sliding_window,
attention_sinks=sinks,
)
else:
self._run_sdpa_forward_extend(
q_extend,
o_ref,
k_buffer,
v_buffer,
req_to_tokens,
b_req_idx,
b_seq_len,
b_seq_len_prefix,
b_seq_len_extend,
scaling=sm_scale,
enable_gqa=enable_gqa,
causal=not is_cross_attn,
is_cross_attn=is_cross_attn,
encoder_lens=encoder_lens,
)
o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
torch.ops.sgl_kernel.extend_attention_cpu(
@@ -216,7 +336,9 @@ class TestExtendAttention(CustomTestCase):
sm_scale,
logit_cap,
is_cross_attn,
sliding_window if sliding_window is not None else 0,
encoder_lens,
sinks if has_sink else None,
)
torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
@@ -227,16 +349,29 @@ class TestExtendAttention(CustomTestCase):
if is_mla and is_cross_attn:
continue
self._test_extend_attention_once(
1, 123, 1, 1, 128, 96, is_mla, is_cross_attn
1, 123, 1, 1, 128, 96, None, False, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 123, 16, 1, 128, 96, is_mla, is_cross_attn
1, 123, 16, 1, 128, 96, None, False, is_mla, is_cross_attn
)
self._test_extend_attention_once(
4, 1230, 16, 4, 128, 96, is_mla, is_cross_attn
4, 1230, 16, 4, 128, 96, None, False, is_mla, is_cross_attn
)
self._test_extend_attention_once(
1, 9000, 16, 1, 32, 32, is_mla, is_cross_attn
1, 9000, 16, 1, 32, 32, None, False, is_mla, is_cross_attn
)
for has_sink in [True, False]:
for sliding_window in [None, 10, 128]:
if not has_sink and sliding_window is not None:
continue
self._test_extend_attention_once(
1, 123, 16, 4, 64, 64, sliding_window, has_sink, False, False
)
self._test_extend_attention_once(
1, 20, 16, 1, 64, 64, sliding_window, has_sink, False, False
)
self._test_extend_attention_once(
1, 20, 1, 1, 64, 64, sliding_window, has_sink, False, False
)
def test_extend_attention_large_seq_causal_mask(self):
test/registered/cpu/test_gemm.py
+66 -122
View File
@@ -1,12 +1,13 @@
import itertools
import unittest
# TODO: use interface in cpu.py
import torch
import torch.nn as nn
from utils import (
MXFP4QuantizeUtil,
convert_weight,
native_w8a8_per_token_matmul,
parametrize,
per_token_quant_int8,
precision,
unpack_and_dequant_awq,
@@ -31,30 +32,15 @@ class Mod(nn.Module):
class TestGemm(CustomTestCase):
M = [1, 101]
N = [16, 32 * 13]
K = [32 * 16]
has_bias = [False, True]
dim = [2, 3, 4, 5]
M_int8 = [2, 128]
N_int8 = [32 * 12]
K_int8 = [32 * 17]
M_fp8 = [1, 11]
N_fp8 = [128, 224]
K_fp8 = [512, 576]
M_awq = [1, 32]
N_awq = [4096]
K_awq = [4096]
M_gptq = [1, 32]
N_gptq = [4096]
K_gptq = [4096]
def _bf16_gemm(self, M, N, K, has_bias, dim):
@parametrize(
M=[1, 101],
N=[16, 32 * 13],
K=[32 * 16],
has_bias=[False, True],
dim=[2, 3, 4, 5],
)
def test_bf16_gemm(self, M, N, K, has_bias, dim):
mat1 = torch.randn(M, K, dtype=torch.bfloat16)
mat2 = torch.randn(N, K, dtype=torch.bfloat16)
if dim == 3:
@@ -84,24 +70,14 @@ class TestGemm(CustomTestCase):
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
torch.testing.assert_close(ref, out2, atol=atol, rtol=rtol)
def test_bf16_gemm(self):
for params in itertools.product(
self.M,
self.N,
self.K,
self.has_bias,
self.dim,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
has_bias=params[3],
dim=params[4],
):
self._bf16_gemm(*params)
def _bf16_gemm_with_small_oc(self, M, N, K, has_bias, use_post_sigmul):
@parametrize(
M=[1, 8, 32, 1024],
N=[12, 1],
K=[32 * 16],
has_bias=[False, True],
use_post_sigmul=[False, True],
)
def bf16_gemm_with_small_oc(self, M, N, K, has_bias, use_post_sigmul):
use_post_sigmul = use_post_sigmul and N == 1
mat_mul = (
None if not use_post_sigmul else torch.randn(M, 2 * K, dtype=torch.bfloat16)
@@ -132,20 +108,8 @@ class TestGemm(CustomTestCase):
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
def test_bf16_gemm_with_small_oc(self):
for params in itertools.product(
[1, 8, 32, 1024], [12, 1], self.K, self.has_bias, [False, True]
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
has_bias=params[3],
use_post_sigmul=params[4],
):
self._bf16_gemm_with_small_oc(*params)
def _int8_gemm(self, M, N, K, has_bias):
@parametrize(M=[2, 128], N=[32 * 12], K=[32 * 17], has_bias=[False, True])
def test_int8_gemm(self, M, N, K, has_bias):
dtype = torch.bfloat16
A = torch.randn((M, K), dtype=dtype) / 10
Aq, As = per_token_quant_int8(A)
@@ -175,35 +139,21 @@ class TestGemm(CustomTestCase):
)
torch.testing.assert_close(ref_out, fused_out, atol=atol, rtol=rtol)
def test_int8_gemm(self):
for params in itertools.product(
self.M_int8,
self.N_int8,
self.K_int8,
self.has_bias,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
has_bias=params[3],
):
self._int8_gemm(*params)
def _fp8_gemm(self, M, N, K, has_bias):
@parametrize(M=[1, 11], N=[128, 224], K=[512, 576], has_bias=[False, True])
def test_fp8_gemm(self, M, N, K, has_bias):
prepack = True
chunk = False
scale_block_size_N = 64
scale_block_size_K = 128
assert scale_block_size_N <= N
assert scale_block_size_K <= K
A_dtype = torch.bfloat16
dtype = torch.bfloat16
model = Mod(K, N, has_bias).eval()
if chunk:
data = torch.randn(M, K + 6, dtype=A_dtype).narrow(1, 0, K)
data = torch.randn(M, K + 6, dtype=dtype).narrow(1, 0, K)
else:
data = torch.randn(M, K, dtype=A_dtype)
data = torch.randn(M, K, dtype=dtype)
weight = model.linear.weight # (N, K)
@@ -211,18 +161,18 @@ class TestGemm(CustomTestCase):
bias = model.linear.bias
fp8_weight, scales, dq_weight = convert_weight(
weight, [scale_block_size_N, scale_block_size_K], A_dtype
weight, [scale_block_size_N, scale_block_size_K], dtype
)
if has_bias:
ref = torch.matmul(data.to(A_dtype), dq_weight.T) + bias.to(A_dtype)
ref = torch.matmul(data.to(dtype), dq_weight.T) + bias.to(dtype)
else:
ref = torch.matmul(data.to(A_dtype), dq_weight.T)
ref = torch.matmul(data.to(dtype), dq_weight.T)
if prepack:
fp8_weight = torch.ops.sgl_kernel.convert_weight_packed(fp8_weight)
opt = torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
out = torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
data,
fp8_weight,
scales,
@@ -232,24 +182,41 @@ class TestGemm(CustomTestCase):
prepack,
)
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, opt, atol=atol, rtol=rtol)
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
def test_fp8_gemm(self):
for params in itertools.product(
self.M_fp8,
self.N_fp8,
self.K_fp8,
self.has_bias,
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
has_bias=params[3],
):
self._fp8_gemm(*params)
@parametrize(M=[1, 11], N=[128, 224], K=[512, 576], has_bias=[False, True])
def test_mxfp4_gemm(self, M, N, K, has_bias):
prepack = True
dtype = torch.bfloat16
def _int4_awq_gemm(self, M, N, K, group_size, has_bias):
A = torch.randn((M, K), dtype=dtype) / 10
# we randomly generate Bq and Bs, then dequantize it to BFloat16 as reference
Bq = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8)
Bs = torch.randint(126, 127, (N, K // 32), dtype=torch.uint8)
Bdq = MXFP4QuantizeUtil.dequantize(Bq, dtype, Bs)
B_packed = torch.ops.sgl_kernel.convert_weight_packed(Bq)
Bs_packed = torch.ops.sgl_kernel.convert_scale_packed(Bs)
bias = torch.randn(N) if has_bias else None
ref = torch.matmul(A.float(), Bdq.float().t()).bfloat16()
if bias is not None:
ref.add_(bias.view(1, -1))
out = torch.ops.sgl_kernel.mxfp4_scaled_mm_cpu(
A, B_packed, Bs_packed, bias, prepack
)
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
@parametrize(
M=[1, 32], N=[4096], K=[4096], group_size=[128], has_bias=[False, True]
)
def test_int4_awq_gemm(self, M, N, K, group_size, has_bias):
awq_weight = torch.randint(-128, 128, (K, N // 8)).to(torch.int)
awq_zero = torch.randint(0, 10, (K // group_size, N // 8)).to(torch.int)
awq_scales = torch.rand(int(K // group_size), N).to(torch.bfloat16)
@@ -281,20 +248,10 @@ class TestGemm(CustomTestCase):
atol = rtol = precision[ref_res.dtype]
torch.testing.assert_close(ref_res, target_res, atol=atol, rtol=rtol)
def test_int4_awq_gemm(self):
for params in itertools.product(
self.M_awq, self.N_awq, self.K_awq, [128], self.has_bias
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
group_size=params[3],
has_bias=params[4],
):
self._int4_awq_gemm(*params)
def _int4_gptq_gemm(self, M, N, K, group_size, has_bias):
@parametrize(
M=[1, 32], N=[4096], K=[4096], group_size=[128], has_bias=[False, True]
)
def test_int4_gptq_gemm(self, M, N, K, group_size, has_bias):
torch.manual_seed(127)
gptq_weight = torch.randint(-128, 128, (K // 8, N)).to(torch.int)
gptq_zero = torch.randint(0, 10, (K // group_size, N // 8)).to(torch.int)
@@ -326,19 +283,6 @@ class TestGemm(CustomTestCase):
atol = rtol = precision[ref_res.dtype]
torch.testing.assert_close(ref_res, target_res, atol=atol, rtol=rtol)
def test_int4_gptq_gemm(self):
for params in itertools.product(
self.M_gptq, self.N_gptq, self.K_gptq, [128], self.has_bias
):
with self.subTest(
M=params[0],
N=params[1],
K=params[2],
group_size=params[3],
has_bias=params[4],
):
self._int4_gptq_gemm(*params)
if __name__ == "__main__":
unittest.main()
test/registered/cpu/test_mla.py
+2
View File
@@ -119,6 +119,8 @@ class TestMLA(CustomTestCase):
sm_scale,
logit_cap,
False,
0,
None,
None,
)
test/registered/cpu/test_moe.py
+183 -98
View File
@@ -1,4 +1,3 @@
import itertools
import math
import unittest
@@ -9,18 +8,21 @@ from sglang.srt.layers.amx_utils import CPUQuantMethod
kernel = torch.ops.sgl_kernel
torch.manual_seed(128)
torch.manual_seed(1234)
from utils import (
BLOCK_K,
BLOCK_N,
MXFP4QuantizeUtil,
factor_for_scale,
fp8_max,
fp8_min,
native_fp8_fused_moe,
parametrize,
precision,
scaled_weight,
torch_naive_fused_moe,
torch_naive_fused_moe_gptoss,
torch_w8a8_per_column_fused_moe,
unpack_and_dequant_awq,
)
@@ -60,37 +62,18 @@ def fused_moe(a, w1, w2, score, topk, renormalize, prepack):
None,
None,
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):
@parametrize(m=[2, 114], n=[32], k=[32], e=[4], topk=[2], renormalize=[False, True])
def test_bf16_moe(self, m, n, k, e, topk, renormalize):
dtype = torch.bfloat16
prepack = True
@@ -105,26 +88,51 @@ class TestFusedExperts(CustomTestCase):
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)
@parametrize(
m=[1, 32], n=[128, 64], k=[128, 64], e=[4], topk=[2], renormalize=[False]
)
def test_bf16_moe_bias(self, m, n, k, e, topk, renormalize):
dtype = torch.bfloat16
def _int8_moe(self, M, N, K, E, topk):
a = torch.randn((m, k), device="cpu", dtype=dtype) / 10
w1 = torch.randn((e, 2 * n, k), device="cpu", dtype=dtype) / 10
w1_b = torch.randn((e, 2 * n), device="cpu", dtype=torch.float) / 10
w2 = torch.randn((e, k, n), device="cpu", dtype=dtype) / 10
w2_b = torch.randn((e, k), device="cpu", dtype=torch.float) / 10
score = torch.randn((m, e), device="cpu", dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
alpha = 1.702
limit = 7.0
torch_output = torch_naive_fused_moe_gptoss(
a, w1, w2, w1_b, w2_b, topk_weight, topk_ids, renormalize, alpha, limit, e
)
packed_w1 = kernel.convert_weight_packed(w1)
packed_w2 = kernel.convert_weight_packed(w2)
fused_output = torch.ops.sgl_kernel.fused_experts_cpu(
a,
packed_w1,
packed_w2,
topk_weight,
topk_ids.to(torch.int),
False, # inplace # See [Note] inplace should be False in fused_experts.
CPUQuantMethod.UNQUANT,
None, # w1_scale
None, # w2_scale
None, # w1_zp
None, # w2_zp
None, # block_size
w1_b,
w2_b,
alpha,
limit,
True, # is_vnni
)
atol = rtol = precision[torch_output.dtype]
torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
@parametrize(M=[1, 39], N=[128], K=[256], E=[8], topk=[3])
def test_int8_moe(self, M, N, K, E, topk):
dtype = torch.bfloat16
prepack = True
@@ -173,6 +181,10 @@ class TestFusedExperts(CustomTestCase):
None,
None,
None,
None,
None,
None,
None,
prepack,
)
@@ -182,24 +194,8 @@ class TestFusedExperts(CustomTestCase):
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):
@parametrize(M=[2, 121], N=[352, 512], K=[256, 320], E=[8], topk=[4])
def test_fp8_moe(self, M, N, K, E, topk):
dtype = torch.bfloat16
a = torch.randn(M, K, dtype=dtype) / math.sqrt(K)
@@ -245,30 +241,132 @@ class TestFusedExperts(CustomTestCase):
None,
None,
[BLOCK_N, BLOCK_K],
None,
None,
None,
None,
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)
@parametrize(M=[2, 121], N=[352, 512], K=[256, 320], E=[8], topk=[4])
def test_mxfp4_moe(self, M, N, K, E, topk):
dtype = torch.bfloat16
def _int4_moe(self, M, N, K, E, topk, group_size=128):
a = torch.randn(M, K, dtype=dtype) / 10
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)
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(w1q)
w2 = kernel.convert_weight_packed(w2q)
w1s = kernel.convert_scale_packed(w1s)
w2s = kernel.convert_scale_packed(w2s)
ref_out = native_fp8_fused_moe(
a, w1dq.float(), w2dq.float(), topk_weight, topk_ids, topk
)
out = kernel.fused_experts_cpu(
a,
w1,
w2,
topk_weight,
topk_ids.to(torch.int32),
False,
CPUQuantMethod.MXFP4,
w1s,
w2s,
None,
None,
None,
None,
None,
None,
None,
True,
)
atol = rtol = precision[dtype]
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
@parametrize(
m=[1, 32], n=[128, 64], k=[128, 64], e=[4], topk=[2], renormalize=[False]
)
def test_mxfp4_moe_bias(self, m, n, k, e, topk, renormalize):
dtype = torch.bfloat16
a = torch.randn((m, k), device="cpu", dtype=dtype) / 10
w1_bf16 = torch.randn((e, 2 * n, k), device="cpu", 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), device="cpu", dtype=torch.float32) / 10
w2_bf16 = torch.randn((e, k, n), device="cpu", 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), device="cpu", dtype=torch.float32) / 10
score = torch.randn((m, e), device="cpu", dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
alpha = 1.702
limit = 7.0
torch_output = torch_naive_fused_moe_gptoss(
a,
w1dq,
w2dq,
w1_b,
w2_b,
topk_weight,
topk_ids,
renormalize,
alpha,
limit,
e,
)
w1 = kernel.convert_weight_packed(w1q)
w2 = kernel.convert_weight_packed(w2q)
w1s = kernel.convert_scale_packed(w1s)
w2s = kernel.convert_scale_packed(w2s)
fused_output = torch.ops.sgl_kernel.fused_experts_cpu(
a,
w1,
w2,
topk_weight,
topk_ids.to(torch.int32),
False, # inplace # See [Note] inplace should be False in fused_experts.
CPUQuantMethod.MXFP4, # use_mxfp4
w1s, # w1_scale
w2s, # w2_scale
None, # w1_zp
None, # w2_zp
None, # block_size
w1_b,
w2_b,
alpha,
limit,
True, # is_vnni
)
atol = rtol = precision[torch_output.dtype]
torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
@parametrize(M=[1, 6], N=[512], K=[256], E=[8], topk=[4])
def test_int4_moe(self, M, N, K, E, topk, group_size=128):
dtype = torch.bfloat16
a = torch.rand(M, K, dtype=dtype) / math.sqrt(K)
@@ -327,29 +425,16 @@ class TestFusedExperts(CustomTestCase):
awq_w13_zero_pack,
awq_w2_zero_pack,
None,
None,
None,
None,
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()
test/registered/cpu/utils.py
+216
View File
@@ -221,6 +221,107 @@ def torch_naive_fused_moe(a, w1, w2, score, topk, renormalize):
).sum(dim=1)
def moe_gptoss_act(x, alpha: float = 1.702, limit: float = 7.0):
x_glu, x_linear = x[..., ::2], x[..., 1::2]
# Clamp the input values
x_glu = x_glu.clamp(min=None, max=limit)
x_linear = x_linear.clamp(min=-limit, max=limit)
out_glu = x_glu * torch.sigmoid(alpha * x_glu)
# Note we add an extra bias of 1 to the linear layer
return out_glu * (x_linear + 1.0)
def torch_naive_gptoss_fused_moe(
x,
w1,
w2,
w1_bias,
w2_bias,
topk_weights,
topk_ids,
activation_alpha,
swiglu_limit,
len_experts,
) -> torch.Tensor:
# Ref code from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/e0828e3cc0a03408724b80c3cc92c8e072db8d01/modeling_deepseek.py#L589
cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
cnts.scatter_(1, topk_ids.to(torch.int64), 1)
tokens_per_expert = cnts.sum(dim=0)
idxs = topk_ids.view(-1).argsort()
sorted_tokens = x[idxs // topk_ids.shape[1]]
tokens_per_expert = tokens_per_expert.cpu().numpy()
outputs = []
start_idx = 0
for i, num_tokens in enumerate(tokens_per_expert):
end_idx = start_idx + num_tokens
if num_tokens == 0:
continue
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
layer_w13_weight = w1[i]
layer_w13_weight_bias = w1_bias[i]
layer_w2_weight_bias = w2_bias[i]
layer_w2_weight = w2[i]
gate_up = F.linear(
tokens_for_this_expert,
layer_w13_weight,
bias=layer_w13_weight_bias.to(torch.bfloat16),
)
gate_up = moe_gptoss_act(gate_up, activation_alpha, swiglu_limit)
expert_out = F.linear(
gate_up, layer_w2_weight, bias=layer_w2_weight_bias.to(torch.bfloat16)
)
outputs.append(expert_out)
start_idx = end_idx
outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
new_x = torch.empty_like(outs)
new_x[idxs] = outs
final_out = (
new_x.view(*topk_ids.shape, -1)
.type(topk_weights.dtype)
.mul_(topk_weights.unsqueeze(dim=-1))
.sum(dim=1)
.type(new_x.dtype)
)
return final_out
def torch_naive_fused_moe_gptoss(
a,
w1,
w2,
w1_bias,
w2_bias,
topk_weight,
topk_ids,
renormalize,
activation_alpha,
swiglu_limit,
len_experts,
):
if renormalize:
topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True)
return torch_naive_gptoss_fused_moe(
a,
w1,
w2,
w1_bias,
w2_bias,
topk_weight,
topk_ids,
activation_alpha,
swiglu_limit,
len_experts,
)
def torch_w8a8_per_column_fused_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, topk):
"""This function performs fused moe with per-column int8 quantization using native torch."""
@@ -294,6 +395,121 @@ def native_fp8_fused_moe(a, w1, w2, topk_weight, topk_ids, topk):
)
# https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/quantization/qtensor/mxfp4_tensor.py
class MXFP4QuantizeUtil:
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:
"""Converting a tensor to a quantized format based on MXFP4 quantization. Only E4M3 is supported.
Args:
input (torch.Tensor): The input tensor to be quantized.
"""
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):
# If the last dimension is odd, pad with zeros
# If this behavior is not desired, please modify the code accordingly
left_side = x[..., 0::2] # Even indices (0, 2, 4...)
right_side = x[..., 1::2] # Odd indices (1, 3, 5...)
new_data = (
right_side.clone() << 4
) # Put odd indices (higher addresses) in high bits
new_data[
..., : left_side.shape[-1]
] += left_side # Put even indices in low bits
return new_data
original_shape = input.shape
original_dtype = input.dtype
input = input.view(-1, cls.block_size)
# get scales
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):
"""Dequantze MXFP4 packed tensor to a target dtype."""
def unfuse_uint8_to_uint4(x):
"""Unfuse uint8 values back to uint4 values.
This is the inverse operation of fuse_uint4_to_uint8.
"""
# Extract the lower 4 bits (even indices)
left_side = x & 0x0F
# Extract the upper 4 bits (odd indices)
right_side = (x >> 4) & 0x0F
# Create a new tensor with alternating values
shape = list(x.shape)
shape[-1] = shape[-1] * 2
result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
# Fill in the values - even indices get low bits, odd indices get high bits
result[..., 0::2] = left_side # Even indices from low bits
result[..., 1::2] = right_side # Odd indices from high bits
return result
e8m0_scale = scale
# Unfuse the uint8 values back to uint4
x_unfused = unfuse_uint8_to_uint4(quantized_data)
# print("@@@ x_unfused: ", x_unfused)
# Extract sign and magnitude
sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(
torch.float32
) # Extract sign bit and convert to +1/-1
magnitude = x_unfused & 0b0111 # Extract magnitude bits
magnitude = magnitude.to(torch.long)
# Create a tensor with the E2M1 values
values = torch.tensor(cls.E2M1_values, device=quantized_data.device)
# Use gather to index the values tensor properly
# We need to reshape magnitude to match the dimensions we want to gather along
original_shape = magnitude.shape
x_float = values[magnitude.reshape(-1)].reshape(original_shape)
# Apply sign and scale
x_float = sign.float() * x_float
# Reshape to apply block-wise scaling
x_float = x_float.reshape(-1, cls.block_size)
# Apply the E8M0 scale
scale_factor = torch.exp2(e8m0_scale.float() - 127)
scale_factor = scale_factor.reshape(-1, 1) # Reshape for proper broadcasting
# Apply scaling and reshape back to original shape
x_float = x_float * scale_factor
# Reshape back to the original shape
return x_float.reshape(original_shape).to(dtype)
def make_non_contiguous(x: torch.Tensor) -> torch.Tensor:
"""
Make a tensor non-contiguous by slicing it via last dimension.