[v0.26.0rc][BugFix][P/D] Backport async KV-delivery preemption fixes - #14723
[v0.26.0rc][BugFix][P/D] Backport async KV-delivery preemption fixes#14723DavidJiang9 wants to merge 2 commits into
Conversation
Signed-off-by: jiangkaiqiang <jiangkaiqiang@huawei.com>
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 backports critical correctness fixes from upstream vLLM regarding async KV-delivery preemption. The changes ensure that in-flight output accounting remains consistent during preemption and prevents the delivery of outputs whose KV blocks have been released, which is vital for maintaining stability in Mooncake KV transfer configurations using async scheduling. 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
|
|
👋 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 [BugFix] 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:\n\nmarkdown\n[Ops][Feature] Backport vLLM #48245 and #50297 stale-output handling and KV delivery preemption\n\n\nSuggested PR Summary:\n\nmarkdown\n### What this PR does / why we need it?\nThis PR backports the stale-output handling from vLLM #48245 and the KV delivery preemption contract from vLLM #50297. It introduces `patch_async_scheduler.py` to handle stale outputs losslessly in async scheduling, updates `patch_balance_schedule.py` with the new preemption and scheduling logic, and adds `requires_kv_delivery` properties to the connectors.\n\nFeedback from the review suggests:\n- Adding defensive programming checks using `getattr` when accessing `_kv_transfer_config` and `requires_kv_delivery` on connectors to prevent potential `AttributeError`s.\n- Forwarding the `is_stale` argument in the `super` call within `_update_request_with_output` in `patch_async_scheduler.py` to ensure correct inheritance behavior.\n\n### Does this PR introduce _any_ user-facing change?\nNo.\n\n### How was this patch tested?\nTested with new CPU regression tests in `tests/ut/patch/platform/test_kv_delivery_preemption.py` and updated AST checks in `tests/ut/patch/platform/test_patch_balance_schedule.py`.\n
| def _requires_kv_delivery(connector: KVConnectorBase_V1) -> bool: | ||
| """Match the producer-role default introduced by vLLM #50297.""" | ||
| return connector._kv_transfer_config.is_kv_producer |
There was a problem hiding this comment.
Defensive programming check: If connector does not have the _kv_transfer_config attribute, or if it is None, accessing connector._kv_transfer_config.is_kv_producer will raise an AttributeError and crash the scheduler during initialization. Using getattr with a default value prevents this potential runtime crash.
| def _requires_kv_delivery(connector: KVConnectorBase_V1) -> bool: | |
| """Match the producer-role default introduced by vLLM #50297.""" | |
| return connector._kv_transfer_config.is_kv_producer | |
| def _requires_kv_delivery(connector: KVConnectorBase_V1) -> bool: | |
| """Match the producer-role default introduced by vLLM #50297.""" | |
| config = getattr(connector, "_kv_transfer_config", None) | |
| return config.is_kv_producer if config is not None else False |
| def _multi_requires_kv_delivery(connector: MultiConnector) -> bool: | ||
| """A MultiConnector must preserve every reliable child hand-off.""" | ||
| return any(child.requires_kv_delivery for child in connector._connectors) |
There was a problem hiding this comment.
Defensive programming check: To prevent potential AttributeError if any child connector in _connectors does not implement or have the requires_kv_delivery attribute, use getattr with a default of False.
| def _multi_requires_kv_delivery(connector: MultiConnector) -> bool: | |
| """A MultiConnector must preserve every reliable child hand-off.""" | |
| return any(child.requires_kv_delivery for child in connector._connectors) | |
| def _multi_requires_kv_delivery(connector: MultiConnector) -> bool: | |
| """A MultiConnector must preserve every reliable child hand-off.""" | |
| return any(getattr(child, "requires_kv_delivery", False) for child in connector._connectors) |
| self._use_consumer_partial_group_hits = not bool( | ||
| kv_transfer_config is not None and kv_transfer_config.is_kv_producer | ||
| ) | ||
| self.requires_kv_delivery = bool(self.connector is not None and self.connector.requires_kv_delivery) |
There was a problem hiding this comment.
Defensive programming check: Use getattr to safely retrieve requires_kv_delivery from self.connector to avoid potential AttributeError if a custom or third-party connector is used that does not implement this property.
| self.requires_kv_delivery = bool(self.connector is not None and self.connector.requires_kv_delivery) | |
| self.requires_kv_delivery = bool(self.connector is not None and getattr(self.connector, "requires_kv_delivery", False)) |
| is_stale: bool = False, | ||
| ) -> tuple[list[int], bool]: | ||
| status_before_update = request.status | ||
| new_token_ids, stopped = super(AsyncScheduler, self)._update_request_with_output(request, new_token_ids) |
There was a problem hiding this comment.
For robustness and correctness of the inheritance chain, the overridden _update_request_with_output method should forward the is_stale argument to the super call, ensuring any parent class or mixin that relies on this parameter receives the correct state.
| new_token_ids, stopped = super(AsyncScheduler, self)._update_request_with_output(request, new_token_ids) | |
| new_token_ids, stopped = super(AsyncScheduler, self)._update_request_with_output( | |
| request, new_token_ids, is_stale=is_stale | |
| ) |
Signed-off-by: jiangkaiqiang <jiangkaiqiang@huawei.com>
3 participants
Purpose
Backport the async-preemption correctness fixes from upstream vLLM #48245 and #50297 to the v0.26 RC branch.
This fixes stale in-flight output accounting after preemption and prevents P/D producers from delivering outputs whose KV blocks have already been released under memory pressure. The change is needed for Mooncake KV transfer configurations using async scheduling.
Changes
requires_kv_deliverycontract for KV connectors and aggregate it throughMultiConnector;Validation
git diff --check origin/releases/v0.26.0rc...HEAD— passed.python -m pytest -q tests/ut/patch/platform/test_kv_delivery_preemption.py tests/ut/patch/platform/test_patch_balance_schedule.pyvllmis not installed;PYTHONPATHalso stopped during collection because the installed Torch does not exposetorch.library.infer_schema.Upstream