minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
AMRUTHA M
2026-09-21 14:07:49 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent fcb080bd40
commit d20cd9d77f
7 changed files with 126 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/hisparse_coordinator.py
+36 -7
View File
@@ -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
+6 -5
View File
@@ -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
+5 -2
View File
@@ -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
+1
View File
@@ -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
+1 -1
View File
@@ -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,
test/registered/kernels/ops/kvcache/test_hisparse.py
+38 -21
View File
@@ -3,26 +3,43 @@ import sys
import pytest
import torch
from sglang.kernels.ops.kvcache.hisparse import (
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
transfer_cache_dsv4_mla,
from sglang.srt.utils import (
get_device,
get_device_module,
is_cuda,
is_hip,
is_xpu,
)
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_amd_ci(est_time=30, stage="stage-b", runner_config="1-gpu-small-amd")
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
if is_xpu():
from sgl_kernel import (
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
transfer_cache_dsv4_mla,
)
else:
from sglang.kernels.ops.kvcache.hisparse import (
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
transfer_cache_dsv4_mla,
)
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available()
or is_npu()
or is_xpu()
or not (is_cuda() or is_hip()),
reason="HiSparse JIT tests require CUDA/ROCm.",
not (is_cuda() or is_hip() or is_xpu()),
reason="HiSparse kernel tests require CUDA/ROCm/XPU.",
)
DEVICE = "cuda"
DEVICE = get_device()
DTYPE = torch.float32
KV_DIM = 8
HOT_BUFFER_SIZE = 4
@@ -131,7 +148,7 @@ def _run_kernel(
block_size=256,
num_real_reqs=torch.tensor([num_real_reqs], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
get_device_module().synchronize()
return out
@@ -171,7 +188,7 @@ def _make_state(
device_buffer[device_buffer_locs[rid, HOT_BUFFER_SIZE]].copy_(
host_cache[newest_token].to(DEVICE, non_blocking=True)
)
torch.cuda.synchronize()
get_device_module().synchronize()
return {
"host_cache": host_cache,
@@ -199,7 +216,7 @@ def test_transfer_cache_dsv4_mla_copies_paged_token() -> None:
src_indices=torch.tensor([src_loc], dtype=torch.int64, device=DEVICE),
dst_indices=torch.tensor([dst_loc], dtype=torch.int64, device=DEVICE),
)
torch.cuda.synchronize()
get_device_module().synchronize()
assert torch.equal(
_read_dsv4_token(dst_cache, dst_loc).to(DEVICE),
@@ -251,7 +268,7 @@ def test_dsv4_swap_in_reads_paged_host_layout() -> None:
block_size=256,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
get_device_module().synchronize()
assert out.item() == swap_loc
assert torch.equal(
@@ -427,7 +444,7 @@ def test_load_cache_to_device_buffer_miss_copy_is_byte_exact(
)
for slot in range(HOT_BUFFER_SIZE):
device_buffer[slot].copy_(host_cache[slot].to(DEVICE))
torch.cuda.synchronize()
get_device_module().synchronize()
top_k_tokens = torch.tensor([[miss_token]], dtype=torch.int32, device=DEVICE)
out = torch.full_like(top_k_tokens, -1)
@@ -454,7 +471,7 @@ def test_load_cache_to_device_buffer_miss_copy_is_byte_exact(
block_size=256,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
get_device_module().synchronize()
# The miss evicts the LRU head (slot 0, physical loc 0) and lands there.
assert torch.equal(out.cpu(), torch.tensor([[0]], dtype=torch.int32))
@@ -593,7 +610,7 @@ def test_load_cache_to_device_buffer_dsv4_mla_miss_copy_layout() -> None:
block_size=256,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
get_device_module().synchronize()
assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32))
@@ -666,7 +683,7 @@ def test_load_cache_to_device_buffer_dsv4_fused_copy_multi_miss() -> None:
block_size=256,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
get_device_module().synchronize()
# Which slot each miss evicts is up to the LRU, so take the destinations
# from the kernel; only require that they are distinct and in range.
@@ -741,7 +758,7 @@ def test_load_cache_to_device_buffer_rocm_large_lru_writeback() -> None:
block_size=1024,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
get_device_module().synchronize()
expected_lru = torch.cat(
[
test/registered/unit/managers/test_hisparse_unit.py
+39 -29
View File
@@ -17,12 +17,26 @@ import torch
from sglang.srt.managers.schedule_batch import ReqKvInfo
from sglang.srt.runtime_context import publish, reset_context
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.srt.utils import (
get_device,
get_device_module,
is_cuda,
is_hip,
is_xpu,
)
from sglang.srt.utils.common import Range
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cuda_ci,
register_xpu_ci,
)
register_cuda_ci(est_time=11, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
DEVICE = get_device()
# ---------------------------------------------------------------------------
# Test configuration (small-scale for fast CI runs)
@@ -74,12 +88,8 @@ class TestHiSparseUnit(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is required for HiSparse tests.")
if is_npu() or is_xpu():
raise unittest.SkipTest("HiSparse tests only support CUDA/ROCm.")
if not (is_cuda() or is_hip()):
raise unittest.SkipTest("CUDA/ROCm not available.")
if not (is_cuda() or is_hip() or is_xpu()):
raise unittest.SkipTest("CUDA/ROCm/XPU not available.")
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29599")
@@ -92,8 +102,8 @@ class TestHiSparseUnit(unittest.TestCase):
alloc_with_pin_memory,
)
cls._original_alloc = ALLOC_MEMORY_FUNCS["cuda"]
ALLOC_MEMORY_FUNCS["cuda"] = alloc_with_pin_memory
cls._original_alloc = ALLOC_MEMORY_FUNCS[DEVICE]
ALLOC_MEMORY_FUNCS[DEVICE] = alloc_with_pin_memory
if is_hip():
from sglang.srt.layers.attention.dsa.utils import (
@@ -116,7 +126,7 @@ class TestHiSparseUnit(unittest.TestCase):
dtype=torch.bfloat16,
qk_rope_head_dim=QK_ROPE_HEAD_DIM,
layer_num=LAYER_NUM,
device="cuda",
device=DEVICE,
index_head_dim=128,
enable_memory_saver=False,
kv_cache_dim=KV_CACHE_DIM,
@@ -126,7 +136,7 @@ class TestHiSparseUnit(unittest.TestCase):
size=SIZE,
page_size=global_page_size,
dtype=torch.bfloat16,
device="cuda",
device=DEVICE,
kvcache=cls.device_pool,
need_sort=False,
host_to_device_ratio=HOST_TO_DEVICE_RATIO,
@@ -137,7 +147,7 @@ class TestHiSparseUnit(unittest.TestCase):
cls.req_to_token_pool = ReqToTokenPool(
size=MAX_NUM_REQS,
max_context_len=MAX_CONTEXT_LEN,
device="cuda",
device=DEVICE,
enable_memory_saver=False,
)
@@ -149,7 +159,7 @@ class TestHiSparseUnit(unittest.TestCase):
token_to_kv_pool_allocator=cls.allocator,
top_k=TOP_K,
device_buffer_size=DEVICE_BUFFER_SIZE,
device="cuda",
device=DEVICE,
tp_group=cls.tp_group,
host_to_device_ratio=HOST_TO_DEVICE_RATIO,
)
@@ -158,7 +168,7 @@ class TestHiSparseUnit(unittest.TestCase):
def tearDownClass(cls):
from sglang.srt.mem_cache.pool_host.common import ALLOC_MEMORY_FUNCS
ALLOC_MEMORY_FUNCS["cuda"] = cls._original_alloc
ALLOC_MEMORY_FUNCS[DEVICE] = cls._original_alloc
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
@@ -260,11 +270,11 @@ class TestHiSparseUnit(unittest.TestCase):
def _populate_host_pool(self, req, fill_len):
"""Allocate host slots, write known patterns, register in coordinator.
Returns host_indices (cuda tensor)."""
Returns host_indices (device tensor)."""
host_pool = self.coordinator.mem_pool_host
host_indices = host_pool.alloc(fill_len)
self.assertIsNotNone(host_indices, "Host alloc failed")
host_indices = host_indices.to(device="cuda")
host_indices = host_indices.to(device=DEVICE)
self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :fill_len] = host_indices
self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = fill_len
for lid in range(LAYER_NUM):
@@ -273,7 +283,7 @@ class TestHiSparseUnit(unittest.TestCase):
return host_indices
def _build_topk_tokens(self, fill_len, *, include_newest=False):
"""Build a 1-D [TOP_K] int32 cuda tensor of token positions.
"""Build a 1-D [TOP_K] int32 device tensor of token positions.
If include_newest=True, fill_len-1 is guaranteed as the last valid slot.
Pads with -1 when fill_len (or fill_len-1) < TOP_K.
@@ -286,25 +296,25 @@ class TestHiSparseUnit(unittest.TestCase):
"""
n = min(fill_len, TOP_K)
if include_newest and n > 1:
tokens = torch.randperm(fill_len - 1, device="cuda")[: n - 1].to(
tokens = torch.randperm(fill_len - 1, device=DEVICE)[: n - 1].to(
torch.int32
)
tokens = torch.cat(
[tokens, torch.tensor([fill_len - 1], dtype=torch.int32, device="cuda")]
[tokens, torch.tensor([fill_len - 1], dtype=torch.int32, device=DEVICE)]
)
else:
tokens = torch.randperm(fill_len, device="cuda")[:n].to(torch.int32)
tokens = torch.randperm(fill_len, device=DEVICE)[:n].to(torch.int32)
if n < TOP_K:
pad = torch.full((TOP_K - n,), -1, dtype=torch.int32, device="cuda")
pad = torch.full((TOP_K - n,), -1, dtype=torch.int32, device=DEVICE)
tokens = torch.cat([tokens, pad])
return tokens
def _make_batch_tensors(self, reqs, fill_lens):
"""Build (req_pool_indices [int64], seq_lens [int32]) on cuda."""
"""Build (req_pool_indices [int64], seq_lens [int32]) on the active device."""
rpi = torch.tensor(
[r.kv.req_pool_idx for r in reqs], dtype=torch.int64, device="cuda"
[r.kv.req_pool_idx for r in reqs], dtype=torch.int64, device=DEVICE
)
sls = torch.tensor(fill_lens, dtype=torch.int32, device="cuda")
sls = torch.tensor(fill_lens, dtype=torch.int32, device=DEVICE)
return rpi, sls
def _assert_kv_correct(self, locs_row, tokens_row, layer_id, count, msg=""):
@@ -467,7 +477,7 @@ class TestHiSparseUnit(unittest.TestCase):
# Step 1: load the first TOP_K positions from host (no newest token —
# the reserved slot is only valid after map_last_loc_to_buffer which is
# called during an actual decode step, not modelled here).
tokens_s1 = torch.arange(TOP_K, dtype=torch.int32, device="cuda")
tokens_s1 = torch.arange(TOP_K, dtype=torch.int32, device=DEVICE)
locs1 = self._swap_in_selected_pages(
rpi, sls, tokens_s1.unsqueeze(0), layer_id=0
)
@@ -481,7 +491,7 @@ class TestHiSparseUnit(unittest.TestCase):
[
tokens_s1[:half], # hits
torch.arange(
new_start, new_start + half, dtype=torch.int32, device="cuda"
new_start, new_start + half, dtype=torch.int32, device=DEVICE
), # misses
]
)
@@ -633,7 +643,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.coordinator.admit_request_into_staging(req)
self.assertTrue(req.hisparse_staging)
torch.cuda.synchronize()
get_device_module().synchronize()
ready = self.coordinator.collect_ready_reqs()
self.assertEqual(len(ready), 1)
self.assertFalse(req.hisparse_staging)
@@ -669,7 +679,7 @@ class TestHiSparseUnit(unittest.TestCase):
self._write_device_patterns(kv_loc, fill_len)
self.coordinator.admit_request_into_staging(req)
torch.cuda.synchronize()
get_device_module().synchronize()
ready = self.coordinator.collect_ready_reqs()
self.assertEqual(ready, [req])