[XPU] Support GPT-OSS MXFP4 checkpoints on Intel XPU (#35751)

Co-authored-by: Meng, Hengyu <hengyu.meng@intel.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Juan Muneton
2026-09-04 14:20:28 +08:00
committed by GitHub
co-authored by Meng, Hengyu Claude Opus 5
parent d122ca99b2
commit 72078cd7f5
5 changed files with 115 additions and 2 deletions
@@ -1151,6 +1151,14 @@ def fused_moe(
a1_scale=a1_scale,
a2_scale=a2_scale,
block_shape=block_shape,
# These were previously dropped, which silently computed a plain
# silu*up for GPT-OSS-style experts instead of the clamped
# gate*sigmoid(gate*alpha)*(up+1) the config asks for.
activation=moe_runner_config.activation,
routed_scaling_factor=moe_runner_config.routed_scaling_factor,
gemm1_alpha=moe_runner_config.gemm1_alpha,
gemm1_limit=moe_runner_config.gemm1_clamp_limit,
swiglu_limit=moe_runner_config.swiglu_limit,
)
return fused_experts(
@@ -92,7 +92,11 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
}
if is_cpu() or is_cuda() or _is_gfx95_supported:
# On XPU the OCP-MoE `Mxfp4Config` path is served by the sgl-kernel-xpu grouped
# GEMM, which consumes the packed e2m1 + ue8m0 g32 checkpoint layout directly.
# Other backends without that kernel keep the existing "unknown quantization
# method" error rather than falling through to a bf16 upcast.
if is_cpu() or is_cuda() or _is_gfx95_supported or is_xpu():
BASE_QUANTIZATION_METHODS.update(
{
"mxfp4": Mxfp4Config,
@@ -61,6 +61,7 @@ from sglang.srt.utils import (
is_gfx95_supported,
is_hip,
is_triton_kernels_available,
is_xpu,
next_power_of_2,
round_up,
set_weight_attrs,
@@ -205,6 +206,7 @@ if TYPE_CHECKING:
_is_cpu = is_cpu()
_is_hip = is_hip()
_is_xpu = is_xpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_aiter_k3_opt = _use_aiter and get_bool_env_var("SGLANG_AITER_K3_OPT")
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
@@ -520,6 +522,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
intermediate_size_per_partition_after_pad
- layer.intermediate_size_per_partition
)
elif _is_xpu:
# The XPU grouped GEMM recovers K/N from the packed weight shapes and
# the group size from the scale shape, so it takes the checkpoint
# dims. Keep this ahead of the triton_kernels branch so an installed
# triton_kernels does not pad the XPU layout.
#
# Align to mxfp4_block anyway: gpt_oss.py shards the intermediate by
# whole mxfp4 blocks (ceil(blocks / tp) * 32), so a rank's slice can
# exceed intermediate_size / tp -- 736 vs 720 for gpt-oss-20b at
# tp=4, which overflows an unaligned buffer during load. round_up to
# 32 reproduces the loader's shard exactly, so no rows are wasted and
# the kernel still sees the true dims.
intermediate_size_per_partition_after_pad = round_up(
intermediate_size_per_partition, mxfp4_block
)
elif has_triton_kernels:
intermediate_size_per_partition_after_pad = round_up(
intermediate_size_per_partition, triton_kernels_padding_alignment
@@ -1057,6 +1074,24 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
layer.w2_weight_bias.float(), requires_grad=False
)
return
elif _is_xpu:
# sgl-kernel-xpu's W4A16 grouped GEMM consumes the checkpoint MXFP4
# layout: packed e2m1 [E, N, K/2] plus N-outer ue8m0 scales
# [E, N, K/32] uint8, with GPT-OSS's interleaved
# [gate_0, up_0, gate_1, up_1, ...] w13 row order (which is exactly
# what the swiglu epilogue expects). Scales and biases are already in
# the expected dtypes (uint8 / bf16 -- the launcher promotes bias to
# fp32 since the kernel accumulates it in fp32), so the only step is
# reinterpreting the packed nibbles as int8, matching the dtype the
# kernel keys the 4-bit path on. That is a free view, and crucially
# there is no bf16 upcast -- the whole point of MXFP4 on XPU.
layer.w13_weight = Parameter(
layer.w13_weight.data.view(torch.int8), requires_grad=False
)
layer.w2_weight = Parameter(
layer.w2_weight.data.view(torch.int8), requires_grad=False
)
return
else:
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp
@@ -1576,6 +1611,35 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
)
return StandardCombineInput(hidden_states=output)
if _is_xpu:
# sgl-kernel-xpu path: moe_grouped_mm_nt_xe20_w4a16 consumes the
# packed MXFP4 weights directly, so no dequantization happens.
from sgl_kernel import fused_experts as sgl_fused_experts
assert TopKOutputChecker.format_is_standard(topk_output)
topk_weights, topk_ids, _ = topk_output
moe_runner_config = self.moe_runner_config
output = sgl_fused_experts(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_ids,
b1=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
use_mxfp4_w4a16=True,
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
activation=moe_runner_config.activation,
routed_scaling_factor=moe_runner_config.routed_scaling_factor,
# GPT-OSS clamped swiglu: gate*sigmoid(gate*alpha)*(up+1). Passing
# gemm1_alpha selects it, and gemm1_limit is required with it.
gemm1_alpha=moe_runner_config.gemm1_alpha,
gemm1_limit=moe_runner_config.gemm1_clamp_limit,
swiglu_limit=moe_runner_config.swiglu_limit,
)
return StandardCombineInput(hidden_states=output)
if self.use_marlin:
assert TopKOutputChecker.format_is_standard(topk_output)
return self._apply_marlin(layer, dispatch_output)