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

[CPU] Add Qwen3.5 model optimization for CPU (#19484)

Browse Source 
Co-authored-by: Zheng, Beilei <beilei.zheng@intel.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
jianan-gu
2026-04-26 10:12:36 -07:00
committed by GitHub
co-authored by Zheng, Beilei Ma Mingfei Xinyuan Tong
parent 7d49564431
commit 10fd0faccd
20 changed files with 768 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/cpu/test_mamba.py
+54 -47
View File
@@ -291,19 +291,20 @@ class TestMambaAttention(CustomTestCase):
def test_fused_gdn_gating(self):
dims = [6, 32]
for dim in dims:
A_log = torch.rand(dim)
a = torch.rand(1024, dim, dtype=torch.bfloat16)
b = torch.rand(1024, dim, dtype=torch.bfloat16)
dt_bias = torch.rand(dim, dtype=torch.bfloat16)
for A_log_dtype in [torch.float32, torch.bfloat16]:
A_log = torch.rand(dim, dtype=A_log_dtype)
a = torch.rand(1024, dim, dtype=torch.bfloat16)
b = torch.rand(1024, dim, dtype=torch.bfloat16)
dt_bias = torch.rand(dim, dtype=torch.bfloat16)
g, beta = torch_gdn_gating(A_log, a, b, dt_bias)
g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
A_log, a, b, dt_bias
)
atol = rtol = precision[g.dtype]
atol2 = rtol2 = precision[beta.dtype]
torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
g, beta = torch_gdn_gating(A_log, a, b, dt_bias)
g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
A_log, a, b, dt_bias
)
atol = rtol = precision[g.dtype]
atol2 = rtol2 = precision[beta.dtype]
torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
def test_fused_sigmoid_gating_delta_rule_update(self):
batch_size = 1
@@ -346,41 +347,47 @@ class TestMambaAttention(CustomTestCase):
if num_value_heads // num_heads > 1:
query_ref = query_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
key_ref = key_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
core_attn_out_ref, last_recurrent_state_ref = sigmoid_gating_delta_rule_update(
query_ref.transpose(0, 1),
key_ref.transpose(0, 1),
value.transpose(0, 1),
A_log,
a,
dt_bias,
b,
initial_state=ssm_states[cache_indices],
output_final_state=True,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
)
core_attn_out = torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
A_log=A_log,
dt_bias=dt_bias,
q=query,
k=key,
v=value,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
softplus_beta=1.0,
softplus_threshold=20.0,
)
last_recurrent_state = ssm_states[cache_indices]
atol = rtol = precision[core_attn_out.dtype]
torch.testing.assert_close(
core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
)
torch.testing.assert_close(
last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
)
for A_log_dtype in [torch.float32, torch.bfloat16]:
A_log = A_log.to(A_log_dtype)
core_attn_out_ref, last_recurrent_state_ref = (
sigmoid_gating_delta_rule_update(
query_ref.transpose(0, 1),
key_ref.transpose(0, 1),
value.transpose(0, 1),
A_log,
a,
dt_bias,
b,
initial_state=ssm_states[cache_indices],
output_final_state=True,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
)
)
core_attn_out = (
torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
A_log=A_log,
dt_bias=dt_bias,
q=query,
k=key,
v=value,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
softplus_beta=1.0,
softplus_threshold=20.0,
)
)
last_recurrent_state = ssm_states[cache_indices]
atol = rtol = precision[core_attn_out.dtype]
torch.testing.assert_close(
core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
)
torch.testing.assert_close(
last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
)
if __name__ == "__main__":
test/srt/cpu/test_qwen3.py
+62
View File
@@ -53,6 +53,34 @@ def fix_query_key_value_ordering_reshape_cat(
return mixed_qkv, z, b, a
def fix_query_key_value_ordering_reshape_cat_contiguous(
mixed_qkvz: torch.Tensor,
mixed_ba: torch.Tensor,
key_dim: int,
value_dim: int,
num_v_heads: int,
head_v_dim: int,
attn_tp_size: int,
):
"""
Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
"""
k_tp = key_dim // attn_tp_size
v_tp = value_dim // attn_tp_size
nv_tp = num_v_heads // attn_tp_size
# Directly split, no head group reshape
query, key, value, z = mixed_qkvz.split([k_tp, k_tp, v_tp, v_tp], dim=-1)
b, a = mixed_ba.split([nv_tp, nv_tp], dim=-1)
# value / z reshape to (seq, num_v_heads/tp, head_v_dim)
value = value.reshape(value.size(0), -1, head_v_dim)
z = z.reshape(z.size(0), -1, head_v_dim)
query, key, value = map(lambda x: x.reshape(x.shape[0], -1), (query, key, value))
mixed_qkv = torch.cat((query, key, value), dim=-1)
return mixed_qkv, z, b, a
class TestQwen3(CustomTestCase):
def test_fused_qkvzba_split_reshape_cat(self):
mixed_qkvz = torch.rand(1024, 12288, dtype=torch.bfloat16)
@@ -82,6 +110,40 @@ class TestQwen3(CustomTestCase):
torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
def test_fused_qkvzba_split_reshape_cat_contiguous(self):
mixed_qkvz = torch.rand(1, 12288, dtype=torch.bfloat16)
mixed_ba = torch.rand(1, 64, dtype=torch.bfloat16)
head_k_dim = 128
head_v_dim = 128
num_v_heads = 32
num_k_heads = 16
attn_tp_size = 1
key_dim = head_k_dim * num_k_heads
value_dim = head_v_dim * num_v_heads
mixed_qkv_ref, z_ref, b_ref, a_ref = (
fix_query_key_value_ordering_reshape_cat_contiguous(
mixed_qkvz,
mixed_ba,
key_dim,
value_dim,
num_v_heads,
head_v_dim,
attn_tp_size,
)
)
num_heads_qk = num_k_heads // attn_tp_size
num_heads_v = num_v_heads // attn_tp_size
mixed_qkv, z, b, a = (
torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu(
mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_k_dim, head_v_dim
)
)
atol = rtol = precision[mixed_qkv.dtype]
torch.testing.assert_close(mixed_qkv, mixed_qkv_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
if __name__ == "__main__":
unittest.main()