[XPU]Enable HiSparse hierarchical sparse KV cache on Intel XPU (#32792)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/srt/managers/hisparse_coordinator.py

@@ -5,11 +5,27 @@ from typing import Dict, List, NamedTuple, Optional, Tuple, Union
 
 import torch
 
-from sglang.kernels.ops.kvcache.hisparse import (
-    copy_cache_planned_mla,
-    load_cache_to_device_buffer_dsv4_mla,
-    load_cache_to_device_buffer_mla,
-)
+from sglang.srt.utils import get_device_module, is_hip, is_xpu
+
+if is_xpu():
+    from sgl_kernel import (
+        load_cache_to_device_buffer_dsv4_mla,
+        load_cache_to_device_buffer_mla,
+    )
+
+    def copy_cache_planned_mla(*args, **kwargs):
+        raise RuntimeError(
+            "HiSparse shared-index prefetch is unsupported on XPU: "
+            "copy_cache_planned_mla has no AOT sgl_kernel implementation."
+        )
+
+else:
+    from sglang.kernels.ops.kvcache.hisparse import (
+        copy_cache_planned_mla,
+        load_cache_to_device_buffer_dsv4_mla,
+        load_cache_to_device_buffer_mla,
+    )
+
 from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import Req
@@ -23,11 +39,11 @@ from sglang.srt.mem_cache.hisparse_memory_pool import (
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.mem_cache.memory_pool_host import DeepSeekV4PagedHostPool
 from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
-from sglang.srt.utils import get_device_module, is_hip
 
 device_module = get_device_module()
 
 _is_hip = is_hip()
+_is_xpu = is_xpu()
 
 logger = logging.getLogger(__name__)
 
@@ -133,6 +149,12 @@ class HiSparseCoordinator:
         # Timing probe: skip the host->device KV bytes to measure the "IO is
         # free" floor. Produces garbage output; benchmarking only.
         self.skip_io = envs.SGLANG_DEBUG_HISPARSE_SKIP_IO.get()
+        if _is_xpu and self.skip_io:
+            raise ValueError(
+                "SGLANG_DEBUG_HISPARSE_SKIP_IO is unsupported on XPU: the AOT swap-in "
+                "ops do not accept skip_io, so timings would silently include the KV "
+                "copy."
+            )
         self.compress_ratio = self.token_to_kv_pool_allocator.compress_ratio
 
         self.is_dsv4_hisparse = isinstance(
@@ -281,6 +303,12 @@ class HiSparseCoordinator:
                 layer_num,
             )
             shared_index_layers = None
+        if shared_index_layers is not None and _is_xpu:
+            logger.warning(
+                "HiSparse shared-index prefetch disabled on XPU: "
+                "copy_cache_planned_mla is unavailable; using synchronous swap-in."
+            )
+            shared_index_layers = None
         self._is_shared_index_layer = list(shared_index_layers or [False] * layer_num)
         self.enable_prefetch = any(self._is_shared_index_layer)
         self._prefetch_groups, self._prefetch_slot = _build_prefetch_groups(
@@ -964,6 +992,7 @@ class HiSparseCoordinator:
             if record_plan
             else {}
         )
+        skip_io_kwargs = {} if _is_xpu else dict(skip_io=self.skip_io)
         swap_in_fn(
             top_k_tokens=top_k_result,
             device_buffer_tokens=self.req_device_buffer_tokens[layer_id],
@@ -981,7 +1010,7 @@ class HiSparseCoordinator:
             page_size=1,
             block_size=self.swap_in_block_size,
             num_real_reqs=self.num_real_reqs,
-            skip_io=self.skip_io,
+            **skip_io_kwargs,
             **plan,
         )
         return top_k_indices

python/sglang/srt/mem_cache/hisparse_memory_pool.py

@@ -10,21 +10,22 @@ from sglang.kernels.ops.kvcache.hisparse_slot_mapping import (
 )
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
-from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils import is_cuda, is_hip, is_xpu
 
 logger = logging.getLogger(__name__)
 
-# sgl_kernel.kvcacheio is only available in CUDA/ROCm sgl-kernel builds (not XPU/MPS/NPU/CPU).
+# sgl_kernel.kvcacheio is only available in CUDA/ROCm/XPU sgl-kernel builds (not MPS/NPU/CPU).
 _is_cuda = is_cuda()
 _is_hip = is_hip()
-if _is_cuda or _is_hip:
+_is_xpu = is_xpu()
+if _is_cuda or _is_hip or _is_xpu:
     from sgl_kernel.kvcacheio import transfer_kv_all_layer_mla
 else:
 
     def transfer_kv_all_layer_mla(*args, **kwargs):
         raise RuntimeError(
-            "HiSparse device KV transfer requires sgl_kernel.kvcacheio (CUDA/ROCm). "
-            "It is not available on this backend."
+            "HiSparse device KV transfer requires sgl_kernel.kvcacheio "
+            "(CUDA/ROCm/XPU). It is not available on this backend."
         )
 
 

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -12,7 +12,6 @@ from sglang.kernels.ops.kvcache.hicache import (
 from sglang.kernels.ops.kvcache.hicache import (
     transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
 )
-from sglang.kernels.ops.kvcache.hisparse import transfer_cache_dsv4_mla
 from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
 
 _is_cuda = is_cuda()
@@ -20,7 +19,11 @@ _is_hip = is_hip()
 _is_npu = is_npu()
 _is_xpu = is_xpu()
 _is_mps = is_mps()
-if _is_cuda or _is_hip:
+if _is_xpu:
+    from sgl_kernel import transfer_cache_dsv4_mla
+else:
+    from sglang.kernels.ops.kvcache.hisparse import transfer_cache_dsv4_mla
+if _is_cuda or _is_hip or _is_xpu:
     from sgl_kernel.kvcacheio import (
         transfer_kv_all_layer_direct_lf_pf,
         transfer_kv_all_layer_mla,

python/sglang/srt/mem_cache/pool_host/common.py

@@ -315,5 +315,6 @@ ALLOC_MEMORY_FUNCS = defaultdict(
     {
         "npu": alloc_with_pin_memory,
         "musa": alloc_with_pin_memory,
+        "xpu": alloc_with_pin_memory,
     },
 )

python/sglang/srt/mem_cache/pool_host/mla.py

@@ -47,7 +47,7 @@ _is_hip = is_hip()
 _is_npu = is_npu()
 _is_xpu = is_xpu()
 _is_mps = is_mps()
-if _is_cuda or _is_hip:
+if _is_cuda or _is_hip or _is_xpu:
     from sgl_kernel.kvcacheio import (
         transfer_kv_all_layer_direct_lf_pf,
         transfer_kv_all_layer_mla,