diff --git a/docs/advanced_features/hicache_best_practices.md b/docs/advanced_features/hicache_best_practices.md index 02749530a..104c2b0e2 100644 --- a/docs/advanced_features/hicache_best_practices.md +++ b/docs/advanced_features/hicache_best_practices.md @@ -39,6 +39,23 @@ Notes: - `page_first`: Only compatible with `kernel` I/O backend, automatically switches to `layer_first` with `direct` backend - `page_first_direct`: Specifically designed for `direct` I/O backend with optimized memory organization +### Heterogeneous TP Support (GQA/MHA models) + +HiCache storage supports cross-cluster KV reuse when different deployments use different TP sizes (for example, `tp=4` and `tp=8`) and share the same storage backend namespace. + +Use `tp_lcm_size` in `--hicache-storage-backend-extra-config`: + +```bash +# Example: heterogeneous TP = {4, 8}, so lcm = 8 +--hicache-storage-backend-extra-config '{"tp_lcm_size": 8}' +``` + +Guidelines: + +- Set `tp_lcm_size` to the least common multiple (LCM) of all TP sizes that will share the same HiCache storage. +- For MHA models with Mooncake and `page_head` layout, HiCache will split head shards based on `tp_lcm_size` to make keys reusable across heterogeneous TP deployments. +- If all clusters use the same TP size, this option is not needed. + ### Prefetch Policies ```bash diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 34f78c889..51030ae1d 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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, ) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index daa79f93f..02680db04 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -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 diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index f3f7c0bb5..0b8199d9f 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -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 diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 7fda96124..892ab0d46 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -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)):