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

[HiCache] Support heterogeneous tp for hicache storage (#18541)

Browse Source 
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
huangtingwei
2026-02-25 16:13:57 +08:00
committed by GitHub
parent 3d879c69e9
commit d40cb2f725
5 changed files with 149 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/managers/cache_controller.py
View File

@@ -583,6 +583,19 @@ class HiCacheController:
# Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool.
is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool)
# Least Common Multiple among heterogeneous tp size
tp_lcm_size = storage_backend_extra_config.pop("tp_lcm_size", None)
should_split_heads = False
if tp_lcm_size:
assert (
tp_lcm_size % self.tp_size == 0
), "tp_lcm_size must be divisible by tp_size."
should_split_heads = (
not is_mla_backend
and self.mem_pool_host.layout == "page_head"
and tp_lcm_size > self.tp_size
)
return HiCacheStorageConfig(
tp_rank=self.tp_rank,
@@ -592,6 +605,8 @@ class HiCacheController:
is_mla_model=is_mla_backend,
is_page_first_layout=self.mem_pool_host.layout == "page_first",
model_name=model_name,
tp_lcm_size=tp_lcm_size,
should_split_heads=should_split_heads,
extra_config=storage_backend_extra_config,
)

2
python/sglang/srt/mem_cache/hicache_storage.py
View File

@@ -53,6 +53,8 @@ class HiCacheStorageConfig:
is_mla_model: bool
is_page_first_layout: bool
model_name: Optional[str]
tp_lcm_size: Optional[int] = None
should_split_heads: bool = False
extra_config: Optional[dict] = None

48
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -620,6 +620,54 @@ class MHATokenToKVPoolHost(HostKVCache):
else:
raise ValueError(f"Unsupported layout: {self.layout}")
def get_split_heads_page_buffer_meta(
self, indices: torch.Tensor, split_factor: int
):
"""
get meta data for zero copy of heterogeneous ranks' KVCache
"""
assert self.layout == "page_head"
assert len(indices) % self.page_size == 0
assert self.head_num % split_factor == 0
ptr_list = []
kv_buffer_data_ptr = self.kv_buffer.data_ptr()
indices = indices.tolist()
v_offset = (
self.layer_num
* self.size
* self.head_num
* self.head_dim
* self.dtype.itemsize
)
for index in range(0, len(indices), self.page_size):
for head_id in range(0, self.head_num, self.head_num // split_factor):
k_ptr = (
kv_buffer_data_ptr
+ indices[index]
* self.layer_num
* self.head_num
* self.head_dim
* self.dtype.itemsize
+ head_id
* self.page_size
* self.layer_num
* self.head_dim
* self.dtype.itemsize
)
v_ptr = k_ptr + v_offset
ptr_list.append(k_ptr)
ptr_list.append(v_ptr)
element_size = (
self.layer_num
* self.dtype.itemsize
* self.page_size
* self.head_num
* self.head_dim
// split_factor
)
element_size_list = [element_size] * len(ptr_list)
return ptr_list, element_size_list
def get_page_buffer_meta(self, indices):
""" "
meta data for zero copy

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

@@ -407,6 +407,21 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
self.mha_suffix = f"{self.local_rank}"
self.mla_suffix = ""
self.storage_config = storage_config
self.split_factor = 0
if self.storage_config.should_split_heads:
self.split_factor = (
self.storage_config.tp_lcm_size // self.storage_config.tp_size
)
base_rank = self.local_rank * self.split_factor
target_ranks = [base_rank + i for i in range(self.split_factor)]
if self.enable_pp:
self.mha_suffix = [
f"{rank}_{self.pp_rank}" for rank in target_ranks
]
else:
self.mha_suffix = [f"{rank}" for rank in target_ranks]
self.gb_per_page = None
self.prefetch_pgs = []
self.backup_pgs = []
@@ -477,6 +492,20 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
bytes_per_page = mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size
self.gb_per_page = bytes_per_page / (1 << 30)
def _get_mha_split_heads_buffer_meta(self, keys, indices):
ptr_list, element_size_list = (
self.mem_pool_host.get_split_heads_page_buffer_meta(
indices, self.split_factor
)
)
key_list = []
for key_ in keys:
for suffix in self.mha_suffix:
key_list.append(f"{key_}_{suffix}_k")
key_list.append(f"{key_}_{suffix}_v")
assert len(key_list) == len(ptr_list)
return key_list, ptr_list, element_size_list
def _get_mha_buffer_meta(self, keys, indices):
ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices)
key_list = []
@@ -500,7 +529,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
if self.is_mla_backend:
return self._get_mla_buffer_meta(keys, host_indices)
else:
return self._get_mha_buffer_meta(keys, host_indices)
if self.storage_config.should_split_heads:
return self._get_mha_split_heads_buffer_meta(keys, host_indices)
else:
return self._get_mha_buffer_meta(keys, host_indices)
def _batch_postprocess(self, results: List[int], is_set_operate=False):
"""
@@ -513,15 +545,29 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
if self.is_mla_backend:
return [k_res == 0 if is_set_operate else k_res > 0 for k_res in results]
else:
kv_pairs = zip(results[::2], results[1::2])
return [
(
(k_res == 0 and v_res == 0)
if is_set_operate
else (k_res > 0 and v_res > 0)
)
for k_res, v_res in kv_pairs
]
if self.storage_config.should_split_heads:
kv_groups = [
results[i : i + self.split_factor * 2]
for i in range(0, len(results), self.split_factor * 2)
]
return [
(
all(res == 0 for res in kv_group)
if is_set_operate
else all(res > 0 for res in kv_group)
)
for kv_group in kv_groups
]
else:
kv_pairs = zip(results[::2], results[1::2])
return [
(
(k_res == 0 and v_res == 0)
if is_set_operate
else (k_res > 0 and v_res > 0)
)
for k_res, v_res in kv_pairs
]
def batch_get_v1(
self,
@@ -712,10 +758,17 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
key_multiplier = 1
else:
query_keys = []
for key in keys:
query_keys.append(f"{key}_{self.mha_suffix}_k")
query_keys.append(f"{key}_{self.mha_suffix}_v")
key_multiplier = 2
if self.storage_config.should_split_heads:
for key in keys:
for suffix in self.mha_suffix:
query_keys.append(f"{key}_{suffix}_k")
query_keys.append(f"{key}_{suffix}_v")
key_multiplier = 2 * self.split_factor
else:
for key in keys:
query_keys.append(f"{key}_{self.mha_suffix}_k")
query_keys.append(f"{key}_{self.mha_suffix}_v")
key_multiplier = 2
exist_result = self._batch_exist(query_keys)
for i in range(len(query_keys)):