From 13e74b573fda9679083c57c50f0dc6a7c585f772 Mon Sep 17 00:00:00 2001 From: LCAIZJ Date: Mon, 3 Nov 2025 11:26:54 +0800 Subject: [PATCH] add alloc_in_same_node as a configurable parameter Signed-off-by: LCAIZJ --- vllm_ascend/distributed/mooncake/config_data.py | 5 ++++- vllm_ascend/distributed/mooncake/mooncake_store.py | 14 +++++++++----- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/vllm_ascend/distributed/mooncake/config_data.py b/vllm_ascend/distributed/mooncake/config_data.py index 745d91131fa..d36b4e58b59 100644 --- a/vllm_ascend/distributed/mooncake/config_data.py +++ b/vllm_ascend/distributed/mooncake/config_data.py @@ -425,6 +425,7 @@ class MooncakeStoreConfig: device_name: str master_server_address: str use_ascend_direct: bool + alloc_in_same_node: bool @staticmethod def from_file(file_path: str) -> "MooncakeStoreConfig": @@ -438,7 +439,9 @@ def from_file(file_path: str) -> "MooncakeStoreConfig": protocol=config.get("protocol", "tcp"), device_name=config.get("device_name", ""), master_server_address=config.get("master_server_address"), - use_ascend_direct=config.get("use_ascend_direct", False)) + use_ascend_direct=config.get("use_ascend_direct", False), + alloc_in_same_node=config.get("alloc_in_same_node", False), + ) @staticmethod def load_from_env() -> "MooncakeStoreConfig": diff --git a/vllm_ascend/distributed/mooncake/mooncake_store.py b/vllm_ascend/distributed/mooncake/mooncake_store.py index cee07c6a251..326c1789ed4 100644 --- a/vllm_ascend/distributed/mooncake/mooncake_store.py +++ b/vllm_ascend/distributed/mooncake/mooncake_store.py @@ -88,11 +88,15 @@ def get_batch(self, keys: list[str], addrs: list[list[int]], def put_batch(self, keys: list[str], addrs: list[list[int]], sizes: list[list[int]], block_ids: list[int]): try: - config = ReplicateConfig() - config.preferred_segment = self.local_seg - config.prefer_alloc_in_same_node = True - res = self.store.batch_put_from_multi_buffers( - keys, addrs, sizes, config) + if self.config.alloc_in_same_node: + config = ReplicateConfig() + config.preferred_segment = self.local_seg + config.prefer_alloc_in_same_node = True + res = self.store.batch_put_from_multi_buffers( + keys, addrs, sizes, config) + else: + res = self.store.batch_put_from_multi_buffers( + keys, addrs, sizes) for value in res: if value < 0: logger.error(f"Failed to put key {keys},res:{res}")