chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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-30 14:03:58 +08:00
committed by GitHub
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

26
python/sglang/srt/model_loader/weight_utils.py
View File

@@ -33,7 +33,10 @@ from tqdm.auto import tqdm
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.distributed import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
from sglang.srt.layers.dp_attention import get_attention_tp_rank
from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
from sglang.srt.layers.quantization.fp8 import Fp8Config
@@ -48,6 +51,7 @@ from sglang.srt.model_loader.ci_weight_validation import (
from sglang.srt.utils import (
BAR_FORMAT,
find_local_repo_dir,
is_cpu,
log_info_on_rank0,
print_warning_once,
)
@@ -1040,9 +1044,25 @@ def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
shard_size = param.data.shape[shard_axis]
start_idx = tp_rank * shard_size
loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
return default_weight_loader(param, loaded_weight)
if (
is_cpu()
and loaded_weight.size(0) % get_tensor_model_parallel_world_size() != 0
and loaded_weight.dim() == 1
):
param_data = param.data # view copy on param for uneven padding
param_data, loaded_weight = narrow_padded_param_and_loaded_weight(
param_data,
loaded_weight,
0, # param_data_start
start_idx,
shard_axis,
shard_size,
)
return default_weight_loader(param_data, loaded_weight)
else:
loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
return default_weight_loader(param, loaded_weight)
return loader