From ceb1e496610c19352aebf67b5231bf36b0d453e9 Mon Sep 17 00:00:00 2001
From: pz1116 <47019764+Pz1116@users.noreply.github.com>
Date: Fri, 17 Apr 2026 10:06:36 +0800
Subject: [PATCH] [BugFix][v0.18.0] fix remote KV waiting promotion in balance
 scheduler (#8280)

<!--  Thanks for sending a pull request!

BEFORE SUBMITTING, PLEASE READ
https://docs.vllm.ai/en/latest/contributing/overview.html

-->
### What this PR does / why we need it?
## Problem

In PD-disaggregated serving with `mooncake_connector` and
`VLLM_ASCEND_BALANCE_SCHEDULING=1`, requests may enter
`WAITING_FOR_REMOTE_KVS` and never be promoted back to runnable state
after remote KV transfer finishes.

The issue is in `BalanceScheduler`'s handling of
`WAITING_FOR_REMOTE_KVS` requests. The current code treats
`_update_waiting_for_remote_kv()` as if it returns a boolean readiness
flag:

```python
is_ready = self._update_waiting_for_remote_kv(request)
if is_ready:
    ...
else:
    ...
```
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such
as API, interface or other behavior changes.
Documentation-only updates are not considered user-facing changes.
-->

### How was this patch tested?
<!--
CI passed with new added/existing test.
If it was tested in a way different from regular unit tests, please
clarify how you tested step by step, ideally copy and paste-able, so
that other reviewers can test and check, and descendants can verify in
the future.
If tests were not added, please describe why they were not added and/or
why it was difficult to add.
-->

Signed-off-by: Pz1116 <zpbzpb123123@gmail.com>
---
 .../patch/platform/patch_balance_schedule.py    | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/vllm_ascend/patch/platform/patch_balance_schedule.py b/vllm_ascend/patch/platform/patch_balance_schedule.py
index a8a70300..9a2cc722 100644
--- a/vllm_ascend/patch/platform/patch_balance_schedule.py
+++ b/vllm_ascend/patch/platform/patch_balance_schedule.py
@@ -275,15 +275,7 @@ class BalanceScheduler(Scheduler):
 
                 # KVTransfer: skip request if still waiting for remote kvs.
                 if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
-                    is_ready = self._update_waiting_for_remote_kv(request)
-                    if is_ready:
-                        if request.num_preemptions:
-                            # We must be loading for a resumed preemption
-                            # rather than a new request.
-                            request.status = RequestStatus.PREEMPTED
-                        else:
-                            request.status = RequestStatus.WAITING
-                    else:
+                    if request.request_id not in self.finished_recving_kv_req_ids:
                         logger.debug(
                             "%s is still in WAITING_FOR_REMOTE_KVS state.",
                             request_id,
@@ -291,6 +283,13 @@ class BalanceScheduler(Scheduler):
                         self.waiting.pop_request()
                         skipped_waiting_requests.prepend_request(request)
                         continue
+                    self._update_waiting_for_remote_kv(request)
+                    if request.num_preemptions:
+                        # We must be loading for a resumed preemption
+                        # rather than a new request.
+                        request.status = RequestStatus.PREEMPTED
+                    else:
+                        request.status = RequestStatus.WAITING
 
                 # Skip request if the structured output request is still waiting
                 # for FSM compilation.