Skip to content

[Store] Add alloc_in_same_node as a configurable parameter #3952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+13 −6
Open

[Store] Add alloc_in_same_node as a configurable parameter #3952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion vllm_ascend/distributed/mooncake/config_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand All @@ -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":
Expand Down
14 changes: 9 additions & 5 deletions vllm_ascend/distributed/mooncake/mooncake_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment on lines +91 to +99
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

There's a potential AttributeError here because self.local_seg is used but it's not guaranteed to be initialized in Mooncakestore.__init__. Specifically, if self.config.protocol == "ascend" and self.config.use_ascend_direct is False, self.local_seg is not set, which would cause a crash if alloc_in_same_node is True.

The suggested change fixes this by explicitly checking for self.local_seg and raising an AttributeError with a helpful message if it's missing. This makes configuration errors easier to debug. It also refactors the logic to be more concise and avoid duplicating the call to batch_put_from_multi_buffers.

Suggested change
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)
put_args = [keys, addrs, sizes]
if self.config.alloc_in_same_node:
if not hasattr(self, "local_seg"):
raise AttributeError("alloc_in_same_node is True, but local_seg is not initialized.")
config = ReplicateConfig()
config.preferred_segment = self.local_seg
config.prefer_alloc_in_same_node = True
put_args.append(config)
res = self.store.batch_put_from_multi_buffers(*put_args)

Copy link
Contributor

@Pz1116 Pz1116 Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we give an example of mooncake config json file with alloc_in_same_node parameter?

Copy link
Contributor Author

@LCAIZJ LCAIZJ Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that's a good idea. For example "alloc_in_same_node": true.

for value in res:
if value < 0:
logger.error(f"Failed to put key {keys},res:{res}")
Expand Down
Loading