chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

[CPU] Optimize Qwen3-next model on CPU (#12525)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
Co-authored-by: Fan Yin <1106310035@qq.com>
This commit is contained in:
jianan-gu
2026-01-29 22:03:58 -08:00
committed by GitHub
co-authored by Ma Mingfei Fan Yin
parent 71e4d3b6bc
commit 336dc4579e
13 changed files with 366 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/qwen3_next.py
+7
View File
@@ -20,8 +20,11 @@ from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.configs.update_config import adjust_tp_num_heads_if_necessary
from sglang.srt.utils import is_cpu
logger = logging.get_logger(__name__)
_is_cpu = is_cpu()
class HybridLayerType(enum.Enum):
@@ -276,6 +279,10 @@ class Qwen3NextConfig(PretrainedConfig):
def mamba2_cache_params(self) -> Mamba2CacheParams:
from sglang.srt.layers.dp_attention import get_attention_tp_size
if _is_cpu:
world_size = get_attention_tp_size()
adjust_tp_num_heads_if_necessary(self, world_size, False)
shape = Mamba2StateShape.create(
tp_world_size=get_attention_tp_size(),
intermediate_size=self.linear_value_head_dim * self.linear_num_value_heads,
python/sglang/srt/configs/update_config.py
+43
View File
@@ -54,6 +54,44 @@ def get_num_heads_padding_size(tp_size, weight_block_size, head_dim):
return pad_size
def adjust_tp_num_heads_if_necessary(model_config, tp_size, is_post_update):
# is_post_update: whether to update an existing config
from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size
# Linear attn check logic
if hasattr(model_config, "linear_num_key_heads") and hasattr(
model_config, "linear_num_value_heads"
):
if (
model_config.linear_num_key_heads % tp_size != 0
or model_config.linear_num_value_heads % tp_size != 0
):
pad_size = tp_size
linear_num_key_heads_cpu = pad_vocab_size(
model_config.linear_num_key_heads, pad_size
)
linear_num_value_heads_cpu = (
linear_num_key_heads_cpu
* model_config.linear_num_value_heads
// model_config.linear_num_key_heads
)
if is_post_update:
model_config.linear_num_key_heads_cpu = linear_num_key_heads_cpu
model_config.linear_num_value_heads_cpu = linear_num_value_heads_cpu
else:
model_config.linear_num_key_heads = linear_num_key_heads_cpu
model_config.linear_num_value_heads = linear_num_value_heads_cpu
else:
if is_post_update:
model_config.linear_num_key_heads_cpu = (
model_config.linear_num_key_heads
)
model_config.linear_num_value_heads_cpu = (
model_config.linear_num_value_heads
)
def update_intermediate_size(model_config, attr_name, intermediate_padding_size):
attr_value = intermediate_padding_size
if hasattr(model_config, "hf_config") and hasattr(
@@ -137,6 +175,8 @@ def adjust_config_with_unaligned_cpu_tp(
model_config.hf_config.num_attention_heads = num_attention_heads
model_config.hf_text_config.num_attention_heads = num_attention_heads
adjust_tp_num_heads_if_necessary(model_config.hf_config, tp_size, True)
intermediate_padding_size = tp_size * get_moe_padding_size(weight_block_size)
model_config = update_intermediate_size(
model_config, "moe_intermediate_size", intermediate_padding_size
@@ -147,6 +187,9 @@ def adjust_config_with_unaligned_cpu_tp(
model_config = update_intermediate_size(
model_config, "intermediate_size_mlp", intermediate_padding_size
)
model_config = update_intermediate_size(
model_config, "shared_expert_intermediate_size", intermediate_padding_size
)
if (
hasattr(model_config.hf_config, "vision_config")
and model_config.hf_config.vision_config.model_type == "siglip_vision_model"