[BugFix][Worker] Preserve decode graph for PD recompute - #17032
iKeybot-code wants to merge 1 commit into
Conversation
|
👋 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! |
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 addresses a performance issue in MRV2 where PD consumer recomputation of the final prompt token is incorrectly classified as prefill. By explicitly tracking these requests and reclassifying the first local recomputation step, the system can maintain uniform-decode graph eligibility for mixed batches, significantly improving throughput and latency for migrated requests. 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
|
![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][BugFix] Handle PD decode tail recomputation as decode to prevent disabling FULL_DECODE_ONLY replaySuggested PR Summary:
### What this PR does / why we need it?
This PR addresses an issue where a PD (Prefill-Decode) decode consumer recomputes the last prompt token after installing the transferred hybrid state. Since upstream `ModelRunnerV2` derives `is_prefilling` solely from the prompt boundary, this one-token recomputation incorrectly classifies a mixed decode batch as a prefill batch, disabling `FULL_DECODE_ONLY` replay.
To resolve this, the PR:
1. Tracks requests undergoing PD decode tail-token recomputation in `MooncakeConnector` and `MooncakeScheduler`.
2. Clears these tracked request IDs upon request completion or preemption to prevent leakage.
3. Updates `NPUModelRunner.gather_batch_req_state` to reclassify these tail-recomputation requests as decode rather than prefill, maintaining a uniform decode batch and keeping `FULL_DECODE_ONLY` replay enabled.
4. Adds comprehensive unit tests for both the connector and the model runner.
Additionally, a review comment suggests improving `_get_kv_transfer_req_ids` in `vllm_ascend/worker/v2/model_runner.py` to safely handle cases where `metadata` attributes (like `metadata` or `pd_decode_recompute_req_ids`) might be explicitly set to `None`, preventing potential `TypeError` exceptions.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Tested with new unit tests added in `tests/ut/kv_offload/test_mooncake_connector.py` and `tests/ut/worker/test_model_runner_v2.py`.| def _get_kv_transfer_req_ids(metadata: object | None) -> set[str]: | ||
| """Return requests explicitly marked for PD tail recomputation.""" | ||
| if metadata is None: | ||
| return set() | ||
|
|
||
| req_ids = set(getattr(metadata, "pd_decode_recompute_req_ids", ())) | ||
| for child_metadata in getattr(metadata, "metadata", ()): | ||
| req_ids.update(_get_kv_transfer_req_ids(child_metadata)) | ||
| return req_ids |
There was a problem hiding this comment.
In _get_kv_transfer_req_ids, calling getattr(metadata, "metadata", ()) can return None if the metadata attribute is explicitly set to None on the object (rather than being absent). This will cause a TypeError: 'NoneType' object is not iterable when the loop attempts to iterate over it. Similarly, pd_decode_recompute_req_ids could also be None in some metadata implementations, which would cause set(None) to raise a TypeError.
Using getattr(..., None) or () is a safer pattern that guarantees an iterable fallback.
| def _get_kv_transfer_req_ids(metadata: object | None) -> set[str]: | |
| """Return requests explicitly marked for PD tail recomputation.""" | |
| if metadata is None: | |
| return set() | |
| req_ids = set(getattr(metadata, "pd_decode_recompute_req_ids", ())) | |
| for child_metadata in getattr(metadata, "metadata", ()): | |
| req_ids.update(_get_kv_transfer_req_ids(child_metadata)) | |
| return req_ids | |
| def _get_kv_transfer_req_ids(metadata: object | None) -> set[str]: | |
| """Return requests explicitly marked for PD tail recomputation.""" | |
| if metadata is None: | |
| return set() | |
| req_ids = set(getattr(metadata, "pd_decode_recompute_req_ids", None) or ()) | |
| for child_metadata in getattr(metadata, "metadata", None) or (): | |
| req_ids.update(_get_kv_transfer_req_ids(child_metadata)) | |
| return req_ids |
173a1e5 to
9fbc96e
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
34d2e50 to
757758d
Compare
757758d to
3633762
Compare
Track Mooncake remote-prefill requests across the no-forward receive step and mark only their first local recompute as decode-like. This keeps mixed decode batches eligible for uniform decode graphs without misclassifying ordinary final prefill tokens. Add lifecycle cleanup and regression coverage for mixed batches, multi-token decode queries, nested connector metadata, aborted requests, and reused request IDs. Signed-off-by: likailong <likailong5@huawei.com>
3633762 to
3c21c0f
Compare
…17128) ## Summary When PD decode uses `recompute_scheduler_enable=true`, the last prompt-token recompute is reported as prefill by the upstream batch state. This keeps `has_prefill` true and makes a mixed decode batch miss `FULL_DECODE_ONLY` graphs. Reclassify only the prompt-boundary recompute on a PD consumer, then refresh `has_prefill` and the uniform decode token count. This path does not require MTP and does not depend on #17032. | Validation | Baseline | Fixed | |---|---:|---:| | Targeted UT | — | 4 passed | | Qwen3-32B-W8A8 requests | 128/128 | 128/128 | | Aggregate output throughput | 121.84 tok/s | 137.59 tok/s (+12.93%) | | Mean join decode latency | 3.2335 s | 2.3539 s (-27.20%) | Full `test_model_runner_v2.py`: 40 passed; its one remaining Spec-PP failure is also reproducible on the unmodified baseline. - vLLM main: vllm-project/vllm@84030bb Signed-off-by: likailong <likailong5@huawei.com>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
…llm-project#17128) ## Summary When PD decode uses `recompute_scheduler_enable=true`, the last prompt-token recompute is reported as prefill by the upstream batch state. This keeps `has_prefill` true and makes a mixed decode batch miss `FULL_DECODE_ONLY` graphs. Reclassify only the prompt-boundary recompute on a PD consumer, then refresh `has_prefill` and the uniform decode token count. This path does not require MTP and does not depend on vllm-project#17032. | Validation | Baseline | Fixed | |---|---:|---:| | Targeted UT | — | 4 passed | | Qwen3-32B-W8A8 requests | 128/128 | 128/128 | | Aggregate output throughput | 121.84 tok/s | 137.59 tok/s (+12.93%) | | Mean join decode latency | 3.2335 s | 2.3539 s (-27.20%) | Full `test_model_runner_v2.py`: 40 passed; its one remaining Spec-PP failure is also reproducible on the unmodified baseline. - vLLM main: vllm-project/vllm@84030bb Signed-off-by: likailong <likailong5@huawei.com>
1 participant
What this PR does / why we need it?
MRV2 classifies a PD consumer's post-transfer tail-token recomputation as prefill, so
has_prefill=Truemakes mixed decode batches missFULL_DECODE_ONLYgraphs. This patch carries an explicit Mooncake request marker across the no-forward receive step and reclassifies only the first local recomputation; ordinary final prefill tokens remain unchanged. The fix is self-contained in vLLM Ascend, requires no vLLM upstream patch, and has been validated with vLLM v0.28.0 and v0.29.0.Does this PR introduce any user-facing change?
Yes. PD decode hybrid batches retain uniform-decode graph eligibility without changing APIs.
How was this patch tested?
The v0.29 run used
VLLM_USE_V2_MODEL_RUNNER=1, Mooncake,FULL_DECODE_ONLY, Qwen3-32B-W8A8, and NPU TP4+TP4. Scheduler-to-worker marker propagation and the decode reclassification branch were also confirmed on all four decode workers.