diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 1d4f48b13280..eb3a896ef34a 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -325,6 +325,7 @@ def build_deepseek_v4_hicache_stack( item_bytes=kvcache.swa_kv_pool.bytes_per_page_padded, num_host_pages=swa_num_host_pages, slot_page_size=kvcache.swa_page_size, + layout=server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) swa_attn_allocator = params.token_to_kv_pool_allocator.swa_attn_allocator @@ -357,6 +358,7 @@ def build_deepseek_v4_hicache_stack( item_bytes=kvcache.c4_kv_pool.bytes_per_page_padded, num_host_pages=num_host_pages, slot_page_size=page_size, + layout=server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) c4_indexer_host_pool = DeepSeekV4PagedHostPool( @@ -368,6 +370,7 @@ def build_deepseek_v4_hicache_stack( ), num_host_pages=num_host_pages, slot_page_size=page_size, + layout=server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) c4_state_host_pool = DeepSeekV4StateHostPool( @@ -378,6 +381,7 @@ def build_deepseek_v4_hicache_stack( ], num_host_pages=swa_num_host_pages, swa_page_size=kvcache.swa_page_size, + layout=server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) c4_indexer_state_host_pool = DeepSeekV4StateHostPool( @@ -388,6 +392,7 @@ def build_deepseek_v4_hicache_stack( ], num_host_pages=swa_num_host_pages, swa_page_size=kvcache.swa_page_size, + layout=server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) entries.extend( @@ -430,6 +435,7 @@ def build_deepseek_v4_hicache_stack( item_bytes=kvcache.c128_kv_pool.bytes_per_page_padded, num_host_pages=num_host_pages, slot_page_size=page_size, + layout=server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) c128_state_host_pool = DeepSeekV4StateHostPool( @@ -440,6 +446,7 @@ def build_deepseek_v4_hicache_stack( ], num_host_pages=swa_num_host_pages, swa_page_size=kvcache.swa_page_size, + layout=server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) entries.extend( diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index c8d3b9a7bf55..b4624ab54a9f 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1754,6 +1754,7 @@ def __init__( item_bytes: int, num_host_pages: int, slot_page_size: int, + layout: str = "layer_first", device: str = "cpu", pin_memory: bool = True, allocator_type: str = "default", @@ -1769,7 +1770,7 @@ def __init__( self.allocator = get_allocator_from_storage(allocator_type) self.page_size = slot_page_size self.size = num_host_pages * slot_page_size - self.layout = "layer_first" + self.layout = layout self.size_per_token = item_bytes self.start_layer = 0 self.end_layer = self.layer_num @@ -1789,26 +1790,62 @@ def __init__( ) alloc_func = ALLOC_MEMORY_FUNCS[self.gpu_device] - self.kv_buffer = [ - alloc_func( - (num_host_pages, self.item_bytes), + self.data_refs = [] + if self.layout == "layer_first": + self.kv_buffer = [ + alloc_func( + (num_host_pages, self.item_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + for _ in range(self.layer_num) + ] + self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + elif self.layout == "page_first": + self.kv_buffer = alloc_func( + (num_host_pages, self.layer_num, self.item_bytes), dtype=self.dtype, device=self.device, pin_memory=self.pin_memory, allocator=self.allocator, ) - for _ in range(self.layer_num) - ] - self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + elif self.layout == "page_first_direct": + self.kv_buffer = alloc_func( + (num_host_pages, self.layer_num, 1, self.item_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") logger.info( "Allocating %.2f GB host memory for V4 paged pool '%s' " - "(layers=%d, pages=%d, item_bytes=%d).", + "(layers=%d, pages=%d, item_bytes=%d, layout=%s).", requested_bytes / 1e9, self.pool_name, self.layer_num, num_host_pages, self.item_bytes, + self.layout, + ) + + self.device_ptrs = torch.tensor( + [x.data_ptr() for x in self.device_buffers], + dtype=torch.uint64, + device=self.gpu_device, + ) + self.data_ptrs = ( + torch.tensor( + [x.data_ptr() for x in self.data_refs], + dtype=torch.uint64, + device=self.gpu_device, + ) + if self.data_refs + else None ) self.clear() @@ -1820,12 +1857,6 @@ def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: ) return indices.reshape(-1, self.slot_page_size)[:, 0] // self.slot_page_size - def _check_io_backend(self, io_backend: str) -> None: - if io_backend != "direct": - raise NotImplementedError( - f"{self.pool_name} supports only direct io_backend, got {io_backend}" - ) - def get_size_per_token(self): return self.item_bytes @@ -1836,7 +1867,7 @@ def init_kv_buffer(self): return self.kv_buffer def get_hybrid_pool_buffer(self): - return self.kv_buffer + return self.kv_buffer if isinstance(self.kv_buffer, list) else [self.kv_buffer] def clear(self): self.free_slots = torch.arange(self.size, dtype=torch.int64) @@ -1867,38 +1898,106 @@ def backup_from_device_all_layer( ): if host_indices is None or device_indices is None: return - self._check_io_backend(io_backend) host_rows = self._to_page_indices(host_indices) device_rows = self._to_page_indices(device_indices) - transfer_kv_direct( - src_layers=self.device_buffers, - dst_layers=self.data_refs, - src_indices=device_rows, - dst_indices=host_rows, - page_size=1, - ) + if io_backend == "kernel" and self.layout == "layer_first": + assert self.data_ptrs is not None + transfer_kv_all_layer_mla( + src_layers=self.device_ptrs, + dst_layers=self.data_ptrs, + src_indices=device_rows, + dst_indices=host_rows, + item_size=self.item_bytes, + num_layers=self.layer_num, + ) + elif io_backend == "kernel" and self.layout == "page_first": + transfer_kv_all_layer_mla_lf_pf( + src_layers=self.device_ptrs, + dst=self.kv_buffer, + src_indices=device_rows, + dst_indices=host_rows, + item_size=self.item_bytes, + dst_layout_dim=self.layer_num * self.item_bytes, + num_layers=self.layer_num, + ) + elif io_backend == "direct" and self.layout == "layer_first": + transfer_kv_direct( + src_layers=self.device_buffers, + dst_layers=self.data_refs, + src_indices=device_rows, + dst_indices=host_rows, + page_size=1, + ) + elif io_backend == "direct" and self.layout == "page_first_direct": + transfer_kv_all_layer_direct_lf_pf( + src_ptrs=self.device_buffers, + dst_ptrs=[self.kv_buffer], + src_indices=device_rows, + dst_indices=host_rows, + page_size=1, + ) + else: + raise ValueError( + f"Unsupported V4 paged host layout/backend: {self.layout}/{io_backend}" + ) def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): if host_indices is None or device_indices is None: return - self._check_io_backend(io_backend) host_rows = self._to_page_indices(host_indices) device_rows = self._to_page_indices(device_indices) - transfer_kv_direct( - src_layers=[self.kv_buffer[layer_id]], - dst_layers=[self.device_buffers[layer_id]], - src_indices=host_rows, - dst_indices=device_rows, - page_size=1, - ) + if io_backend == "kernel" and self.layout == "layer_first": + transfer_kv_per_layer_mla( + src=self.data_refs[layer_id], + dst=self.device_buffers[layer_id], + src_indices=host_rows, + dst_indices=device_rows, + item_size=self.item_bytes, + ) + elif io_backend == "kernel" and self.layout == "page_first": + transfer_kv_per_layer_mla_pf_lf( + src=self.kv_buffer, + dst=self.device_buffers[layer_id], + src_indices=host_rows, + dst_indices=device_rows, + layer_id=layer_id, + item_size=self.item_bytes, + src_layout_dim=self.layer_num * self.item_bytes, + ) + elif io_backend == "direct" and self.layout == "layer_first": + transfer_kv_direct( + src_layers=[self.data_refs[layer_id]], + dst_layers=[self.device_buffers[layer_id]], + src_indices=host_rows, + dst_indices=device_rows, + page_size=1, + ) + elif io_backend == "direct" and self.layout == "page_first_direct": + transfer_kv_per_layer_direct_pf_lf( + src_ptrs=[self.kv_buffer], + dst_ptrs=[self.device_buffers[layer_id]], + src_indices=host_rows, + dst_indices=device_rows, + layer_id=layer_id, + page_size=1, + ) + else: + raise ValueError( + f"Unsupported V4 paged host layout/backend: {self.layout}/{io_backend}" + ) def get_data_page(self, index, flat=True): index = int(index) // self.slot_page_size - data_page = torch.stack( - [self.kv_buffer[i][index] for i in range(self.layer_num)] - ) + if self.layout == "layer_first": + data_page = torch.stack( + [self.kv_buffer[i][index] for i in range(self.layer_num)] + ) + elif self.layout in ["page_first", "page_first_direct"]: + data_page = self.kv_buffer[index] + else: + raise ValueError(f"Unsupported layout: {self.layout}") return data_page.flatten() if flat else data_page def get_dummy_flat_data_page(self): @@ -1911,22 +2010,41 @@ def get_dummy_flat_data_page(self): def set_from_flat_data_page(self, index, data_page): index = int(index) // self.slot_page_size - data = data_page.view(self.dtype).reshape(self.layer_num, self.item_bytes) - for i in range(self.layer_num): - self.kv_buffer[i][index].copy_(data[i]) + if self.layout == "layer_first": + data = data_page.view(self.dtype).reshape(self.layer_num, self.item_bytes) + for i in range(self.layer_num): + self.kv_buffer[i][index].copy_(data[i]) + elif self.layout == "page_first": + self.kv_buffer[index].copy_( + data_page.view(self.dtype).reshape(self.layer_num, self.item_bytes) + ) + elif self.layout == "page_first_direct": + self.kv_buffer[index].copy_( + data_page.view(self.dtype).reshape(self.layer_num, 1, self.item_bytes) + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") def get_page_buffer_meta(self, indices): ptr_list = [] rows = self._to_page_indices(indices).tolist() - for row in rows: - for layer_id in range(self.layer_num): - ptr = ( - self.kv_buffer[layer_id].data_ptr() - + int(row) * self.item_bytes * self.dtype.itemsize - ) - ptr_list.append(ptr) - element_size = self.item_bytes * self.dtype.itemsize - return ptr_list, [element_size] * len(ptr_list) + if self.layout == "layer_first": + for row in rows: + page_index = int(row) + for layer_id in range(self.layer_num): + ptr = ( + self.kv_buffer[layer_id].data_ptr() + + page_index * self.item_bytes * self.dtype.itemsize + ) + ptr_list.append(ptr) + element_size = self.item_bytes * self.dtype.itemsize + return ptr_list, [element_size] * len(ptr_list) + if self.layout in ["page_first", "page_first_direct"]: + page_bytes = self.layer_num * self.item_bytes * self.dtype.itemsize + for row in rows: + ptr_list.append(self.kv_buffer[int(row)].data_ptr()) + return ptr_list, [page_bytes] * len(ptr_list) + raise ValueError(f"Unsupported layout: {self.layout}") class DeepSeekV4StateHostPool(HostKVCache): @@ -1938,6 +2056,7 @@ def __init__( state_pools: list, num_host_pages: int, swa_page_size: int, + layout: str = "layer_first", device: str = "cpu", pin_memory: bool = True, allocator_type: str = "default", @@ -1956,7 +2075,7 @@ def __init__( self.allocator = get_allocator_from_storage(allocator_type) self.page_size = swa_page_size self.size = num_host_pages * swa_page_size - self.layout = "layer_first" + self.layout = layout self.start_layer = 0 self.end_layer = self.layer_num self.lock = threading.RLock() @@ -1979,25 +2098,60 @@ def __init__( ) alloc_func = ALLOC_MEMORY_FUNCS[self.gpu_device] - self.kv_buffer = [ - alloc_func( - (num_host_pages, self.state_page_bytes), + self.data_refs = [] + if self.layout == "layer_first": + self.kv_buffer = [ + alloc_func( + (num_host_pages, self.state_page_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + for _ in range(self.layer_num) + ] + self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + elif self.layout == "page_first": + self.kv_buffer = alloc_func( + (num_host_pages, self.layer_num, self.state_page_bytes), dtype=self.dtype, device=self.device, pin_memory=self.pin_memory, allocator=self.allocator, ) - for _ in range(self.layer_num) - ] - self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + elif self.layout == "page_first_direct": + self.kv_buffer = alloc_func( + (num_host_pages, self.layer_num, 1, self.state_page_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") logger.info( "Allocating %.2f GB host memory for V4 state pool '%s' " - "(layers=%d, pages=%d, state_page_bytes=%d).", + "(layers=%d, pages=%d, state_page_bytes=%d, layout=%s).", requested_bytes / 1e9, self.pool_name, self.layer_num, num_host_pages, self.state_page_bytes, + self.layout, + ) + self.device_ptrs = torch.tensor( + [x.data_ptr() for x in self.device_page_views], + dtype=torch.uint64, + device=self.gpu_device, + ) + self.data_ptrs = ( + torch.tensor( + [x.data_ptr() for x in self.data_refs], + dtype=torch.uint64, + device=self.gpu_device, + ) + if self.data_refs + else None ) def _init_device_page_views(self) -> None: @@ -2041,12 +2195,6 @@ def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: ) return indices.reshape(-1, self.swa_page_size)[:, 0] // self.swa_page_size - def _check_io_backend(self, io_backend: str) -> None: - if io_backend != "direct": - raise NotImplementedError( - f"{self.pool_name} supports only direct io_backend, got {io_backend}" - ) - def get_size_per_token(self): return self.state_page_bytes @@ -2057,7 +2205,7 @@ def init_kv_buffer(self): return self.kv_buffer def get_hybrid_pool_buffer(self): - return self.kv_buffer + return self.kv_buffer if isinstance(self.kv_buffer, list) else [self.kv_buffer] def clear(self): pass @@ -2084,38 +2232,106 @@ def backup_from_device_all_layer( ): if host_indices is None or device_indices is None: return - self._check_io_backend(io_backend) host_rows = self._to_page_indices(host_indices) device_rows = self._to_page_indices(device_indices) - transfer_kv_direct( - src_layers=self.device_page_views, - dst_layers=self.data_refs, - src_indices=device_rows, - dst_indices=host_rows, - page_size=1, - ) + if io_backend == "kernel" and self.layout == "layer_first": + assert self.data_ptrs is not None + transfer_kv_all_layer_mla( + src_layers=self.device_ptrs, + dst_layers=self.data_ptrs, + src_indices=device_rows, + dst_indices=host_rows, + item_size=self.state_page_bytes, + num_layers=self.layer_num, + ) + elif io_backend == "kernel" and self.layout == "page_first": + transfer_kv_all_layer_mla_lf_pf( + src_layers=self.device_ptrs, + dst=self.kv_buffer, + src_indices=device_rows, + dst_indices=host_rows, + item_size=self.state_page_bytes, + dst_layout_dim=self.layer_num * self.state_page_bytes, + num_layers=self.layer_num, + ) + elif io_backend == "direct" and self.layout == "layer_first": + transfer_kv_direct( + src_layers=self.device_page_views, + dst_layers=self.data_refs, + src_indices=device_rows, + dst_indices=host_rows, + page_size=1, + ) + elif io_backend == "direct" and self.layout == "page_first_direct": + transfer_kv_all_layer_direct_lf_pf( + src_ptrs=self.device_page_views, + dst_ptrs=[self.kv_buffer], + src_indices=device_rows, + dst_indices=host_rows, + page_size=1, + ) + else: + raise ValueError( + f"Unsupported V4 state host layout/backend: {self.layout}/{io_backend}" + ) def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): if host_indices is None or device_indices is None: return - self._check_io_backend(io_backend) host_rows = self._to_page_indices(host_indices) device_rows = self._to_page_indices(device_indices) - transfer_kv_direct( - src_layers=[self.kv_buffer[layer_id]], - dst_layers=[self.device_page_views[layer_id]], - src_indices=host_rows, - dst_indices=device_rows, - page_size=1, - ) + if io_backend == "kernel" and self.layout == "layer_first": + transfer_kv_per_layer_mla( + src=self.data_refs[layer_id], + dst=self.device_page_views[layer_id], + src_indices=host_rows, + dst_indices=device_rows, + item_size=self.state_page_bytes, + ) + elif io_backend == "kernel" and self.layout == "page_first": + transfer_kv_per_layer_mla_pf_lf( + src=self.kv_buffer, + dst=self.device_page_views[layer_id], + src_indices=host_rows, + dst_indices=device_rows, + layer_id=layer_id, + item_size=self.state_page_bytes, + src_layout_dim=self.layer_num * self.state_page_bytes, + ) + elif io_backend == "direct" and self.layout == "layer_first": + transfer_kv_direct( + src_layers=[self.data_refs[layer_id]], + dst_layers=[self.device_page_views[layer_id]], + src_indices=host_rows, + dst_indices=device_rows, + page_size=1, + ) + elif io_backend == "direct" and self.layout == "page_first_direct": + transfer_kv_per_layer_direct_pf_lf( + src_ptrs=[self.kv_buffer], + dst_ptrs=[self.device_page_views[layer_id]], + src_indices=host_rows, + dst_indices=device_rows, + layer_id=layer_id, + page_size=1, + ) + else: + raise ValueError( + f"Unsupported V4 state host layout/backend: {self.layout}/{io_backend}" + ) def get_data_page(self, index, flat=True): index = int(index) // self.swa_page_size - data_page = torch.stack( - [self.kv_buffer[i][index] for i in range(self.layer_num)] - ) + if self.layout == "layer_first": + data_page = torch.stack( + [self.kv_buffer[i][index] for i in range(self.layer_num)] + ) + elif self.layout in ["page_first", "page_first_direct"]: + data_page = self.kv_buffer[index] + else: + raise ValueError(f"Unsupported layout: {self.layout}") return data_page.flatten() if flat else data_page def get_dummy_flat_data_page(self): @@ -2128,22 +2344,47 @@ def get_dummy_flat_data_page(self): def set_from_flat_data_page(self, index, data_page): index = int(index) // self.swa_page_size - data = data_page.view(self.dtype).reshape(self.layer_num, self.state_page_bytes) - for i in range(self.layer_num): - self.kv_buffer[i][index].copy_(data[i]) + if self.layout == "layer_first": + data = data_page.view(self.dtype).reshape( + self.layer_num, self.state_page_bytes + ) + for i in range(self.layer_num): + self.kv_buffer[i][index].copy_(data[i]) + elif self.layout == "page_first": + self.kv_buffer[index].copy_( + data_page.view(self.dtype).reshape( + self.layer_num, self.state_page_bytes + ) + ) + elif self.layout == "page_first_direct": + self.kv_buffer[index].copy_( + data_page.view(self.dtype).reshape( + self.layer_num, 1, self.state_page_bytes + ) + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") def get_page_buffer_meta(self, indices): ptr_list = [] rows = self._to_page_indices(indices).tolist() - for row in rows: - for layer_id in range(self.layer_num): - ptr = ( - self.kv_buffer[layer_id].data_ptr() - + int(row) * self.state_page_bytes * self.dtype.itemsize - ) - ptr_list.append(ptr) - element_size = self.state_page_bytes * self.dtype.itemsize - return ptr_list, [element_size] * len(ptr_list) + if self.layout == "layer_first": + for row in rows: + page_index = int(row) + for layer_id in range(self.layer_num): + ptr = ( + self.kv_buffer[layer_id].data_ptr() + + page_index * self.state_page_bytes * self.dtype.itemsize + ) + ptr_list.append(ptr) + element_size = self.state_page_bytes * self.dtype.itemsize + return ptr_list, [element_size] * len(ptr_list) + if self.layout in ["page_first", "page_first_direct"]: + page_bytes = self.layer_num * self.state_page_bytes * self.dtype.itemsize + for row in rows: + ptr_list.append(self.kv_buffer[int(row)].data_ptr()) + return ptr_list, [page_bytes] * len(ptr_list) + raise ValueError(f"Unsupported layout: {self.layout}") @dataclass diff --git a/python/sglang/test/kl_multiturn_utils.py b/python/sglang/test/kl_multiturn_utils.py index bd21c502321b..97219b6d3ef7 100644 --- a/python/sglang/test/kl_multiturn_utils.py +++ b/python/sglang/test/kl_multiturn_utils.py @@ -2,6 +2,7 @@ from __future__ import annotations +import time from typing import Callable from sglang.test.kl_test_utils import ( @@ -145,7 +146,13 @@ def _interleave_order(n: int, branches_per_group: int) -> list[int] | None: def _generate_maybe_interleaved( - base_url, inputs, max_new_tokens, order=None, sampling_temperature: float = 1 + base_url, + inputs, + max_new_tokens, + order=None, + sampling_temperature: float = 1, + request_batch_size: int | None = None, + inter_batch_delay_s: float = 0, ): """Generate with optional interleaved submission order. @@ -153,22 +160,31 @@ def _generate_maybe_interleaved( original order so the caller always sees results[i] corresponds to inputs[i]. """ - if order is None: - return _generate( - base_url, - inputs, - max_new_tokens, - return_logprob=True, - temperature=sampling_temperature, - ) - ordered = [inputs[i] for i in order] - results = _generate( - base_url, - ordered, - max_new_tokens, - return_logprob=True, - temperature=sampling_temperature, + ordered = inputs if order is None else [inputs[i] for i in order] + if not ordered: + return [] + + batch_size = ( + request_batch_size + if request_batch_size is not None and request_batch_size > 0 + else len(ordered) ) + results = [] + for start in range(0, len(ordered), batch_size): + results.extend( + _generate( + base_url, + ordered[start : start + batch_size], + max_new_tokens, + return_logprob=True, + temperature=sampling_temperature, + ) + ) + if batch_size < len(ordered) and inter_batch_delay_s > 0: + time.sleep(inter_batch_delay_s) + + if order is None: + return results unordered = [None] * len(results) for idx, orig in enumerate(order): unordered[orig] = results[idx] @@ -423,6 +439,8 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( branches_per_group: int = 0, replay_batch_size: int = 1, sampling_temperature: float = 1, + request_batch_size: int | None = None, + inter_batch_delay_s: float = 0, ): """Verify logprobs when decode cache is hit. @@ -453,12 +471,13 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( # Turn 1: populate cache, no assertion, no interleaving _flush_cache(base_url) - results = _generate( + results = _generate_maybe_interleaved( base_url, first_turn_input_ids, max_new_tokens, - return_logprob=True, - temperature=sampling_temperature, + sampling_temperature=sampling_temperature, + request_batch_size=request_batch_size, + inter_batch_delay_s=inter_batch_delay_s, ) assert len(results) == n @@ -478,6 +497,8 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( max_new_tokens, order, sampling_temperature=sampling_temperature, + request_batch_size=request_batch_size, + inter_batch_delay_s=inter_batch_delay_s, ) assert len(results) == n diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl.py b/test/registered/radix_cache/test_unified_radix_cache_kl.py index 2d8ab360f3be..d7e4c5e05a51 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl.py @@ -49,6 +49,8 @@ class UnifiedRadixTreeTestMixin: prefill_cache_assert = None decode_cache_assert = None sampling_temperature: float = 1 + decode_hit_request_batch_size: int | None = None + decode_hit_inter_batch_delay_s: float = 0 gsm8k_threshold: float = 0.93 mmlu_threshold: float = 0.8 @@ -163,6 +165,8 @@ def test_multiturn_decode_cache_hit_branching(self): branches_per_group=branches, max_new_tokens=self.max_new_tokens, sampling_temperature=self.sampling_temperature, + request_batch_size=self.decode_hit_request_batch_size, + inter_batch_delay_s=self.decode_hit_inter_batch_delay_s, ) diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py index 2f3fc32ba071..d43b6fffc66b 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py @@ -92,8 +92,13 @@ def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label): class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): """DeepSeek V4 Flash FP8 + HiCache + UnifiedRadixCache.""" + hicache_io_backend = "direct" + hicache_mem_layout = "page_first_direct" + max_running_requests = 4 kl_threshold = 0.005 sampling_temperature = 0 + decode_hit_request_batch_size = 3 + decode_hit_inter_batch_delay_s = 0.5 decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens) gsm8k_threshold = 0.90 num_gsm8k_questions = 100 @@ -129,15 +134,15 @@ def setUpClass(cls): "--hicache-write-policy", "write_through", "--hicache-io-backend", - "direct", + cls.hicache_io_backend, "--hicache-mem-layout", - "page_first_direct", + cls.hicache_mem_layout, "--swa-full-tokens-ratio", "0.25", "--max-total-tokens", "20000", "--max-running-requests", - "2", + str(cls.max_running_requests), ], env={ "SGLANG_DSV4_FP4_EXPERTS": "0", @@ -151,5 +156,14 @@ def tearDownClass(cls): kill_process_tree(cls.process.pid) +class TestUnifiedDeepSeekV4FlashHiCachePageFirstDirect( + TestUnifiedDeepSeekV4FlashHiCache +): + """DeepSeek V4 Flash HiCache layout smoke: page_first_direct + direct.""" + + hicache_io_backend = "kernel" + hicache_mem_layout = "layer_first" + + if __name__ == "__main__": unittest.main()