HiSparse swap-in kernels (transfer_cache_dsv4_mla, load_cache_to_device_buffer)

[Amrutha-M05]

Summary

Port the HiSparse KV-offload swap-in kernels from CUDA to Intel XPU as AOT-compiled SYCL, used by DeepSeek DSA / V4 hierarchical sparse attention:

• transfer_cache_dsv4_mla — bulk-copy DSv4 C4 tokens between page-padded C4 buffers, one set per model layer (evict / backup path).
• load_cache_to_device_buffer_mla / load_cache_to_device_buffer_dsv4_mla — per-request swap-in of the current top-k tokens into a small hot device buffer, maintaining LRU ordering and streaming misses in from the host cache.

The SYCL ports mirror the CUDA sources in sglang jit_kernel/csrc/hisparse.cuh and reuse the same page-padded C4 layout (kValueBytes=576, kScaleBytes=8, kPageSize=64), so outputs match the CUDA reference bit-for-bit.

Both ops are built into common_ops and registered on torch::kXPU through TORCH_LIBRARY_FRAGMENT, so they are reachable as torch.ops.sgl_kernel.* with no compiler on the runtime host and no first-call compile latency.

Changes

• src/sycl/HiSparse.cpp — new translation unit holding both SGL_KERNEL_EXPORT entry points: dtype/contiguity/device validation, shared-memory layout computation, and the queue submissions. Picked up automatically by the file(GLOB ...) in src/CMakeLists.txt.
• include/sgl_kernel_ops.h — declare the two ops.
• src/torch_extension_sycl.cc — m.def schemas plus m.impl(..., torch::kXPU, ...) bindings.
• include/sgl_kernel/hisparse/ (moved out of jit_kernel/)
  • c4_layout.hpp — shared paged-C4 pointer resolution, layout constants, and cooperative sub-group transfer_item.
  • transfer_cache_dsv4_mla.hpp — bulk transfer, one sub-group per item, grid-stride over items, all layers per item. Templated on BLOCK_SIZE; the host side dispatches to 256 / 512 / 1024 from the runtime block_size argument (default 1024).
  • load_cache_to_device_buffer.hpp — per-request LRU swap-in: hashes top-k, compacts hits/evictables, assigns misses to evicted slots, rewrites LRU order, then each sub-group copies one miss host→device.
• python/sgl_kernel/hisparse.py — thin torch.ops wrappers for the three entry points; API matches sglang.jit_kernel.hisparse.
• python/sgl_kernel/__init__.py — export the three entry points.
• tests/test_hisparse.py, tests/run_suite.py — tests plus registration in the default suite so they run in CI.
• benchmark/bench_hisparse.py — all three entry points vs a pure-PyTorch implementation of the same op (addresses the review comment asking for benchmark tests).
• cmake/BuildFlags.cmake — add -std=c++20 to SYCL_KERNEL_OPTIONS. The device-compiled sources include ATen/torch headers, which hard-#error below C++20; -sycl-std sets the SYCL spec version, not the C++ language standard, so it has to be set explicitly to match SYCL_HOST_FLAGS.

CUDA → SYCL mapping used

CUDA	SYCL
warp (32 lanes)	sub-group pinned to kSubGroupSize = 32 via [[sycl::reqd_sub_group_size(32)]]
__ballot + popc(&before)	exclusive_scan_over_group
popc(mask)	reduce_over_group
__shfl_up_sync / __shfl_sync	inclusive_scan_over_group / group_broadcast
atomicCAS on shared memory	sycl::atomic_ref<..., memory_scope::work_group, address_space::local_space>
extern __shared__	sycl::local_accessor<char, 1>
__syncthreads()	item.barrier(access::fence_space::local_space)

Test plan

• pytest tests/test_hisparse.py -v on an Intel XPU host — 11 passed in 1.60s. Registered in tests/run_suite.py so CI picks it up.
• Coverage: LRU hit/evict compaction and MRU/LRU write-back, miss classification + evict-slot reuse, seq_len <= hot_buffer_size fast path, CUDA-graph padding (num_real_reqs), DSv4 page-padded addressing on both paths, seq_lens / req_pool_indices in both i32 and i64.
• Both ops resolve as torch.ops.sgl_kernel.transfer_cache_dsv4_mla / load_cache_to_device_buffer_mla from a plain import sgl_kernel, with no compiler on the host.
• src/sycl/HiSparse.cpp passes icpx -fsycl -fsyntax-only -std=c++20, and both IsDsv4Layout instantiations are exercised by the tests.
• No .wait() in either launcher — kernels stay async on the current XPU stream.
• benchmark/bench_hisparse.py completes and reports both kernels vs torch eager.
• Benchmark eager references cross-checked against the kernels: transfer_cache_dsv4_mla output byte-identical (torch.equal), and both swap-in paths leave every requested top-k token resident.

CC: @siju-samuel @rbabukv

[Copilot]

Pull request overview

Adds Intel XPU (SYCL) JIT-compiled ports of the HiSparse KV-offload “swap-in” kernels used by hierarchical sparse attention, wiring them into the sgl_kernel.jit Python API and providing accuracy tests to validate CUDA-equivalent behavior.

Changes:

• Introduces SYCL JIT kernel headers for DSv4 paged-C4 transfers and MLA swap-in with LRU maintenance under include/sgl_kernel/jit_kernel/hisparse/.
• Adds Python JIT loaders + ctypes bindings exposing transfer_cache_dsv4_mla and load_cache_to_device_buffer_{mla,dsv4_mla} under python/sgl_kernel/jit/, and exports them from python/sgl_kernel/jit/__init__.py when is_xpu().
• Adds a dedicated pytest suite validating paging correctness, LRU ordering, miss/evict behavior, and padding semantics; updates SYCL kernel compile flags to explicitly set -std=c++20.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tests/test_hisparse_jit.py	New XPU-only accuracy tests covering DSv4 paging, swap-in miss/hit behavior, and LRU updates.
python/sgl_kernel/jit/hisparse.py	New JIT module compilation + ctypes wrappers for the HiSparse kernels.
python/sgl_kernel/jit/__init__.py	Exports the new HiSparse JIT entry points under the XPU path.
include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp	Shared DSv4 paged-C4 pointer resolution and cooperative item copy helper.
include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp	SYCL port of the DSv4 bulk transfer kernel with C-exported entry points.
include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp	SYCL port of the per-request swap-in kernel (hashing, hit/evict compaction, miss copy, LRU rewrite).
cmake/BuildFlags.cmake	Adds -std=c++20 to SYCL device compilation options to match host C++ standard expectations.

[Amrutha-M05]

@kareemshaik80 Could you please review this PR .

[kareemshaik80]

please add benchmark tests.

[airMeng]

There will be another PR on SGLang side?

[Amrutha-M05]

There will be another PR on SGLang side?

Hi @airMeng sgl-project/sglang#32792 is the corresponding PR on SGLang

[airMeng]

As I comment in another JIT PR, I'd prefer to copy all dependency files here to avoid circular dependency. Morever, to avoid SGLang side changes to break us.

Like the cpp files which will include include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp

[kareemshaik80]

Add only kernel UT and avoid srt imports here.

[kareemshaik80]

avoid srt imports here.