diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index 79f3aa2e9..17168d414 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -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, diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py b/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py index 2840eb8fc..12d1d5b1d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index 14fae2623..583b23ada 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -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), diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 255c2e591..a44fe5dc4 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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: diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 4be793fad..898d0c4b0 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -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, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index cb18136be..9281163ef 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -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: diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 6dde83ab7..c75eae0b5 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -74,6 +74,7 @@ suite_xeon = { suite_xpu = { "per-commit-xpu": [ TestFile("xpu/test_intel_xpu_backend.py"), + TestFile("xpu/test_deepseek_ocr.py"), ], } diff --git a/test/srt/xpu/test_deepseek_ocr.py b/test/srt/xpu/test_deepseek_ocr.py new file mode 100644 index 000000000..9d5da10e3 --- /dev/null +++ b/test/srt/xpu/test_deepseek_ocr.py @@ -0,0 +1,121 @@ +""" +python3 -m unittest test_deepseek_ocr.py +""" + +import json +import os +import unittest + +import requests +from transformers import AutoTokenizer + +from sglang.srt.utils import kill_process_tree +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestDeepSeekOCR(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = "deepseek-ai/DeepSeek-OCR" + cls.tokenizer = AutoTokenizer.from_pretrained(cls.model, use_fast=False) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.common_args = [ + "--device", + "xpu", + "--attention-backend", + "intel_xpu", + ] + os.environ["SGLANG_USE_SGL_XPU"] = "1" + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *cls.common_args, + ], + ) + + @classmethod + def tearDownClass(cls): + """Fixture that is run once after all tests in the class.""" + kill_process_tree(cls.process.pid) + + def get_request_json(self, max_new_tokens=32, n=1): + response = requests.post( + self.base_url + "/generate", + json={ + "text": "\n<|grounding|>Convert the document to pure text.", + "image_data": "../../examples/assets/example_image.png", + "sampling_params": { + "temperature": 0 if n == 1 else 0.5, + "max_new_tokens": max_new_tokens, + }, + }, + ) + return response.json() + + def run_decode( + self, + max_new_tokens=128, + n=1, + ): + + ret = self.get_request_json(max_new_tokens=max_new_tokens, n=n) + print(json.dumps(ret, indent=2)) + + def assert_one_item(item): + if item["meta_info"]["finish_reason"]["type"] == "stop": + self.assertEqual( + item["meta_info"]["finish_reason"]["matched"], + self.tokenizer.eos_token_id, + ) + elif item["meta_info"]["finish_reason"]["type"] == "length": + self.assertEqual( + len(item["output_ids"]), item["meta_info"]["completion_tokens"] + ) + self.assertEqual(len(item["output_ids"]), max_new_tokens) + + # Determine whether to assert a single item or multiple items based on n + if n == 1: + assert_one_item(ret) + else: + self.assertEqual(len(ret), n) + for i in range(n): + assert_one_item(ret[i]) + + print("=" * 100) + + def test_moe(self): + self.run_decode() + + +class TestDeepSeekOCRTriton(TestDeepSeekOCR): + @classmethod + def setUpClass(cls): + cls.model = "deepseek-ai/DeepSeek-OCR" + cls.tokenizer = AutoTokenizer.from_pretrained(cls.model, use_fast=False) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.common_args = [ + "--device", + "xpu", + "--attention-backend", + "intel_xpu", + ] + os.environ["SGLANG_USE_SGL_XPU"] = "0" + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *cls.common_args, + ], + ) + + +if __name__ == "__main__": + unittest.main()