From 9cee0a31d1a7f39526ba02b8ffb9e3b2854fbef2 Mon Sep 17 00:00:00 2001 From: K Chiranjeevi Date: Fri, 28 Aug 2026 14:48:03 +0530 Subject: [PATCH] [XPU] Lazily import tvm_ffi-dependent all_reduce kernel in minimax_m2 (#35290) Signed-off-by: Chiranjeevi Koppula --- python/sglang/srt/models/minimax_m2.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 431ff22b6..8d959e8d8 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -27,10 +27,6 @@ from torch import nn from transformers import PretrainedConfig from sglang.kernels.kernel_api_logging import debug_kernel_api -from sglang.kernels.ops.communication.all_reduce import ( - fused_parallel_qknorm, - get_fused_parallel_qknorm_max_occupancy, -) from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( get_pp_group, @@ -103,6 +99,7 @@ from sglang.srt.utils import ( is_cuda, is_non_idle_and_non_empty, is_npu, + is_xpu, make_layers, ) from sglang.srt.utils.custom_op import register_custom_op @@ -113,6 +110,13 @@ _is_cpu = is_cpu() _is_amx_available = cpu_has_amx_support() _is_cuda = is_cuda() _is_npu = is_npu() +_is_xpu = is_xpu() + +if not _is_xpu: + from sglang.kernels.ops.communication.all_reduce import ( + fused_parallel_qknorm, + get_fused_parallel_qknorm_max_occupancy, + ) if _is_npu: from sgl_kernel_npu.norm.split_qkv_tp_rmsnorm_rope import split_qkv_tp_rmsnorm_rope