[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:
co-authored by
Meng, Hengyu
Claude Opus 5
parent
d122ca99b2
commit
72078cd7f5
@@ -47,7 +47,7 @@ The following table summarizes quantization method support across NVIDIA and AMD
|
|||||||
<td>Yes</td>
|
<td>Yes</td>
|
||||||
<td>Yes</td>
|
<td>Yes</td>
|
||||||
<td>Yes (A5)</td>
|
<td>Yes (A5)</td>
|
||||||
<td>On GPU: requires CDNA3/CDNA4 with MXFP support (uses Aiter). On Ascend NPU (A5): W4A4 MXFP4 for Qwen3 dense and MoE LLMs (MXFP4 weights + activations) — dense models support online dual-level MXFP4; offline <code>W4A4_MXFP4</code> dense and MoE checkpoints (single-level) are auto-detected via <code>modelslim</code></td>
|
<td>On GPU: requires CDNA3/CDNA4 with MXFP support (uses Aiter). On Ascend NPU (A5): W4A4 MXFP4 for Qwen3 dense and MoE LLMs (MXFP4 weights + activations) — dense models support online dual-level MXFP4; offline <code>W4A4_MXFP4</code> dense and MoE checkpoints (single-level) are auto-detected via <code>modelslim</code>. On Intel GPUs (XPU): W4A16 MoE experts on Xe2/BMG via <code>sgl-kernel-xpu</code>, enabled automatically with <code>--device xpu</code> (see <a href="../hardware-platforms/xpu">Intel GPUs (XPU)</a>)</td>
|
||||||
</tr>
|
</tr>
|
||||||
<tr>
|
<tr>
|
||||||
<td><code>mxfp8</code></td>
|
<td><code>mxfp8</code></td>
|
||||||
|
|||||||
@@ -41,6 +41,9 @@ A list of LLMs have been optimized on Intel GPU, and more are on the way:
|
|||||||
**Note:** The model identifiers listed in the table above
|
**Note:** The model identifiers listed in the table above
|
||||||
have been verified on [Intel® Arc™ B580 Graphics](https://www.intel.com/content/www/us/en/products/sku/241598/intel-arc-b580-graphics/specifications.html).
|
have been verified on [Intel® Arc™ B580 Graphics](https://www.intel.com/content/www/us/en/products/sku/241598/intel-arc-b580-graphics/specifications.html).
|
||||||
|
|
||||||
|
Quantized MoE models are covered separately in
|
||||||
|
[MXFP4 MoE Quantization](#mxfp4-moe-quantization) below.
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|
||||||
### Install From Source
|
### Install From Source
|
||||||
@@ -116,6 +119,40 @@ sglang serve \
|
|||||||
--page-size \ # intel_xpu attention backend supports [32, 64, 128]
|
--page-size \ # intel_xpu attention backend supports [32, 64, 128]
|
||||||
```
|
```
|
||||||
|
|
||||||
|
## MXFP4 MoE Quantization
|
||||||
|
|
||||||
|
Native MXFP4 MoE checkpoints (OCP microscaling FP4: packed `e2m1` weights plus
|
||||||
|
per-32-element `ue8m0` block scales) run on Intel GPUs through the
|
||||||
|
`sgl-kernel-xpu` W4A16 grouped GEMM. The expert weights stay in the checkpoint's
|
||||||
|
packed layout end to end — there is no dequantization to BF16 — so GPT-OSS-20B
|
||||||
|
loads in roughly 13 GB rather than the ~42 GB a BF16 upcast would need, which is
|
||||||
|
what lets it fit on a single 24 GB card.
|
||||||
|
|
||||||
|
The `mxfp4` method is registered automatically on `--device xpu`, which already
|
||||||
|
requires `sgl-kernel-xpu` — no extra flag is needed:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
sglang serve \
|
||||||
|
--model-path openai/gpt-oss-20b \
|
||||||
|
--device xpu \
|
||||||
|
--attention-backend intel_xpu
|
||||||
|
```
|
||||||
|
|
||||||
|
**Tested models:**
|
||||||
|
|
||||||
|
| Model | Notes |
|
||||||
|
|:---|:---|
|
||||||
|
| [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | 32 experts, top-k 4, hidden 2880, intermediate 2880; clamped-swiglu activation (`gemm1_alpha` 1.702, limit 7.0) with expert biases |
|
||||||
|
|
||||||
|
**Requirements and limitations:**
|
||||||
|
|
||||||
|
| Item | Status |
|
||||||
|
|---|---|
|
||||||
|
| GPU architecture | Xe2 / BMG only (Intel® Arc™ B-Series and Arc™ Pro B-Series) |
|
||||||
|
| Quantized layers | Fused MoE experts only; attention, router, embeddings and LM head stay BF16 (per the checkpoint's `modules_to_not_convert`) |
|
||||||
|
| Activations | BF16 (W4A16). Clamped swiglu (GPT-OSS) and plain SiLU are supported |
|
||||||
|
| Expert parallelism | Not yet validated on this path |
|
||||||
|
|
||||||
## Benchmarking with Requests
|
## Benchmarking with Requests
|
||||||
|
|
||||||
You can benchmark the performance via the `bench_serving` script.
|
You can benchmark the performance via the `bench_serving` script.
|
||||||
|
|||||||
@@ -1151,6 +1151,14 @@ def fused_moe(
|
|||||||
a1_scale=a1_scale,
|
a1_scale=a1_scale,
|
||||||
a2_scale=a2_scale,
|
a2_scale=a2_scale,
|
||||||
block_shape=block_shape,
|
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(
|
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(
|
BASE_QUANTIZATION_METHODS.update(
|
||||||
{
|
{
|
||||||
"mxfp4": Mxfp4Config,
|
"mxfp4": Mxfp4Config,
|
||||||
|
|||||||
@@ -61,6 +61,7 @@ from sglang.srt.utils import (
|
|||||||
is_gfx95_supported,
|
is_gfx95_supported,
|
||||||
is_hip,
|
is_hip,
|
||||||
is_triton_kernels_available,
|
is_triton_kernels_available,
|
||||||
|
is_xpu,
|
||||||
next_power_of_2,
|
next_power_of_2,
|
||||||
round_up,
|
round_up,
|
||||||
set_weight_attrs,
|
set_weight_attrs,
|
||||||
@@ -205,6 +206,7 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
_is_cpu = is_cpu()
|
_is_cpu = is_cpu()
|
||||||
_is_hip = is_hip()
|
_is_hip = is_hip()
|
||||||
|
_is_xpu = is_xpu()
|
||||||
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
_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")
|
_aiter_k3_opt = _use_aiter and get_bool_env_var("SGLANG_AITER_K3_OPT")
|
||||||
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
|
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
|
||||||
@@ -520,6 +522,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
intermediate_size_per_partition_after_pad
|
intermediate_size_per_partition_after_pad
|
||||||
- layer.intermediate_size_per_partition
|
- 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:
|
elif has_triton_kernels:
|
||||||
intermediate_size_per_partition_after_pad = round_up(
|
intermediate_size_per_partition_after_pad = round_up(
|
||||||
intermediate_size_per_partition, triton_kernels_padding_alignment
|
intermediate_size_per_partition, triton_kernels_padding_alignment
|
||||||
@@ -1057,6 +1074,24 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
layer.w2_weight_bias.float(), requires_grad=False
|
layer.w2_weight_bias.float(), requires_grad=False
|
||||||
)
|
)
|
||||||
return
|
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:
|
else:
|
||||||
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp
|
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp
|
||||||
|
|
||||||
@@ -1576,6 +1611,35 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
)
|
)
|
||||||
return StandardCombineInput(hidden_states=output)
|
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:
|
if self.use_marlin:
|
||||||
assert TopKOutputChecker.format_is_standard(topk_output)
|
assert TopKOutputChecker.format_is_standard(topk_output)
|
||||||
return self._apply_marlin(layer, dispatch_output)
|
return self._apply_marlin(layer, dispatch_output)
|
||||||
|
|||||||
Reference in New Issue
Block a user