From 0ac23cffacc22ff7bcbe32eee425cfaa6c56fffa Mon Sep 17 00:00:00 2001 From: Chandrakant Khandelwal Date: Wed, 29 Apr 2026 07:51:52 +0530 Subject: [PATCH] Add intel_xpu as backend for GptOssForCausalLM, enabled for bf16 models with torch native backend (#12771) Co-authored-by: Ma Mingfei --- .../sglang/srt/layers/moe/fused_moe_native.py | 40 ++++++++++++++++++- python/sglang/srt/server_args.py | 15 +++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_native.py b/python/sglang/srt/layers/moe/fused_moe_native.py index 4a9070fe3..592205075 100644 --- a/python/sglang/srt/layers/moe/fused_moe_native.py +++ b/python/sglang/srt/layers/moe/fused_moe_native.py @@ -7,6 +7,7 @@ import torch from torch.nn import functional as F from sglang.srt.layers.activation import GeluAndMul, SiluAndMul +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import swiglu_with_alpha_and_limit from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.token_dispatcher import ( StandardCombineInput, @@ -76,6 +77,9 @@ def moe_forward_native( else: raise ValueError(f"Unsupported activation: {moe_runner_config.activation=}") + # Get bias terms if available + w13_bias = getattr(layer, "w13_weight_bias", None) + w2_bias = getattr(layer, "w2_weight_bias", None) outputs = [] start_idx = 0 for i, num_tokens in enumerate(tokens_per_expert): @@ -87,9 +91,43 @@ def moe_forward_native( layer_w13_weight = layer.w13_weight[i] layer_w2_weight = layer.w2_weight[i] + # Store original dtype + original_dtype = tokens_for_this_expert.dtype + + # Get bias terms if available for this expert + layer_w13_bias = w13_bias[i] if w13_bias is not None else None + layer_w2_bias = w2_bias[i] if w2_bias is not None else None + + # Apply w13 linear gate_up = F.linear(tokens_for_this_expert, layer_w13_weight) - gate_up = act(gate_up) + + # Add bias if present (for models like GPT-OSS) + if layer_w13_bias is not None: + gate_up_fp32 = gate_up.float() + layer_w13_bias + gate_up = gate_up_fp32.to(original_dtype) + + # Apply activation + if ( + moe_runner_config.activation == "silu" + and moe_runner_config.gemm1_alpha is not None + ): + assert moe_runner_config.gemm1_clamp_limit is not None + gate_up = swiglu_with_alpha_and_limit( + gate_up, + moe_runner_config.gemm1_alpha, + moe_runner_config.gemm1_clamp_limit, + ) + else: + gate_up = act(gate_up) + + # Apply w2 linear expert_out = F.linear(gate_up, layer_w2_weight) + + # Add bias if present (for models like GPT-OSS) + if layer_w2_bias is not None: + expert_out = expert_out.float() + layer_w2_bias + expert_out = expert_out.to(original_dtype) + outputs.append(expert_out) start_idx = end_idx diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 817ae43fb..5e3f592d0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1860,17 +1860,32 @@ class ServerArgs: self.attention_backend = "trtllm_mha" elif is_sm90_supported(): self.attention_backend = "fa3" + elif is_xpu(): + self.attention_backend = "intel_xpu" elif is_hip(): self.attention_backend = "aiter" else: self.attention_backend = "triton" + if is_xpu(): + # Check for bf16 dtype on Intel XPU + if self.dtype == "auto": + logger.warning( + "GptOssForCausalLM on Intel XPU currently supports bfloat16 dtype only" + ) + elif self.dtype not in ["bfloat16"]: + raise NotImplementedError( + f"GptOssForCausalLM on Intel XPU only supports bfloat16 dtype, " + f"but got '{self.dtype}'. Please use --dtype bfloat16 or remove --dtype to use auto." + ) + supported_backends = [ "triton", "trtllm_mha", "fa3", "fa4", "ascend", + "intel_xpu", "aiter", ] prefill_attn_backend, decode_attn_backend = self.get_attention_backends()