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

[CPU] Add support for Gemma4 on Xeon (#22498)

Browse Source 
Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
Co-authored-by: Haotong Zou <haotong.zou@intel.com>
This commit is contained in:
blzheng
2026-08-17 10:52:26 +08:00
committed by GitHub
co-authored by Copilot jianan-gu Haotong Zou
parent 3adc70bb5e
commit b6d7602914
27 changed files with 514 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cpu/test_extend.py
+24 -2
View File
@@ -196,6 +196,7 @@ class TestExtendAttention(CustomTestCase):
*,
b_seq_len_prefix=None,
b_seq_len_extend=None,
kv_from_cache=False,
):
dtype = torch.bfloat16
@@ -322,8 +323,8 @@ class TestExtendAttention(CustomTestCase):
o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
torch.ops.sgl_kernel.extend_attention_cpu(
q_extend,
k_extend,
v_extend,
None if kv_from_cache else k_extend,
None if kv_from_cache else v_extend,
o_extend,
k_buffer,
v_buffer,
@@ -374,6 +375,27 @@ class TestExtendAttention(CustomTestCase):
1, 20, 1, 1, 64, 64, sliding_window, has_sink, False, False
)
def test_extend_attention_kv_from_cache(self):
# KV-shared layers pass no extend K/V, so the kernel masks the extend
# range causally itself; sizes straddle several BLOCK_N.
# Window only tested with a sink - _run_sdpa_forward_extend models none,
# the same restriction test_extend_attention applies.
for sliding_window, has_sink in [(None, False), (128, True)]:
for prefix, extend in [([0], [343]), ([100], [343]), ([0], [1500])]:
self._test_extend_attention_once(
B=1,
N_CTX=4096,
H_Q=16,
H_KV=4,
D=64,
DV=64,
sliding_window=sliding_window,
has_sink=has_sink,
b_seq_len_prefix=prefix,
b_seq_len_extend=extend,
kv_from_cache=True,
)
def test_extend_attention_large_seq_causal_mask(self):
self._test_extend_attention_once(
B=1,
test/registered/cpu/test_moe.py
+43 -6
View File
@@ -54,6 +54,7 @@ def run_fused_experts(
alpha=None,
limit=None,
is_vnni=True,
activation=None,
inplace=False,
):
return kernel.fused_experts_cpu(
@@ -74,6 +75,7 @@ def run_fused_experts(
alpha,
limit,
is_vnni,
activation,
)
@@ -138,13 +140,35 @@ def make_mxfp4_weights(e, out_dim, in_dim, dtype, with_bias=False):
class TestFusedExperts:
def test_unsupported_activation_is_rejected(self):
m, n, k, e, topk = 2, 32, 32, 4, 2
a = torch.randn((m, k), dtype=dtype) / 10
w1 = make_bf16_weights(e, 2 * n, k)
w2 = make_bf16_weights(e, k, n)
topk_weights, topk_ids = make_routing(m, e, topk, dtype=dtype)
packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
with pytest.raises(RuntimeError, match="Unsupported activation"):
run_fused_experts(
a,
packed_w1,
packed_w2,
topk_weights,
topk_ids,
quant=CPUQuantMethod.UNQUANT,
is_vnni=prepack,
activation="relu",
)
@pytest.mark.parametrize("m", [2, 114])
@pytest.mark.parametrize("n", [32])
@pytest.mark.parametrize("k", [32])
@pytest.mark.parametrize("e", [4])
@pytest.mark.parametrize("topk", [2])
@pytest.mark.parametrize("renormalize", [False, True])
def test_bf16_moe(self, m, n, k, e, topk, renormalize):
@pytest.mark.parametrize("activation", ["silu", "gelu"])
def test_bf16_moe(self, m, n, k, e, topk, renormalize, activation):
a = torch.randn((m, k), dtype=dtype) / 10
w1 = make_bf16_weights(e, 2 * n, k)
w2 = make_bf16_weights(e, k, n)
@@ -156,7 +180,9 @@ class TestFusedExperts:
renormalize=renormalize,
return_score=True,
)
torch_output = torch_naive_fused_moe(a, w1, w2, score, topk, renormalize)
torch_output = torch_naive_fused_moe(
a, w1, w2, score, topk, renormalize, activation=activation
)
packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
@@ -168,6 +194,7 @@ class TestFusedExperts:
topk_ids,
quant=CPUQuantMethod.UNQUANT,
is_vnni=prepack,
activation=activation,
inplace=True,
)
@@ -276,7 +303,8 @@ class TestFusedExperts:
@pytest.mark.parametrize("K", [256, 320])
@pytest.mark.parametrize("E", [8])
@pytest.mark.parametrize("topk", [4])
def test_fp8_moe(self, M, N, K, E, topk):
@pytest.mark.parametrize("activation", ["silu", "gelu"])
def test_fp8_moe(self, M, N, K, E, topk, activation):
a = torch.randn(M, K, dtype=dtype) / math.sqrt(K)
w1, w1s, w1_scaled = make_fp8_weights(E, 2 * N, K)
@@ -288,7 +316,7 @@ class TestFusedExperts:
w2 = kernel.convert_weight_packed(w2)
ref_out = native_fp8_fused_moe(
a, w1_scaled, w2_scaled, topk_weight, topk_ids, topk
a, w1_scaled, w2_scaled, topk_weight, topk_ids, topk, activation=activation
)
out = run_fused_experts(
a,
@@ -301,6 +329,7 @@ class TestFusedExperts:
w2_scale=w2s,
block_size=[BLOCK_N, BLOCK_K],
is_vnni=True,
activation=activation,
inplace=False,
)
@@ -372,7 +401,8 @@ class TestFusedExperts:
@pytest.mark.parametrize("K", [256, 320])
@pytest.mark.parametrize("E", [8])
@pytest.mark.parametrize("topk", [4])
def test_mxfp4_moe(self, M, N, K, E, topk):
@pytest.mark.parametrize("activation", ["silu", "gelu"])
def test_mxfp4_moe(self, M, N, K, E, topk, activation):
a = torch.randn(M, K, dtype=dtype) / 10
w1dq, w1_packed, w1s_packed = make_mxfp4_weights(E, 2 * N, K, dtype=dtype)
@@ -381,7 +411,13 @@ class TestFusedExperts:
topk_weight, topk_ids = make_routing(M, E, topk, dtype=dtype)
ref_out = native_fp8_fused_moe(
a, w1dq.float(), w2dq.float(), topk_weight, topk_ids, topk
a,
w1dq.float(),
w2dq.float(),
topk_weight,
topk_ids,
topk,
activation=activation,
)
out = run_fused_experts(
a,
@@ -393,6 +429,7 @@ class TestFusedExperts:
w1_scale=w1s_packed,
w2_scale=w2s_packed,
is_vnni=True,
activation=activation,
inplace=False,
)
test/registered/cpu/test_rope.py
+82 -5
View File
@@ -70,9 +70,9 @@ class TestROPE(CustomTestCase):
with torch.no_grad(), torch.amp.autocast("cpu", enabled=enable_autocast):
q = torch.randn(seq_len, num_heads * head_size, dtype=dtype)
q_clone = q.clone()
q_sgl = q.clone()
k = torch.randn(seq_len, num_kv_heads * head_size, dtype=dtype)
k_clone = k.clone()
k_sgl = k.clone()
# ref kernel
q_ref, k_ref = rope.forward_native(
@@ -81,10 +81,10 @@ class TestROPE(CustomTestCase):
positions=positions,
)
# fused rope kernel
q_sgl, k_sgl = torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu(
torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu(
positions,
q_clone,
k_clone,
q_sgl,
k_sgl,
rope.head_size,
rope.cos_sin_cache,
rope.mrope_section,
@@ -286,6 +286,83 @@ class TestROPE(CustomTestCase):
torch.testing.assert_close(q_out_ref, q_out_sgl, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(k_out_ref, k_out_sgl, atol=1e-2, rtol=1e-2)
def test_apply_multidimensional_rope(self):
"""Test apply_multidimensional_rope_cpu against the native Python reference."""
def _rotate_half(x):
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1)
def _apply_rotary(x, cos, sin):
return (x * cos) + (_rotate_half(x) * sin)
def _apply_multidimensional_rope_ref(x, cos, sin):
ndim = 2
chunk_size = x.shape[-1] // ndim
cos_3d = cos.unsqueeze(1)
sin_3d = sin.unsqueeze(1)
x_parts = x.split(chunk_size, dim=-1)
cos_parts = cos_3d.split(chunk_size, dim=-1)
sin_parts = sin_3d.split(chunk_size, dim=-1)
y_parts = [
_apply_rotary(x_parts[k], cos_parts[k], sin_parts[k])
for k in range(ndim)
]
return torch.cat(y_parts, dim=-1)
test_configs = [
# (num_tokens, num_heads, head_dim, dtype, sincos_dtype)
(4, 8, 64, torch.bfloat16, torch.bfloat16),
(32, 16, 128, torch.bfloat16, torch.bfloat16),
(128, 4, 256, torch.bfloat16, torch.bfloat16),
(1, 1, 32, torch.bfloat16, torch.float32),
(32, 16, 128, torch.bfloat16, torch.float32),
(2520, 12, 64, torch.bfloat16, torch.bfloat16),
(2520, 12, 64, torch.bfloat16, torch.float32),
# head_dim 160 -> 40 elements per rotary half, so the 32-wide
# vector loop runs once and leaves an 8-element scalar tail
(17, 3, 160, torch.bfloat16, torch.bfloat16),
(17, 3, 160, torch.float16, torch.float32),
(32, 16, 128, torch.float16, torch.float16),
]
for num_tokens, num_heads, head_dim, dtype, sincos_dtype in test_configs:
with self.subTest(
num_tokens=num_tokens,
num_heads=num_heads,
head_dim=head_dim,
dtype=dtype,
sincos_dtype=sincos_dtype,
):
torch.manual_seed(42)
query = torch.randn(
num_tokens, num_heads, head_dim, dtype=dtype, device="cpu"
)
key = torch.randn(
num_tokens, num_heads, head_dim, dtype=dtype, device="cpu"
)
cos = torch.randn(
num_tokens, head_dim, dtype=sincos_dtype, device="cpu"
)
sin = torch.randn(
num_tokens, head_dim, dtype=sincos_dtype, device="cpu"
)
q_expected = _apply_multidimensional_rope_ref(
query.float(), cos.float(), sin.float()
).to(dtype)
k_expected = _apply_multidimensional_rope_ref(
key.float(), cos.float(), sin.float()
).to(dtype)
torch.ops.sgl_kernel.apply_multidimensional_rope_cpu(
query, key, cos, sin
)
atol = rtol = precision[dtype]
torch.testing.assert_close(query, q_expected, atol=atol, rtol=rtol)
torch.testing.assert_close(key, k_expected, atol=atol, rtol=rtol)
if __name__ == "__main__":
unittest.main()