modify local_rank_size

Author: qyh111

ucm/integration/vllm/ucm_connector.py

@@ -98,6 +98,7 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         self.num_layers = self._vllm_config.model_config.get_num_layers(
             self._vllm_config.parallel_config
         )
+        self.tp_size = self._vllm_config.parallel_config.tensor_parallel_size
         self.kv_cache_dtype: torch.dtype = None
 
         if current_platform.is_cuda_alike():
@@ -218,6 +219,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         config = self.connector_configs[0].get("ucm_connector_config") or {}
         config["device"] = self.local_rank
         config["role"] = "worker"
+        config["local_rank_size"] = self.tp_size if self.is_mla or self.is_dsa else 1
         if len(sample_kv_layer) == 2:
             k_io_size = (
                 sample_kv_layer[0][0].numel() * sample_kv_layer[0][0].element_size()

ucm/store/pcstore/pcstore_connector.py

@@ -52,8 +52,7 @@ def __init__(self, config: Dict):
             param.transferIoDirect = config.get("use_direct", False)
             param.transferStreamNumber = config.get("stream_number", 8)
             param.transferBufferNumber = config.get("buffer_number", 4096)
-            param.transferLocalRankSize = config.get("local_rank_size", 1)
-            param.transferScatterGatherEnable = config.get("use_scatter_gatter", False)
+            param.transferScatterGatherEnable = config.get("use_scatter_gatter", True)
         ret = self.store.Setup(param)
         if ret != 0:
             msg = f"Failed to initialize ucmpcstore, errcode: {ret}."