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

Co-authored-by: hzh0425 <hzh0425@apache.org>

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

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,
         )
 

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
 
 

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

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])
+            if self.storage_config.should_split_heads:
-            return [
+                kv_groups = [
-                (
+                    results[i : i + self.split_factor * 2]
-                    (k_res == 0 and v_res == 0)
+                    for i in range(0, len(results), self.split_factor * 2)
-                    if is_set_operate
+                ]
-                    else (k_res > 0 and v_res > 0)
+                return [
-                )
+                    (
-                for k_res, v_res in kv_pairs
+                        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:
+            if self.storage_config.should_split_heads:
-                query_keys.append(f"{key}_{self.mha_suffix}_k")
+                for key in keys:
-                query_keys.append(f"{key}_{self.mha_suffix}_v")
+                    for suffix in self.mha_suffix:
-            key_multiplier = 2
+                        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)):