diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 555ef5215..693e7fffe 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -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 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index c6c8020a1..ee975a9b7 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 634f2eae5..39c42d4e9 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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, ] diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 97c87f926..3f138f3ac 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -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 diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index 15b315837..2192381fd 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -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): diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index b4624ab54..805b137a3 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -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): diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index a8538c42c..56fc32a16 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -55,6 +55,9 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None): inflight_middle_chunks=0, ) req.finished = lambda: req.finished_reason is not None + req.set_extend_input_len = lambda extend_input_len: setattr( + req, "extend_input_len", extend_input_len + ) return req @@ -161,6 +164,7 @@ class TestHiSparseUnit(unittest.TestCase): self.coordinator.req_to_device_buffer.zero_() self.coordinator.req_device_buffer_size.zero_() self.coordinator.req_to_host_pool.fill_(-1) + self.coordinator.req_to_host_pool_allocated_len.zero_() self.coordinator.req_device_buffer_tokens.fill_(-1) self.coordinator.req_device_buffer_token_locs.fill_(-1) self.coordinator.lru_slots[:] = self.coordinator._lru_init.view(1, 1, -1) @@ -241,6 +245,7 @@ class TestHiSparseUnit(unittest.TestCase): self.assertIsNotNone(host_indices, "Host alloc failed") host_indices = host_indices.to(device="cuda") self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len] = host_indices + self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx] = fill_len for lid in range(LAYER_NUM): for i in range(fill_len): host_pool.kv_buffer[lid][host_indices[i]] = self._kv_pattern(lid, i) @@ -554,6 +559,56 @@ class TestHiSparseUnit(unittest.TestCase): self._cleanup_req(req, kv_loc) self._assert_sizes_restored(initial, "staging_path") + # ================================================================== + # Test: Single-node staging host page allocation + # ================================================================== + def test_single_node_staging_allocates_paged_host_slots(self): + """Single-node staging should allocate host slots at page granularity.""" + initial = self._get_initial_sizes() + fill_len = self.page_size * 2 + 1 + rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size + req = _make_req("single-node-staging-pages", list(range(fill_len))) + self._alloc_req_slot(req) + + kv_loc = self._alloc_kv(req, fill_len) + self._write_device_patterns(kv_loc, fill_len) + + self.coordinator.admit_request_into_staging(req) + torch.cuda.synchronize() + ready = self.coordinator.collect_ready_reqs() + self.assertEqual(ready, [req]) + + host_row = self.coordinator.req_to_host_pool[req.req_pool_idx, :rounded_len] + self.assertTrue(torch.all(host_row >= 0)) + self.assertEqual(torch.unique(host_row).numel(), rounded_len) + self.assertEqual( + int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + rounded_len, + ) + + available_size = self.coordinator.mem_pool_host.available_size() + next_host_index = self.coordinator.mem_pool_host.alloc_paged_token_slots( + self.coordinator.req_to_host_pool, + self.coordinator.req_to_host_pool_allocated_len, + req.req_pool_idx, + fill_len, + 1, + ) + self.assertEqual( + self.coordinator.mem_pool_host.available_size(), available_size + ) + self.assertTrue(torch.all(next_host_index >= 0)) + + allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices( + self.coordinator.req_to_host_pool, + req.req_pool_idx, + int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + ) + self.assertEqual(allocated_host_indices.numel(), rounded_len) + + self._cleanup_req(req, kv_loc) + self._assert_sizes_restored(initial, "single_node_staging_pages") + # ================================================================== # Test: Direct-to-host (PD separated) path # ================================================================== @@ -589,6 +644,61 @@ class TestHiSparseUnit(unittest.TestCase): self._cleanup_req(req, kv_loc, logical_only=True) self._assert_sizes_restored(initial, "direct_path") + # ================================================================== + # Test: PD decode prealloc host page allocation + # ================================================================== + def test_pd_decode_prealloc_hisparse_host_slots(self): + """PD decode prealloc should allocate RDMA targets through the host pool.""" + initial = self._get_initial_sizes() + fill_len = self.page_size * 2 + 1 + req = _make_req("pd-decode-prealloc", list(range(fill_len))) + + from sglang.srt.disaggregation.decode import DecodePreallocQueue + + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.req_to_token_pool = self.req_to_token_pool + queue.token_to_kv_pool_allocator = self.allocator + queue.tree_cache = SimpleNamespace( + evictable_size=lambda: 0, + protected_size=lambda: 0, + ) + queue.scheduler = SimpleNamespace( + enable_hisparse=True, + hisparse_coordinator=self.coordinator, + server_args=SimpleNamespace(disaggregation_decode_enable_radix_cache=False), + ) + + host_indices = queue._pre_alloc(req) + self.assertEqual(host_indices.numel(), fill_len) + self.assertTrue(torch.all(host_indices >= 0)) + self.assertTrue( + torch.equal( + host_indices, + self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len], + ) + ) + self.assertEqual(req.kv_allocated_len, fill_len) + self.assertEqual(req.kv_committed_len, fill_len) + self.assertEqual(req.extend_input_len, fill_len) + + rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size + self.assertEqual( + int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + rounded_len, + ) + allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices( + self.coordinator.req_to_host_pool, + req.req_pool_idx, + int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + ) + self.assertEqual(allocated_host_indices.numel(), rounded_len) + + kv_loc = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : req.kv_allocated_len + ].clone() + self._cleanup_req(req, kv_loc, logical_only=True) + self._assert_sizes_restored(initial, "pd_decode_prealloc_hisparse") + # ================================================================== # Test: Batch multiple requests # ==================================================================