From a5a9d66bafa930d78770232fd83dfb72777ac538 Mon Sep 17 00:00:00 2001 From: Rohit Kumar Singh <9626333+SKRohit@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:13:41 +0530 Subject: [PATCH] [XPU] Fix/kimi linear xpu (#34546) Co-authored-by: github-actions[bot] Co-authored-by: Singh Co-authored-by: Singh --- .../srt/layers/attention/linear/kernels/kda_triton.py | 7 +++++-- python/sglang/srt/layers/moe/topk.py | 3 ++- python/sglang/srt/models/kimi_linear.py | 4 ++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py index ebf06d49f..85774ec36 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py @@ -5,7 +5,7 @@ import torch from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( LinearAttnKernelBase, ) -from sglang.srt.utils import is_cpu, is_npu +from sglang.srt.utils import is_cpu, is_npu, is_xpu if not is_cpu(): from sglang.kernels.ops.attention.fla.fused_recurrent import ( @@ -23,7 +23,10 @@ if not is_cpu(): class TritonKDAKernel(LinearAttnKernelBase): """Triton-based kernel for KDA (Kimi Delta Attention) linear attention.""" - supports_packed_decode: bool = not is_cpu() and not is_npu() + # XPU has no tvm_ffi CUDA JIT kernel for KDA packed decode; route XPU to the + # non-packed Triton decode() path (fused_sigmoid_gating_delta_rule_update), + # the same fallback CPU/NPU use. Batched decode is handled via query_start_loc. + supports_packed_decode: bool = not is_cpu() and not is_npu() and not is_xpu() def packed_decode( self, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 951dc2238..e5511d15f 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1741,7 +1741,8 @@ def biased_grouped_topk_gpu( topk_indices, gating_output, renormalize, - correction_bias, + # The XPU topk_sigmoid AOT kernel requires an fp32 correction bias. + correction_bias.to(torch.float32), scale, ) diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index b8fc10636..506353b0e 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -49,7 +49,7 @@ from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA as KimiMLAAtten from sglang.srt.models.llama import LlamaMLP as KimiMLP from sglang.srt.models.transformers import maybe_prefix from sglang.srt.runtime_context import get_parallel, get_stream -from sglang.srt.utils import make_layers +from sglang.srt.utils import is_xpu, make_layers from sglang.srt.utils.common import BumpAllocator, add_prefix, set_weight_attrs @@ -550,7 +550,7 @@ class KimiLinearModel(nn.Module): else: self.embed_tokens = PPMissingLayer() - self.alt_stream = get_stream("alt") + self.alt_stream = None if is_xpu() else get_stream("alt") self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers,