diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 8597f89ee..f0a171d4c 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -11,6 +11,7 @@ from sglang.kernels.jit.utils import ( load_jit, make_cpp_args, ) +from sglang.srt.utils import is_xpu from .utils import make_name @@ -58,6 +59,10 @@ def topk_transform_512( torch.ops.sgl_kernel.deepseek_v4_topk_transform_512( scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices ) + elif is_xpu(): + torch.ops.sgl_kernel.topk_transform( + scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices + ) else: module = _jit_topk_v1_module() module.topk_transform( @@ -111,6 +116,15 @@ def topk_transform_ragged_v2( They are invalid for that row and the buffer must have no other consumer. ``seq_lens`` entries must be NON-NEGATIVE, as for the paged entry point. """ + if is_xpu(): + torch.ops.sgl_kernel.topk_transform_ragged( + scores, + seq_lens, + out_indices, + out_offsets, + row_starts, + ) + return module = _jit_topk_v2_module() module.topk_transform_ragged(scores, seq_lens, row_starts, out_offsets, out_indices) @@ -143,6 +157,16 @@ def topk_transform_512_v2( the valid way to express "no tokens": the row takes the trivial path and the output is all -1. """ + if is_xpu(): + torch.ops.sgl_kernel.topk_transform_paged( + scores, + seq_lens, + page_tables, + out_page_indices, + page_size, + metadata, + ) + return module = _jit_topk_v2_module() module.topk_transform_paged( scores, diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index d245ddce3..60f57e1d8 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -152,9 +152,9 @@ class PagedIndexerMetadata: assert isinstance(self.deep_gemm_metadata, torch.Tensor) - from sglang.kernels.ops.attention.dsv4 import plan_topk_v2 + if envs.SGLANG_OPT_USE_TOPK_V2.get() and not is_xpu(): + from sglang.kernels.ops.attention.dsv4 import plan_topk_v2 - if envs.SGLANG_OPT_USE_TOPK_V2.get(): self.topk_metadata = plan_topk_v2(self.c4_seq_lens) else: self.topk_metadata = torch.empty((0,))