[XPU]Enable HiSparse hierarchical sparse KV cache on Intel XPU - #32792
Open
Amrutha-M05 wants to merge 2 commits into
Open
[XPU]Enable HiSparse hierarchical sparse KV cache on Intel XPU#32792Amrutha-M05 wants to merge 2 commits into
Amrutha-M05 wants to merge 2 commits into
Conversation
Amrutha-M05
requested review from
Ying1123,
alphabetc1,
hanming-lu,
hnyls2002,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
July 29, 2026 15:34
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Amrutha-M05
force-pushed
the
hisparse-xpu
branch
from
July 30, 2026 08:44
734d7d0 to
f18129c
Compare
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.
Amrutha-M05
force-pushed
the
hisparse-xpu
branch
from
August 2, 2026 14:57
f18129c to
4c99140
Compare
Draft
5 tasks
7 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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}andtransfer_cache_dsv4_mlain the companionsgl-kernel-xpurepo undersgl_kernel.jit.kvcache.hisparse, which mirrors the CUDA-side layoutsglang.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_andtransfer_cache_dsv4_mla.Modifications
Backend gates widened to include XPU:
python/sglang/srt/mem_cache/pool_host/common.py— register"xpu": alloc_with_pin_memoryinALLOC_MEMORY_FUNCSso XPU takes torch's built-inpin_memory=Truepath instead of callingcudaHostRegister. Also guardalloc_with_host_registerwithnot is_xpu()for the mmap-backed custom-allocator path.python/sglang/srt/mem_cache/pool_host/mla.pyandmemory_pool_host.py— widenif _is_cuda or _is_hip:to... or _is_xpu:sosgl_kernel.kvcacheio.transfer_kv_all_layer_*is imported on XPU.python/sglang/srt/mem_cache/hisparse_memory_pool.py— same widening for the DSAtransfer_kv_all_layer_mlaimport; updated the fallback error message.Kernel dispatch:
python/sglang/srt/managers/hisparse_coordinator.py— branch the swap-in kernel import at load time:is_xpuis imported before the conditional.Tests made device-agnostic:
test/registered/kernels/ops/kvcache/test_hisparse.py(JIT kernel tests) andtest/registered/unit/managers/test_hisparse_unit.py— useget_device()/get_device_module()fromsglang.srt.utils; import the XPU kernels fromsgl_kernel.jit.kvcache.hisparseon XPU; registered under the XPU CI suitestage-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).The unit tests include the kernel-vs-
naive_load_topkoracle 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-xpurepo.Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-cicc: @siju-samuel @rbabukv
CI States
Latest PR Test (Base): ❌ Run #32102299980
Latest PR Test (Extra): ❌ Run #32102299772