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

[XPU] Integrate MoE and minor improvements in XPU attention backend (#13561)

Browse Source
This commit is contained in:
Meng, Hengyu
2026-02-05 15:09:59 +08:00
committed by GitHub
parent dff3ba202a
commit 368936a62b
8 changed files with 233 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View File

@@ -20,6 +20,8 @@ from sglang.srt.utils import (
is_cpu,
is_cuda,
is_hip,
is_xpu,
use_intel_xpu_backend,
)
from sglang.srt.utils.custom_op import register_custom_op
@@ -40,6 +42,8 @@ _is_cuda = is_cuda()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_is_xpu = is_xpu()
_use_sgl_xpu = use_intel_xpu_backend()
if _is_cuda:
from sgl_kernel import gelu_and_mul, moe_sum_reduce, silu_and_mul
@@ -55,6 +59,8 @@ elif _is_hip:
raise ImportError("aiter is required when SGLANG_USE_AITER is set to True")
else:
from vllm import _custom_ops as vllm_ops
elif _is_xpu:
from sgl_kernel import moe_sum_reduce, silu_and_mul
padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
@@ -493,7 +499,7 @@ def fused_experts_impl(
intermediate_cache2 = _swiglu_silu_clamp_mul(
intermediate_cache1.view(-1, N), gemm1_limit
)
elif _is_cuda or _is_hip:
elif _is_cuda or _is_hip or _is_xpu:
if not filter_expert:
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
else:
@@ -621,6 +627,12 @@ def fused_experts_impl(
out_hidden_states[begin_chunk_idx:end_chunk_idx],
routed_scaling_factor,
)
elif _is_xpu:
moe_sum_reduce(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states[begin_chunk_idx:end_chunk_idx],
routed_scaling_factor,
)
else:
vllm_ops.moe_sum(
intermediate_cache3.view(*intermediate_cache3.shape),
@@ -690,6 +702,27 @@ def fused_moe(
Returns:
- torch.Tensor: The output tensor after applying the MoE layer.
"""
if _use_sgl_xpu:
topk_weight, topk_ids, _ = topk_output
from sgl_kernel import fused_experts as sgl_fused_experts
return sgl_fused_experts(
hidden_states,
w1,
w2,
topk_weight,
topk_ids,
b1=b1,
b2=b2,
use_fp8_w8a8=use_fp8_w8a8,
w1_scale=w1_scale,
w2_scale=w2_scale,
w1_zp=w1_zp,
w2_zp=w2_zp,
a1_scale=a1_scale,
a2_scale=a2_scale,
block_shape=block_shape,
)
return fused_experts(
hidden_states,

5
python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py
View File

@@ -5,12 +5,13 @@ from typing import Tuple
import torch
import triton
from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils import is_cuda, is_hip, is_xpu
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_xpu = is_xpu()
if _is_cuda or _is_hip:
if _is_cuda or _is_hip or _is_xpu:
from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size

16
python/sglang/srt/layers/moe/moe_runner/triton.py
View File

@@ -19,7 +19,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
register_pre_permute,
)
from sglang.srt.layers.moe.utils import MoeRunnerBackend
from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_hip
from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_hip, is_xpu
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher.standard import (
@@ -33,6 +33,7 @@ _is_cuda = is_cuda()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
_use_aiter = bool(int(os.getenv("SGLANG_USE_AITER", "0")))
_is_xpu = is_xpu()
_MOE_PADDING_SIZE = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
@@ -51,8 +52,11 @@ if _is_cuda or _is_hip:
from vllm import _custom_ops as vllm_ops # moe_sum
elif _is_cpu and _is_cpu_amx_available:
pass
elif _is_xpu:
from sgl_kernel import moe_sum_reduce, silu_and_mul
if _is_cuda or _is_hip:
if _is_cuda or _is_hip or _is_xpu:
from sgl_kernel import ( # noqa: F401
moe_align_block_size as sgl_moe_align_block_size,
)
@@ -211,7 +215,7 @@ class TritonRunnerCore(MoeRunnerCore):
intermediate_cache2 = _swiglu_silu_clamp_mul(
intermediate_cache1.view(-1, N), gemm1_limit
)
elif _is_cuda or _is_hip:
elif _is_cuda or _is_hip or _is_xpu:
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
else:
vllm_ops.silu_and_mul(
@@ -315,6 +319,12 @@ class TritonRunnerCore(MoeRunnerCore):
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
)
elif _is_xpu:
moe_sum_reduce(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
routed_scaling_factor,
)
else:
vllm_ops.moe_sum(
intermediate_cache3.view(*intermediate_cache3.shape),

1
python/sglang/srt/layers/moe/topk.py
View File

@@ -72,6 +72,7 @@ _is_cpu = is_cpu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_xpu = is_xpu()
_is_npu = is_npu()
_is_xpu = is_xpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _is_cuda:

49
python/sglang/srt/layers/quantization/unquant.py
View File

@@ -32,6 +32,7 @@ from sglang.srt.utils import (
next_power_of_2,
set_weight_attrs,
use_intel_amx_backend,
use_intel_xpu_backend,
)
if TYPE_CHECKING:
@@ -470,6 +471,54 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
)
return StandardCombineInput(hidden_states=output)
def forward_xpu(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
moe_runner_config = self.moe_runner_config
assert moe_runner_config.activation in [
"silu",
"gelu",
], f"activation = {moe_runner_config.activation} is not supported."
backend = self.runner.runner_backend
if use_intel_xpu_backend():
# sgl-kernel-xpu path
from sgl_kernel import fused_experts
topk_weights, topk_ids, _ = topk_output
output = 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),
activation=moe_runner_config.activation,
)
return StandardCombineInput(hidden_states=output)
else:
assert backend.is_triton()
assert (
moe_runner_config.activation == "silu"
), f"activation = {moe_runner_config.activation} is not supported \
for Triton PATH, please set ENV SGLANG_USE_SGL_XPU=1."
quant_info = TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
b13=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
)
return self.runner.run(dispatch_output, quant_info)
def forward_npu(
self,
layer: torch.nn.Module,

12
python/sglang/srt/utils/common.py
View File

@@ -299,6 +299,10 @@ def xpu_has_xmx_support():
return False
def use_intel_xpu_backend():
return get_bool_env_var("SGLANG_USE_SGL_XPU") and is_xpu()
@lru_cache(maxsize=1)
def is_flashinfer_available():
"""
@@ -2182,7 +2186,13 @@ def direct_register_custom_op(
try:
my_lib.define(op_name + schema_str)
my_lib.impl(op_name, op_func, "CUDA" if not is_npu() else "PrivateUse1")
if is_npu():
# https://github.com/sgl-project/sglang/pull/12287/files#r2499583982
my_lib.impl(op_name, op_func, "PrivateUse1")
elif is_xpu():
my_lib.impl(op_name, op_func, "XPU")
else:
my_lib.impl(op_name, op_func, "CUDA")
if fake_impl is not None:
my_lib._register_fake(op_name, fake_impl)
except RuntimeError as error: