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

[HiSparse & PD] Support hisparse memory pool host page > 1 (#23606)

Browse Source 
Co-authored-by: hzh0425 <hzh0425@apache.org>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
huangtingwei
2026-05-19 01:29:35 -07:00
committed by GitHub
co-authored by hzh0425 Zhiqiang Xie
parent 1f7bf155c3
commit 67fd005b97
7 changed files with 253 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/common/conn.py
+5 -13
View File
@@ -235,19 +235,11 @@ class CommonKVManager(BaseKVManager):
# Sanity checks
if info.page_size is not None and info.page_size != self.kv_args.page_size:
if self.server_args.enable_hisparse:
# HiSparse: decode host pool page_size=1, prefill device pool page_size >= 1.
# Transfer will use send_kvcache_hisparse with per-token item_lens.
logger.info(
f"HiSparse PD transfer mode: prefill page_size={info.page_size}, "
f"decode host page_size={self.kv_args.page_size}"
)
else:
raise RuntimeError(
f"Page size mismatch: prefill server has page_size={info.page_size}, "
f"but decode server has page_size={self.kv_args.page_size}. "
f"Both servers must use the same --page-size value."
)
raise RuntimeError(
f"Page size mismatch: prefill server has page_size={info.page_size}, "
f"but decode server has page_size={self.kv_args.page_size}. "
f"Both servers must use the same --page-size value."
)
if (
info.kv_cache_dtype is not None
python/sglang/srt/disaggregation/decode.py
+18 -27
View File
@@ -380,16 +380,14 @@ class DecodePreallocQueue:
kv_args.pp_rank = self.pp_rank
kv_args.system_dp_rank = self.scheduler.ps.dp_rank
if self.scheduler.enable_hisparse:
# Direct-to-host: register host pool pointers so P writes to D's host memory
host_pool = self.scheduler.hisparse_coordinator.mem_pool_host
kv_data_ptrs, kv_data_lens, kv_item_lens = (
host_pool.get_contiguous_buf_infos()
)
else:
kv_data_ptrs, kv_data_lens, kv_item_lens = (
self.token_to_kv_pool.get_contiguous_buf_infos()
)
transfer_kv_pool = (
self.scheduler.hisparse_coordinator.mem_pool_host
if self.scheduler.enable_hisparse
else self.token_to_kv_pool
)
kv_data_ptrs, kv_data_lens, kv_item_lens = (
transfer_kv_pool.get_contiguous_buf_infos()
)
if self.draft_token_to_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are
# always shared with a target model.
@@ -403,10 +401,7 @@ class DecodePreallocQueue:
kv_args.kv_data_ptrs = kv_data_ptrs
kv_args.kv_data_lens = kv_data_lens
kv_args.kv_item_lens = kv_item_lens
# HiSparse Host pool has page_size=1; use it when hisparse is enabled
kv_args.page_size = (
1 if self.scheduler.enable_hisparse else self.token_to_kv_pool.page_size
)
kv_args.page_size = self.token_to_kv_pool.page_size
kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
self.metadata_buffers.get_buf_infos()
@@ -913,6 +908,7 @@ class DecodePreallocQueue:
swa_allocatable_tokens -= swa_required
decode_req.req.cache_protected_len = prefix_len
page_size = self.token_to_kv_pool_allocator.page_size
if self.scheduler.enable_hisparse:
# Must cast to int32 for ZMQ serialization -- from_zmq reads np.int32.
kv_indices = (
@@ -921,7 +917,6 @@ class DecodePreallocQueue:
.numpy()
.astype(np.int32)
)
page_size = 1 # host pool page_size
else:
# Only send delta indices (beyond prefix) to prefill.
kv_indices = (
@@ -931,7 +926,6 @@ class DecodePreallocQueue:
.cpu()
.numpy()
)
page_size = self.token_to_kv_pool_allocator.page_size
seq_len = len(decode_req.req.origin_input_ids)
@@ -1270,14 +1264,13 @@ class DecodePreallocQueue:
extend_num_tokens=fill_len,
)
# Allocate host indices for the RDMA transfer target.
host_indices = coordinator.mem_pool_host.alloc(fill_len)
if host_indices is None:
raise RuntimeError(
f"HiSparse host mem pool alloc failed for {fill_len} tokens "
f"in _pre_alloc (req {req.rid})"
)
host_indices = host_indices.to(device=coordinator.device)
coordinator.req_to_host_pool[req.req_pool_idx, :fill_len] = host_indices
host_indices = coordinator.mem_pool_host.alloc_paged_token_slots(
coordinator.req_to_host_pool,
coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx,
0,
fill_len,
)
elif self.token_to_kv_pool_allocator.page_size == 1:
kv_loc = self.token_to_kv_pool_allocator.alloc(delta_len)
else:
@@ -1793,6 +1786,4 @@ class SchedulerDisaggregationDecodeMixin:
for req in transferred_reqs:
# Direct-to-host: KV data already in host pool, skip staging
self.hisparse_coordinator.admit_request_direct(req)
self.waiting_queue.extend(transferred_reqs)
else:
self.waiting_queue.extend(transferred_reqs)
self.waiting_queue.extend(transferred_reqs)
python/sglang/srt/disaggregation/mooncake/conn.py
+8 -69
View File
@@ -129,8 +129,6 @@ class KVArgsRegisterInfo:
# for mamba state different tp slice transfer
dst_state_item_lens: List[List[int]]
dst_state_dim_per_tensor: List[List[int]]
# HiSparse: decode host pool stores KV at token granularity
enable_hisparse: bool = False
# Note: always put the staging field at the final (since the staging field is optional and contains multiple inputs)
staging: Optional[StagingRegisterInfo] = None
@@ -153,11 +151,8 @@ class KVArgsRegisterInfo:
dst_state_dim_per_tensor=(
unpack_int_lists(msg[11], "I") if len(msg) > 11 else []
),
enable_hisparse=(
msg[12].decode("ascii") == "1" if len(msg) > 12 else False
),
# Note: always put the staging field at the final
staging=StagingRegisterInfo.from_zmq_fields(msg, 13),
staging=StagingRegisterInfo.from_zmq_fields(msg, 12),
)
@@ -704,49 +699,6 @@ class MooncakeKVManager(CommonKVManager):
executor=executor,
)
def send_kvcache_hisparse(
self,
mooncake_session_id: str,
prefill_kv_indices: npt.NDArray[np.int32],
dst_kv_ptrs: list[int],
dst_kv_indices: npt.NDArray[np.int32],
page_index_slice: slice,
executor: concurrent.futures.ThreadPoolExecutor,
):
"""HiSparse transfer: prefill page_size > decode host page_size=1.
Receives page-level prefill_kv_indices and the full token-level
dst_kv_indices. Expands both to token granularity before transfer.
"""
page_size = self.kv_args.page_size
per_token_item_lens = [il // page_size for il in self.kv_args.kv_item_lens]
# Expand page-level src indices to token-level
base = np.repeat(prefill_kv_indices * page_size, page_size)
offsets = np.tile(np.arange(page_size, dtype=np.int32), len(prefill_kv_indices))
expanded_src = base + offsets
# Expand page-level index_slice to token-level for dst
token_start = page_index_slice.start * page_size
token_end = min(page_index_slice.stop * page_size, len(dst_kv_indices))
expanded_dst = dst_kv_indices[token_start:token_end]
# Clip src to match dst length (last page may be partial)
expanded_src = expanded_src[: len(expanded_dst)]
logger.debug(
f"Send KVCache for hisparse: {expanded_src.shape} -> {expanded_dst.shape}"
)
return self._send_kvcache_generic(
mooncake_session_id=mooncake_session_id,
src_data_ptrs=self.kv_args.kv_data_ptrs,
dst_data_ptrs=dst_kv_ptrs,
item_lens=per_token_item_lens,
prefill_data_indices=expanded_src,
dst_data_indices=expanded_dst,
executor=executor,
)
def send_kvcache_slice(
self,
mooncake_session_id: str,
@@ -1269,23 +1221,13 @@ class MooncakeKVManager(CommonKVManager):
self.attn_tp_size
== target_rank_registration_info.dst_attn_tp_size
):
if target_rank_registration_info.enable_hisparse:
ret = self.send_kvcache_hisparse(
req.mooncake_session_id,
kv_chunk.prefill_kv_indices,
target_rank_registration_info.dst_kv_ptrs,
req.dst_kv_indices,
kv_chunk.index_slice,
executor,
)
else:
ret = self.send_kvcache(
req.mooncake_session_id,
kv_chunk.prefill_kv_indices,
target_rank_registration_info.dst_kv_ptrs,
chunked_dst_kv_indice,
executor,
)
ret = self.send_kvcache(
req.mooncake_session_id,
kv_chunk.prefill_kv_indices,
target_rank_registration_info.dst_kv_ptrs,
chunked_dst_kv_indice,
executor,
)
elif (
self.enable_staging
and staging_strategy is not None
@@ -1789,8 +1731,6 @@ class MooncakeKVReceiver(CommonKVReceiver):
dst_tp_rank = str(tp_rank).encode("ascii")
dst_attn_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii")
dst_kv_item_len = str(kv_item_len).encode("ascii")
enable_hisparse = b"1" if self.kv_mgr.server_args.enable_hisparse else b"0"
if (
self.kv_mgr.enable_staging
and self.kv_mgr._staging_ctx.allocator is not None
@@ -1818,7 +1758,6 @@ class MooncakeKVReceiver(CommonKVReceiver):
dst_kv_item_len,
packed_state_item_lens,
packed_state_dim_per_tensor,
enable_hisparse,
packed_staging_base_ptr,
staging_total_size_str,
]
python/sglang/srt/managers/hisparse_coordinator.py
+40 -31
View File
@@ -67,7 +67,9 @@ class HiSparseCoordinator:
self.mem_pool_device = self.token_to_kv_pool_allocator.hisparse_kvcache
host_size = self.token_to_kv_pool_allocator.size_full // self.compress_ratio
self.mem_pool_host = DeepSeekV4SingleKVPoolHost(
self.mem_pool_device, host_size, 1
self.mem_pool_device,
host_size,
page_size=self.mem_pool_device.page_size,
)
self.item_size_bytes = (
self.mem_pool_host.kv_cache_total_dim
@@ -84,11 +86,12 @@ class HiSparseCoordinator:
device_pool=self.mem_pool_device,
host_to_device_ratio=host_to_device_ratio,
host_size=0,
page_size=1,
page_size=self.mem_pool_device.page_size,
layout="layer_first",
override_kv_cache_dim=self.mem_pool_device.kv_cache_dim,
)
self.item_size_bytes = self.mem_pool_host.token_stride_size
self.page_size = self.mem_pool_device.page_size
max_num_req_slots = req_to_token_pool.req_to_token.shape[0]
max_context_len = req_to_token_pool.max_context_len
@@ -110,11 +113,14 @@ class HiSparseCoordinator:
max_num_req_slots, dtype=torch.int64, device="cpu"
)
self.req_to_host_pool = torch.full(
(max_num_req_slots, max_compressed_context_len),
(max_num_req_slots, max_compressed_context_len + self.page_size),
-1,
dtype=torch.int64,
device=device,
)
self.req_to_host_pool_allocated_len = torch.zeros(
max_num_req_slots, dtype=torch.int64, device="cpu"
)
self.write_staging_stream = device_module.Stream()
self.decode_backup_stream = device_module.Stream()
@@ -200,18 +206,13 @@ class HiSparseCoordinator:
)
prefill_len = len(device_indices)
host_indices = self.mem_pool_host.alloc(prefill_len)
if host_indices is None:
logger.error(
"HiSparse: host mem pool alloc failed for %d tokens (req %s)",
prefill_len,
req.rid,
)
raise RuntimeError(
f"HiSparse host mem pool alloc failed for {prefill_len} tokens"
)
host_indices = host_indices.to(device=self.device)
self.req_to_host_pool[req.req_pool_idx, :prefill_len] = host_indices
host_indices = self.mem_pool_host.alloc_paged_token_slots(
self.req_to_host_pool,
self.req_to_host_pool_allocated_len,
req.req_pool_idx,
0,
prefill_len,
)
start_event = device_module.Event()
finish_event = device_module.Event()
@@ -549,17 +550,19 @@ class HiSparseCoordinator:
device_locs = self.req_to_device_buffer[backup_req_indices, buffer_slot]
host_locs = self.mem_pool_host.alloc(len(device_locs))
if host_locs is None:
logger.error(
"HiSparse: host mem pool alloc failed for %d decode backup tokens",
len(device_locs),
host_locs_list = []
for i in backup_indices:
req_idx = int(req_pool_indices_cpu[i])
start_pos = (int(seq_lens_cpu[i]) - 1) // self.compress_ratio - 1
host_locs = self.mem_pool_host.alloc_paged_token_slots(
self.req_to_host_pool,
self.req_to_host_pool_allocated_len,
req_idx,
start_pos,
1,
)
raise RuntimeError(
f"HiSparse host mem pool alloc failed for {len(device_locs)} decode backup tokens"
)
host_locs = host_locs.to(device=self.device)
self.req_to_host_pool[backup_req_indices, actual_compressed_pos] = host_locs
host_locs_list.append(host_locs)
host_locs = torch.cat(host_locs_list)
self.wait_for_pending_backup()
schedule_stream = device_module.current_stream()
@@ -702,12 +705,15 @@ class HiSparseCoordinator:
self.token_to_kv_pool_allocator.free_hisparse(allocated_locs)
# Free host memory that was allocated during admit_request_into_staging
compressed_len = prefill_len // self.compress_ratio
host_indices = self.req_to_host_pool[req.req_pool_idx, :compressed_len]
host_indices = host_indices[host_indices >= 0]
host_indices = self.mem_pool_host.allocated_host_indices(
self.req_to_host_pool,
req.req_pool_idx,
self.req_to_host_pool_allocated_len[req.req_pool_idx],
)
if host_indices.numel() > 0:
self.mem_pool_host.free(host_indices)
self.req_to_host_pool[req.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0
self._skip_first_backup[req.req_pool_idx] = False
req.hisparse_staging = False
@@ -730,7 +736,6 @@ class HiSparseCoordinator:
# subsequent release_kv_cache -> allocator.free -> free_hisparse path
# re-frees them (double-free into the page allocator's free list).
allocated_len = req.kv_allocated_len
compressed_len = allocated_len // self.compress_ratio
# release memory -- only free actually-allocated buffer indices
current_cap = int(self.req_device_buffer_size[req.req_pool_idx])
@@ -748,8 +753,11 @@ class HiSparseCoordinator:
)
self.mem_pool_device.full_to_hisparse_device_index_mapping[compressed_locs] = 0
host_indices = self.req_to_host_pool[req.req_pool_idx, :compressed_len]
host_indices = host_indices[host_indices >= 0]
host_indices = self.mem_pool_host.allocated_host_indices(
self.req_to_host_pool,
req.req_pool_idx,
self.req_to_host_pool_allocated_len[req.req_pool_idx],
)
if host_indices.numel() > 0:
self.mem_pool_host.free(host_indices)
@@ -759,6 +767,7 @@ class HiSparseCoordinator:
self.req_to_device_buffer[req.req_pool_idx, :] = 0
self.req_device_buffer_size[req.req_pool_idx] = 0
self.req_to_host_pool[req.req_pool_idx, :] = -1
self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0
self.lru_slots[:, req.req_pool_idx, :].copy_(self._lru_init)
self._skip_first_backup[req.req_pool_idx] = False
python/sglang/srt/mem_cache/hisparse_memory_pool.py
+4 -3
View File
@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
HiSparseC4DevicePool,
)
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import HiSparseHostPoolMixin
from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils.common import get_num_new_pages
@@ -385,7 +386,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
)
class DeepSeekV4SingleKVPoolHost:
class DeepSeekV4SingleKVPoolHost(HiSparseHostPoolMixin):
def __init__(
self,
@@ -397,12 +398,12 @@ class DeepSeekV4SingleKVPoolHost:
):
assert host_size > 0, "Host size must be specified and greater than 0"
assert page_size == 1, "Host page size must be 1 for DeepSeekV4SingleKVPoolHost"
self.device_pool = device_pool
self.size = host_size
self.page_size = page_size
self.num_pages = (self.size + self.page_size - 1) // self.page_size
self.size = self.num_pages * self.page_size
self.pin_memory = pin_memory
self.device = device
@@ -421,7 +422,7 @@ class DeepSeekV4SingleKVPoolHost:
def clear(self):
self.free_slots = torch.arange(
1, self.num_pages + 1, dtype=torch.int64, device="cpu"
1, self.size + 1, dtype=torch.int64, device="cpu"
)
def init_kv_buffer(self):
python/sglang/srt/mem_cache/memory_pool_host.py
+68 -4
View File
@@ -37,13 +37,14 @@ from sglang.srt.mem_cache.memory_pool import (
MLATokenToKVPool,
NSATokenToKVPool,
)
from sglang.srt.utils import is_cuda, is_mps, is_npu, is_xpu
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
_is_xpu = is_xpu()
_is_mps = is_mps()
if not (_is_npu or _is_xpu or _is_mps):
if _is_cuda or _is_hip:
from sgl_kernel.kvcacheio import (
transfer_kv_all_layer,
transfer_kv_all_layer_direct_lf_pf,
@@ -91,6 +92,69 @@ class HostTensorAllocator(abc.ABC):
return tensor
class HiSparseHostPoolMixin:
def _round_up_to_page_size(self, size: int) -> int:
return (size + self.page_size - 1) // self.page_size * self.page_size
def alloc_page(self, num_pages: int) -> Optional[torch.Tensor]:
return self.alloc(num_pages * self.page_size)
def alloc_paged_token_slots(
self,
req_to_host_pool: torch.Tensor,
req_to_host_pool_allocated_len: torch.Tensor,
req_pool_idx: int,
start_pos: int,
num_tokens: int,
) -> torch.Tensor:
"""Allocate request host slots by page and return token-granular slots."""
device = req_to_host_pool.device
if num_tokens <= 0:
return torch.empty((0,), dtype=torch.int64, device=device)
allocated_len = int(req_to_host_pool_allocated_len[req_pool_idx])
end_pos = start_pos + num_tokens
page_end = self._round_up_to_page_size(end_pos)
assert start_pos <= allocated_len
if page_end > allocated_len:
num_new_pages = (page_end - allocated_len) // self.page_size
host_locs = self.alloc_page(num_new_pages)
if host_locs is None:
logger.error(
"HiSparse: host mem pool alloc failed for %d host pages "
"(req_pool_idx=%d, start_pos=%d, num_tokens=%d)",
num_new_pages,
req_pool_idx,
start_pos,
num_tokens,
)
raise RuntimeError(
f"HiSparse host mem pool alloc failed for {num_new_pages} pages"
)
req_to_host_pool[req_pool_idx, allocated_len:page_end] = host_locs.to(
device=device, non_blocking=True
)
req_to_host_pool_allocated_len[req_pool_idx] = page_end
return req_to_host_pool[req_pool_idx, start_pos:end_pos]
def allocated_host_indices(
self,
req_to_host_pool: torch.Tensor,
req_pool_idx: int,
allocated_len: int,
) -> torch.Tensor:
allocated_len = int(allocated_len)
host_len = min(
self._round_up_to_page_size(allocated_len),
req_to_host_pool.shape[1],
)
host_indices = req_to_host_pool[req_pool_idx, :host_len]
return host_indices[host_indices >= 0]
def get_allocator_from_storage(allocator_type):
if allocator_type == "mooncake":
try:
@@ -785,7 +849,7 @@ class MHATokenToKVPoolHost(HostKVCache):
return ptr_list, element_size_list
class MLATokenToKVPoolHost(HostKVCache):
class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
device_pool: MLATokenToKVPool
def __init__(
@@ -833,7 +897,7 @@ class MLATokenToKVPoolHost(HostKVCache):
for registering host memory with the disaggregation transfer engine."""
data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)]
data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
item_lens = [self.token_stride_size] * self.layer_num
item_lens = [self.token_stride_size * self.page_size] * self.layer_num
return data_ptrs, data_lens, item_lens
def get_size_per_token(self):