[Feature][KV Transfer] Support multiple memfabric TransDataOpType and A5 - #15570
[Feature][KV Transfer] Support multiple memfabric TransDataOpType and A5#15570dsxsteven wants to merge 1 commit into
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the KV transfer mechanism by decoupling the MemFabric data-path protocol from hardware detection. By introducing a configurable protocol setting, users can now explicitly specify the transfer mode (sdma, device_rdma, or device_urma) via the Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
github-actions
Bot
added
documentation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [Feature] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][Feature] Support MemFabric transfer protocol configuration for A5 nodesSuggested PR Summary:
### What this PR does / why we need it?
This PR adds support for A5 series nodes in layerwise and sparse KV cache offloading by introducing a configurable MemFabric transfer protocol (`memfabric_transfer_protocol`). Users can now specify `sdma` (default) or `device_rdma` for A3 nodes, and `device_urma` for A5 nodes. The PR also updates the documentation and unit tests accordingly.
Feedback on the changes includes:
- Adding type validation for `transfer_protocol` to prevent potential `AttributeError`s during string operations.
- Wrapping the import of `TransDataOpType` in a `try...except` block to gracefully handle older versions of `memfabric_hybrid` that might raise an `ImportError`.
### Does this PR introduce _any_ user-facing change?
Yes, users can now configure `memfabric_transfer_protocol` in `kv_connector_extra_config` to support A5 nodes.
### How was this patch tested?
The patch was tested using updated unit tests in `tests/ut/kv_offload/test_memfabric_transfer_engine.py`.| protocol = (transfer_protocol or _DEFAULT_MEMFABRIC_TRANSFER_PROTOCOL).strip().lower() | ||
| if protocol not in _VALID_MEMFABRIC_TRANSFER_PROTOCOLS: | ||
| raise ValueError( | ||
| f"Invalid MemFabric transfer_protocol={transfer_protocol!r}; " | ||
| f"expected one of {_VALID_MEMFABRIC_TRANSFER_PROTOCOLS}" | ||
| ) |
There was a problem hiding this comment.
To prevent unexpected runtime crashes with generic AttributeErrors, we should validate that transfer_protocol is a string (or None) before performing string operations like .strip().lower().
if transfer_protocol is not None and not isinstance(transfer_protocol, str):
raise TypeError(
f"MemFabric transfer_protocol must be a string, got {type(transfer_protocol).__name__}"
)
protocol = (transfer_protocol or _DEFAULT_MEMFABRIC_TRANSFER_PROTOCOL).strip().lower()
if protocol not in _VALID_MEMFABRIC_TRANSFER_PROTOCOLS:
raise ValueError(
f"Invalid MemFabric transfer_protocol={transfer_protocol!r}; "
f"expected one of {_VALID_MEMFABRIC_TRANSFER_PROTOCOLS}"
)| Unknown names fail fast here as well: a wrong protocol otherwise only | ||
| surfaces later as an obscure engine initialization failure. | ||
| """ | ||
| from memfabric_hybrid import TransDataOpType |
There was a problem hiding this comment.
If an older version of memfabric_hybrid is installed, importing TransDataOpType will raise an ImportError. Since _get_transfer_protocol is called outside the main try...except ImportError block in _build_engine, this will result in an unhandled exception with a generic error message. Wrapping this import in a try...except block allows us to provide a clear and helpful error message advising the user to upgrade the library.
try:
from memfabric_hybrid import TransDataOpType
except ImportError as exc:
raise ImportError(
"Please upgrade memfabric_hybrid (memfabric-hybrid) to a version that supports TransDataOpType."
) from exc
- Select data path via kv_connector_extra_config["memfabric_transfer_protocol"]: sdma (default) / device_rdma for A3 nodes, device_urma for A5 nodes; no machine-type checks in code - Fail fast on invalid protocol values; log resolved data_op_type at init - Fix UTs broken by hardcoded DEVICE_URMA and add protocol coverage - Document the new option in the layerwise/sparse KV offload guide and design doc Signed-off-by: dsxsteven <dsxsteven@sina.com>
b492f0b to
391d6b5
Compare
1 participant
What this PR does / why we need it?
kv_connector_extra_config["memfabric_transfer_protocol"]:sdma(default) /device_rdmafor A3 nodes,device_urmafor A5 nodes; no machine-type checks in codeDoes this PR introduce any user-facing change?
Yes: new
kv_connector_extra_config["memfabric_transfer_protocol"]option (sdma(default) /device_rdma/device_urma) to select the MemFabric transfer data path for A3/A5 nodes; documented in the layerwise/sparse KV offload guide and design doc.How was this patch tested?
Signed-off-by: dsxsteven dsxsteven@sina.com