Skip to content

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

Open
Amrutha-M05 wants to merge 2 commits into
sgl-project:mainfrom
Amrutha-M05:hisparse-xpu
Open

[XPU]Enable HiSparse hierarchical sparse KV cache on Intel XPU#32792
Amrutha-M05 wants to merge 2 commits into
sgl-project:mainfrom
Amrutha-M05:hisparse-xpu

Conversation

@Amrutha-M05

@Amrutha-M05 Amrutha-M05 commented Jul 29, 2026

Copy link
Copy Markdown

Motivation

HiSparse (hierarchical sparse KV cache) is a CUDA/ROCm-only feature today. This PR enables it on Intel XPU so DSA / DeepSeek-V4-class models can run with hierarchical sparse attention on Battlemage/Crescent Island GPUs.

The upstream HiSparse path hard-gates on is_cuda() or is_hip() in several places and dispatches its host-side pin-memory allocation through a CUDA-only registrar. The XPU-side SYCL JIT kernels for the two hot kernels — load_cache_to_device_buffer_{mla,dsv4_mla} and transfer_cache_dsv4_mla in the companion sgl-kernel-xpu repo under sgl_kernel.jit.kvcache.hisparse, which mirrors the CUDA-side layout sglang.kernels.ops.kvcache.hisparse. What was missing was the sglang-side wiring, backend gates, and device-agnostic tests.

This PR depends on sgl-kernel-xpu PR 335 for load_cache_to_device_buffer_ and transfer_cache_dsv4_mla.

Modifications

Backend gates widened to include XPU:

  • python/sglang/srt/mem_cache/pool_host/common.py — register "xpu": alloc_with_pin_memory in ALLOC_MEMORY_FUNCS so XPU takes torch's built-in pin_memory=True path instead of calling cudaHostRegister. Also guard alloc_with_host_register with not is_xpu() for the mmap-backed custom-allocator path.
  • python/sglang/srt/mem_cache/pool_host/mla.py and memory_pool_host.py — widen if _is_cuda or _is_hip: to ... or _is_xpu: so sgl_kernel.kvcacheio.transfer_kv_all_layer_* is imported on XPU.
  • python/sglang/srt/mem_cache/hisparse_memory_pool.py — same widening for the DSA transfer_kv_all_layer_mla import; updated the fallback error message.

Kernel dispatch:

  • python/sglang/srt/managers/hisparse_coordinator.py — branch the swap-in kernel import at load time:
    if is_xpu():
        from sgl_kernel.jit.kvcache.hisparse import load_cache_to_device_buffer_{mla,dsv4_mla}
    else:
        from sglang.kernels.ops.kvcache.hisparse import ...
    Also reordered imports so is_xpu is imported before the conditional.

Tests made device-agnostic:

  • test/registered/kernels/ops/kvcache/test_hisparse.py (JIT kernel tests) and test/registered/unit/managers/test_hisparse_unit.py — use get_device() / get_device_module() from sglang.srt.utils; import the XPU kernels from sgl_kernel.jit.kvcache.hisparse on XPU; registered under the XPU CI suite stage-b-test-1-gpu-xpu.

Accuracy Tests

Existing HiSparse correctness tests pass unchanged on XPU (Intel B60), and CUDA behavior is unaffected (the changes are additive or is_xpu() gates plus an XPU import branch).

$ pytest test/registered/kernels/ops/kvcache/test_hisparse.py
9 passed, 1 skipped in 25.06s      # 1 skipped = ROCm-only test

$ pytest test/registered/unit/managers/test_hisparse_unit.py
9 passed, 2 skipped in 43.15s      # 2 skipped = page_size=1 ROCm-only paths

The unit tests include the kernel-vs-naive_load_topk oracle comparison (fast-path, long-seq host-DMA, LRU replacement across decode steps, and batched multi-request), so kernel numerical parity between the SYCL and CUDA implementations is exercised end-to-end on the same test bodies.

Speed Tests and Profiling

Not applicable — this PR is enablement only; it does not change the hot path on CUDA/ROCm. The XPU path uses the same swap-in and evict/backup shapes as CUDA and benchmarking of the SYCL kernels themselves is tracked in the companion sgl-kernel-xpu repo.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

cc: @siju-samuel @rbabukv


CI States

Latest PR Test (Base): ❌ Run #32102299980
Latest PR Test (Extra): ❌ Run #32102299772

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Wires the DSA/DSv4 HiSparse path onto XPU by widening backend guards,
switching pin-memory allocation to torch's built-in pinned allocator, and
routing the XPU coordinator import at the XPU-side SYCL JIT namespace
(sgl_kernel.jit.kvcache.hisparse) that mirrors sglang.kernels.ops.kvcache.hisparse.

- pool_host/{common,mla}.py, memory_pool_host.py, hisparse_memory_pool.py:
  extend cuda/hip backend gates to also cover xpu; register "xpu" in
  ALLOC_MEMORY_FUNCS so mmap-backed allocations skip cudaHostRegister.
- managers/hisparse_coordinator.py: import is_xpu before the conditional
  kernel import; branch to sgl_kernel.jit.kvcache.hisparse on XPU.
- test/registered/jit/test_hisparse.py,
  test/registered/unit/managers/test_hisparse_unit.py: make device selection
  device-agnostic via get_device()/get_device_module() and register XPU CI.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant