EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[HiCacheStorage] mooncake store support page_first_direct layout (#10591)

Browse Source
This commit is contained in:
huangtingwei
2025-09-29 11:45:48 +08:00
committed by GitHub
parent 43fa9f22bd
commit e05555fad8
4 changed files with 30 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -48,9 +48,9 @@ class HiRadixCache(RadixCache):
if hicache_io_backend == "direct":
if hicache_mem_layout == "page_first":
hicache_mem_layout = "layer_first"
hicache_mem_layout = "page_first_direct"
logger.warning(
"Page first layout is not supported with direct IO backend, switching to layer first layout"
"Page first layout is not supported with direct IO backend, switching to page first direct layout"
)
self.kv_cache = token_to_kv_pool_allocator.get_kvcache()

7
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
View File

@@ -190,9 +190,10 @@ class MooncakeStore(HiCacheStorage):
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
super().register_mem_pool_host(mem_pool_host)
assert (
self.mem_pool_host.layout == "page_first"
), "mooncake store storage backend only support page first layout"
assert self.mem_pool_host.layout in [
"page_first",
"page_first_direct",
], "mooncake store storage backend only support page first or page first direct layout"
buffer = self.mem_pool_host.kv_buffer
try:
buffer_ptr = buffer.data_ptr()

11
python/sglang/srt/server_args.py
View File

@@ -930,8 +930,15 @@ class ServerArgs:
def _handle_hicache(self):
if self.hicache_storage_backend == "mooncake":
self.hicache_io_backend = "kernel"
self.hicache_mem_layout = "page_first"
if self.hicache_mem_layout == "layer_first":
if self.hicache_io_backend == "direct":
self.hicache_mem_layout = "page_first_direct"
elif self.hicache_io_backend == "kernel":
self.hicache_mem_layout = "page_first"
logger.warning(
f"Mooncake storage backend does not support layer_first layout, "
f"switching to {self.hicache_mem_layout} layout for {self.hicache_io_backend} io backend"
)
if self.hicache_mem_layout == "page_first_direct":
if self.hicache_io_backend != "direct":