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
python/sglang/srt/layers/amx_utils.py
+2
View File
@@ -15,6 +15,7 @@ class CPUQuantMethod(IntEnum):
INT8_W8A8 = 1
FP8_W8A16 = 2
INT4_W4A8 = 3
MXFP4 = 4
class CPUQuantAlgo(IntEnum):
@@ -96,6 +97,7 @@ def dtype_is_supported(weight):
return weight.dtype in [
torch.float16,
torch.bfloat16,
torch.uint8,
torch.int8,
torch.float8_e4m3fn,
]
python/sglang/srt/layers/attention/intel_amx_backend.py
+6
View File
@@ -98,6 +98,7 @@ class IntelAMXAttnBackend(AttentionBackend):
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache=True,
sinks=None,
):
if layer.qk_head_dim != layer.v_head_dim:
o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
@@ -128,7 +129,9 @@ class IntelAMXAttnBackend(AttentionBackend):
layer.scaling,
layer.logit_cap,
layer.is_cross_attention,
layer.sliding_window_size + 1,
forward_batch.encoder_lens,
sinks,
)
return o
@@ -140,6 +143,7 @@ class IntelAMXAttnBackend(AttentionBackend):
layer: RadixAttention,
forward_batch: ForwardBatch,
save_kv_cache=True,
sinks=None,
):
attn_logits, _ = self.forward_metadata
@@ -169,7 +173,9 @@ class IntelAMXAttnBackend(AttentionBackend):
layer.scaling,
layer.logit_cap,
layer.is_cross_attention,
layer.sliding_window_size + 1,
forward_batch.encoder_lens,
sinks,
)
return o
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+5
View File
@@ -218,6 +218,7 @@ class FusedMoE(torch.nn.Module):
self.use_presharded_weights = use_presharded_weights
self.use_triton_kernels = get_moe_runner_backend().is_triton_kernels()
self.use_flashinfer_trtllm_moe = (
get_moe_runner_backend().is_flashinfer_trtllm()
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
@@ -465,6 +466,8 @@ class FusedMoE(torch.nn.Module):
start = 0
if self.use_padded_loading:
if _is_cpu and is_bias:
shard_dim = 1
expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
expert_data,
loaded_weight,
@@ -534,6 +537,8 @@ class FusedMoE(torch.nn.Module):
shard_size = expert_data.shape[shard_dim]
if self.use_padded_loading:
if _is_cpu and is_bias:
shard_dim = 1
expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
expert_data,
loaded_weight,
python/sglang/srt/layers/quantization/__init__.py
+3 -1
View File
@@ -53,6 +53,7 @@ from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config
from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config
from sglang.srt.utils import (
cpu_has_amx_support,
is_cpu,
is_cuda,
is_hip,
is_mps,
@@ -95,7 +96,7 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
}
if is_cuda() or (_is_mxfp_supported and is_hip()):
if is_cpu() or is_cuda() or (_is_mxfp_supported and is_hip()):
BASE_QUANTIZATION_METHODS.update(
{
"mxfp4": Mxfp4Config,
@@ -126,6 +127,7 @@ CPU_QUANTIZATION_METHODS = {
"compressed-tensors": CompressedTensorsConfig,
"awq": AWQCPUConfig,
"gptq": CPUGPTQConfig,
"mxfp4": Mxfp4Config,
}
QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS}
python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py
+4
View File
@@ -99,6 +99,10 @@ class AWQIntelAMXMoEKernel:
layer.w13_qzeros,
layer.w2_qzeros,
None, # block_size
None, # w1 bias
None, # w3 bias
None, # alpha
None, # limit
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/layers/quantization/fp8.py
+4
View File
@@ -1828,6 +1828,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
None, # w1_zp
None, # w2_zp
self.quant_config.weight_block_size, # block_size
None, # w1 bias
None, # w3 bias
None, # alpha
None, # limit
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/layers/quantization/gptq_cpu.py
+4
View File
@@ -370,6 +370,10 @@ class GPTQMoEIntelAMXMethod(FusedMoEMethodBase):
layer.w13_qzeros,
layer.w2_qzeros,
None, # block_size
None, # w1 bias
None, # w3 bias
None, # alpha
None, # limit
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/layers/quantization/mxfp4.py
+60
View File
@@ -33,6 +33,10 @@ from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.amx_utils import (
CPUQuantMethod,
_amx_process_weight_after_loading,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
@@ -46,6 +50,8 @@ from sglang.srt.layers.quantization.base_config import (
from sglang.srt.layers.quantization.utils import is_layer_skipped
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
cpu_has_amx_support,
is_cpu,
is_flashinfer_available,
is_gfx95_supported,
is_hip,
@@ -57,6 +63,7 @@ from sglang.srt.utils import (
next_power_of_2,
round_up,
set_weight_attrs,
use_intel_amx_backend,
)
from sglang.srt.utils.common import get_bool_env_var
from sglang.srt.utils.custom_op import register_custom_op
@@ -138,9 +145,11 @@ if TYPE_CHECKING:
StandardDispatchOutput,
)
_is_cpu = is_cpu()
_is_hip = is_hip()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
_is_cpu_amx_available = cpu_has_amx_support()
_sm120_mxfp4_min_warps_patched = False
if _is_hip:
@@ -849,6 +858,30 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
self.w2_weight_triton_tensor = w2_weight
del layer.w13_weight
del layer.w2_weight
elif _is_cpu and _is_cpu_amx_available:
_amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
if use_intel_amx_backend(layer):
packed_w13_weight_scale = torch.ops.sgl_kernel.convert_scale_packed(
layer.w13_weight_scale
)
packed_w2_weight_scale = torch.ops.sgl_kernel.convert_scale_packed(
layer.w2_weight_scale
)
layer.w13_weight_scale = Parameter(
packed_w13_weight_scale, requires_grad=False
)
layer.w2_weight_scale = Parameter(
packed_w2_weight_scale, requires_grad=False
)
if hasattr(layer, "w13_weight_bias"):
layer.w13_weight_bias = Parameter(
layer.w13_weight_bias.float(), requires_grad=False
)
if hasattr(layer, "w2_weight_bias"):
layer.w2_weight_bias = Parameter(
layer.w2_weight_bias.float(), requires_grad=False
)
return
else:
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp
@@ -1107,6 +1140,33 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
if use_intel_amx_backend(layer):
from sglang.srt.layers.moe.topk import apply_topk_weights_cpu
topk_weights, topk_ids, _ = dispatch_output.topk_output
x, topk_weights = apply_topk_weights_cpu(
self.moe_runner_config.apply_router_weight_on_input, topk_weights, x
)
output = torch.ops.sgl_kernel.fused_experts_cpu(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_ids,
False, # inplace See [Note] inplace should be False in fused_experts.
CPUQuantMethod.MXFP4,
layer.w13_weight_scale, # w1_scale
layer.w2_weight_scale, # w2_scale
None, # w1_zp
None, # w2_zp
None, # block_size
getattr(layer, "w13_weight_bias", None),
getattr(layer, "w2_weight_bias", None),
layer.moe_runner_config.gemm1_alpha,
layer.moe_runner_config.gemm1_clamp_limit,
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)
if self.use_marlin:
assert TopKOutputChecker.format_is_standard(topk_output)
python/sglang/srt/layers/quantization/unquant.py
+12
View File
@@ -268,6 +268,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
# Pack weight for get better performance on CPU
if _is_cpu and _is_cpu_amx_available:
_amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
if hasattr(layer, "w13_weight_bias"):
layer.w13_weight_bias = Parameter(
layer.w13_weight_bias.float(), requires_grad=False
)
if hasattr(layer, "w2_weight_bias"):
layer.w2_weight_bias = Parameter(
layer.w2_weight_bias.float(), requires_grad=False
)
if (
self.use_deep_gemm
@@ -579,6 +587,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
None, # w1_zp
None, # w2_zp
None, # block_size
getattr(layer, "w13_weight_bias", None),
getattr(layer, "w2_weight_bias", None),
layer.moe_runner_config.gemm1_alpha,
layer.moe_runner_config.gemm1_clamp_limit,
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/layers/quantization/w8a8_int8.py
+4
View File
@@ -375,6 +375,10 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
None, # w1_zp
None, # w2_zp
None, # block_size
None, # w1 bias
None, # w3 bias
None, # alpha
None, # limit
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/models/deepseek.py
+4
View File
@@ -207,6 +207,10 @@ class DeepseekMoE(nn.Module):
None, # w1_zp
None, # w2_zp
None, # block_size
None, # w1_bias
None, # w2_bias
None, # alpha
None, # limit
True, # is_vnni
)
else:
python/sglang/srt/models/gpt_oss.py
+20 -1
View File
@@ -81,6 +81,7 @@ from sglang.srt.utils import (
add_prefix,
get_cuda_version,
is_blackwell_supported,
is_cpu,
is_cuda,
is_flashinfer_available,
is_npu,
@@ -89,6 +90,7 @@ from sglang.srt.utils import (
)
from sglang.srt.utils.custom_op import register_custom_op
_is_cpu = is_cpu()
_is_npu = is_npu()
_is_cuda = is_cuda()
_is_tinygemm_supported = (
@@ -881,7 +883,8 @@ class GptOssForCausalLM(nn.Module):
moe_ep_rank_end = (moe_ep_rank + 1) * moe_num_local_experts
for name, weight in weights:
weight = weight.cuda()
if _is_cuda:
weight = weight.cuda()
if "gate_up_proj_blocks" in name:
# Handle MLP gate and up projection weights
@@ -1163,6 +1166,22 @@ class GptOssForCausalLM(nn.Module):
param = params_dict[name]
if "sinks" in name:
start = get_attention_tp_rank() * param.numel()
tp_size = get_tensor_model_parallel_world_size()
full_shard_size = param.numel() * tp_size
# This handles TP padding: if the checkpoint dim is not divisible by tp_size,
# the last TP shard extends beyond `loaded_weight`, pad with zeros before slicing.
if (
_is_cpu
and full_shard_size > loaded_weight.size(0)
and start + param.numel() >= loaded_weight.size(0)
):
pad_size = start + param.numel() - loaded_weight.size(0)
pad_tensor = torch.zeros(pad_size).to(
loaded_weight.dtype
)
loaded_weight = torch.cat(
[loaded_weight, pad_tensor], dim=0
).to(loaded_weight.dtype)
param.data.copy_(
loaded_weight[start : start + param.numel()]
)
python/sglang/srt/server_args.py
+4
View File
@@ -2059,6 +2059,8 @@ class ServerArgs:
self.attention_backend = "trtllm_mha"
elif is_sm90_supported():
self.attention_backend = "fa3"
elif is_cpu() and cpu_has_amx_support():
self.attention_backend = "intel_amx"
elif is_xpu():
self.attention_backend = "intel_xpu"
elif is_hip():
@@ -2084,6 +2086,7 @@ class ServerArgs:
"fa3",
"fa4",
"ascend",
"intel_amx",
"intel_xpu",
"aiter",
]
@@ -2147,6 +2150,7 @@ class ServerArgs:
self.ep_size == 1
and is_triton_kernels_available()
and self.quantization is None
and not (is_cpu() and cpu_has_amx_support())
):
# The triton_kernels package segfaults on Blackwell (B200)
# with NVIDIA driver >= 595. Fall back to triton backend.