feat(connector): add PegaPdConnector for vLLM PD disaggregated serving#174
feat(connector): add PegaPdConnector for vLLM PD disaggregated serving#174wz1qqx wants to merge 2 commits into
Conversation
Add PegaPdConnector that extends MultiConnector to combine PegaKVConnector
(L2 CPU/SSD cache) with NixlConnector (RDMA KV transfer) for prefill-decode
disaggregated inference with Dynamo.
Delegation logic:
- Prefix matching: PegaKVConnector only (knows the L2 cache state)
- Block allocation: PegaKVConnector gets real blocks, NixlConnector empty
- Handshake: NixlConnector only (PegaFlow has no P/D handshake)
- Save/load: both connectors participate via MultiConnector
Also registers PegaPdConnector in KVConnectorFactory via the vLLM plugin
entry point, so it can be used via --kv-transfer-config:
{"kv_connector": "PegaPdConnector", ...}
Tested with Kimi K2.5 (MLA, TP8) on 2P2D RoCE cluster:
- PegaFlow save/load/query_prefetch all functional
- 90% cache hit rate with approximate KV-aware routing (temperature=0)
- TTFT avg 2.16s at 0.2 req/s arrival rate
There was a problem hiding this comment.
Pull request overview
Adds a new vLLM KV connector composition (PegaPdConnector) to support Dynamo prefill/decode disaggregated serving by combining PegaFlow’s L2 KV cache management with NIXL-based P→D KV transfer, and registers it via the existing vLLM general plugin hook.
Changes:
- Register
PegaPdConnectorinKVConnectorFactorythrough thevllm.general_pluginsentry point. - Introduce
PegaPdConnector(aMultiConnector) that delegates prefix matching/allocation toPegaKVConnectorand handshake toNixlConnector.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
python/pegaflow/vllm_plugin.py |
Extends plugin registration to include PegaPdConnector lazily in KVConnectorFactory. |
python/pegaflow/connector/pd_connector.py |
Adds PegaPdConnector implementation combining PegaFlow + NIXL behaviors for PD disaggregated serving. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| Skips MultiConnector.bind_connector_metadata() to avoid double-binding | ||
| children. Calls KVConnectorBase_V1 directly for has_connector_metadata(). | ||
| """ | ||
| assert isinstance(connector_metadata, PegaPdConnectorMetadata) |
There was a problem hiding this comment.
bind_connector_metadata uses an assert for runtime type checking. Assertions can be stripped with python -O, which would allow an unexpected metadata type through and likely fail later in less obvious ways. Prefer an explicit isinstance check that raises TypeError (or ValueError) when the metadata is not PegaPdConnectorMetadata.
| assert isinstance(connector_metadata, PegaPdConnectorMetadata) | |
| if not isinstance(connector_metadata, PegaPdConnectorMetadata): | |
| raise TypeError( | |
| "Expected PegaPdConnectorMetadata for connector_metadata, " | |
| f"got {type(connector_metadata).__name__}" | |
| ) |
| assert isinstance(connector_metadata, PegaPdConnectorMetadata) | ||
| KVConnectorBase_V1.bind_connector_metadata(self, connector_metadata) | ||
| if connector_metadata.extra_async_saves: | ||
| self._extra_async_saves.update(connector_metadata.extra_async_saves) |
There was a problem hiding this comment.
for c, cm in zip(self._connectors, connector_metadata.metadata) will silently ignore extra connectors or extra metadata entries if their lengths differ. That can lead to only a subset of child connectors being bound, causing hard-to-debug behavior. Add an explicit length check (and raise) before iterating so misconfiguration is detected immediately.
| self._extra_async_saves.update(connector_metadata.extra_async_saves) | |
| self._extra_async_saves.update(connector_metadata.extra_async_saves) | |
| # Ensure we do not silently drop connectors or metadata entries. | |
| num_connectors = len(self._connectors) | |
| num_metadata = len(connector_metadata.metadata) | |
| if num_connectors != num_metadata: | |
| raise ValueError( | |
| "Mismatch between number of connectors and metadata entries: " | |
| f"{num_connectors} connector(s) vs {num_metadata} metadata entr" | |
| "ies in PegaPdConnector.bind_connector_metadata" | |
| ) |
| if len(self._connectors) != 2: | ||
| raise ValueError( | ||
| "PegaPdConnector requires exactly two connectors " | ||
| f"(got {len(self._connectors)})" | ||
| ) | ||
|
|
||
| if not isinstance(self._connectors[0], PegaKVConnector): | ||
| raise TypeError( | ||
| "Expected first connector to be PegaKVConnector, " | ||
| f"got {type(self._connectors[0]).__name__}" | ||
| ) | ||
| if not isinstance(self._connectors[1], NixlConnector): | ||
| raise TypeError( | ||
| "Expected second connector to be NixlConnector, " | ||
| f"got {type(self._connectors[1]).__name__}" |
There was a problem hiding this comment.
This implementation relies on MultiConnector internals (self._connectors, self._extra_async_saves). Because these are private attributes of an external dependency, they may change across vLLM versions and break this connector unexpectedly. If vLLM exposes public accessors/hooks for these, prefer those; otherwise consider adding a short compatibility note (or version constraint) and defensive checks (e.g., hasattr) to fail fast with a clear error.
| if len(self._connectors) != 2: | |
| raise ValueError( | |
| "PegaPdConnector requires exactly two connectors " | |
| f"(got {len(self._connectors)})" | |
| ) | |
| if not isinstance(self._connectors[0], PegaKVConnector): | |
| raise TypeError( | |
| "Expected first connector to be PegaKVConnector, " | |
| f"got {type(self._connectors[0]).__name__}" | |
| ) | |
| if not isinstance(self._connectors[1], NixlConnector): | |
| raise TypeError( | |
| "Expected second connector to be NixlConnector, " | |
| f"got {type(self._connectors[1]).__name__}" | |
| # Defensive check against changes in vLLM's MultiConnector internals. | |
| if not hasattr(self, "_connectors"): | |
| raise RuntimeError( | |
| "PegaPdConnector is incompatible with this vLLM version: " | |
| "MultiConnector does not expose a '_connectors' attribute." | |
| ) | |
| connectors = self._connectors | |
| if len(connectors) != 2: | |
| raise ValueError( | |
| "PegaPdConnector requires exactly two connectors " | |
| f"(got {len(connectors)})" | |
| ) | |
| if not isinstance(connectors[0], PegaKVConnector): | |
| raise TypeError( | |
| "Expected first connector to be PegaKVConnector, " | |
| f"got {type(connectors[0]).__name__}" | |
| ) | |
| if not isinstance(connectors[1], NixlConnector): | |
| raise TypeError( | |
| "Expected second connector to be NixlConnector, " | |
| f"got {type(connectors[1]).__name__}" |
- Remove quoted type annotations (UP037): use bare types with `from __future__ import annotations` - Add `strict=True` to `zip()` call (B905)
wz1qqx
changed the title
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Summary
PegaPdConnectorextendingMultiConnectorto combinePegaKVConnector+NixlConnectorfor Dynamo PD disaggregated inferencePegaPdConnectorinKVConnectorFactoryvia the vLLM plugin entry pointDesign
PegaPdConnectordelegates:Usage in
--kv-transfer-config:{ "kv_connector": "PegaPdConnector", "kv_connector_module_path": "pegaflow.connector.pd_connector", "kv_connector_extra_config": { "connectors": [ {"kv_connector": "PegaKVConnector", "kv_role": "kv_both"}, {"kv_connector": "NixlConnector", "kv_role": "kv_both"} ] } }Test plan
🤖 Generated with Claude Code