Skip to content

[Bugfix] Set KVTransferConfig.engine_id in post_init #18576

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

Merged
merged 3 commits into from
May 23, 2025
Merged

[Bugfix] Set KVTransferConfig.engine_id in post_init #18576

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b5880c9
add test
lk-chen May 22, 2025
c004c60
fix
lk-chen May 23, 2025
51096a0
add assertion
lk-chen May 23, 2025
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
8 changes: 8 additions & 0 deletions tests/v1/kv_connector/unit/test_multi_connector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,3 +239,11 @@ def get_connector_events() -> dict[str, list[str]]:
print(f"[ERROR] Could not read connector events for {name}: {e}")

return connector_events


def test_engine_id_conflict():
configs = [KVTransferConfig() for _ in range(2)]
ids = [config.engine_id for config in configs]
assert ids[0] != ids[1], (
"Engine IDs should be different for different configs. "
f"Got {ids}")
5 changes: 4 additions & 1 deletion vllm/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3492,7 +3492,7 @@ class KVTransferConfig:
"""The KV connector for vLLM to transmit KV caches between vLLM instances.
"""

engine_id: str = str(uuid.uuid4())
engine_id: Optional[str] = None
"""The engine id for KV transfers."""

kv_buffer_device: Optional[str] = "cuda"
Expand Down Expand Up @@ -3549,6 +3549,9 @@ def compute_hash(self) -> str:
return hash_str

def __post_init__(self) -> None:
if self.engine_id is None:
self.engine_id = str(uuid.uuid4())
Comment on lines +3552 to +3553
Copy link
Member

@tlrmchlsmth tlrmchlsmth May 23, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

If we create the engine_id here, will different DP ranks end up with different values? Currently all ranks end up with the same ID, which causes issues (cc @wseaton)

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Different DP should trigger this post_init separately which should have different value at the end.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

After second thought, I think it's tricky, depends on how you generate EngineArgs for different DP

case 1

for _ in range(dp_size):
  eng_args = EngArgs(
    kv_transfer_config=KVTransferConfig(
      ..
    )
    ..
)

this will have different id

case 2

eng_args = EngArgs(..)
for _ in range(dp_size):
  copy = deepcopy(eng_args)
  # modify dp rank etc. here

this will have identical id

So if you explicitly want different id, plz explicitly set it


if self.kv_role is not None and self.kv_role not in get_args(KVRole):
raise ValueError(f"Unsupported kv_role: {self.kv_role}. "
f"Supported roles are {get_args(KVRole)}")
Expand Down
1 change: 1 addition & 0 deletions vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -537,6 +537,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

def add_remote_agent(self, nixl_agent_meta: NixlAgentMetadata):
engine_id = nixl_agent_meta.engine_id
assert engine_id != self.engine_id, "Conflict engine id found!"
if engine_id in self._remote_agents:
return

Expand Down