[v0.26.0rc][BugFix][P/D] Backport async KV-delivery preemption fixes

tests/ut/patch/platform/test_kv_delivery_preemption.py

@@ -0,0 +1,214 @@
+"""CPU regression tests adapted from vLLM PRs #48245 and #50297.
+
+The upstream tests use vLLM's in-tree ``tests/v1/core`` helpers, which are not
+shipped in its wheel. These tests cover the same compatibility-boundary state
+transitions without requiring an NPU model runner.
+"""
+
+import inspect
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_ascend.patch.platform import patch_async_scheduler as async_backport
+from vllm_ascend.patch.platform import patch_balance_schedule as backport
+
+
+class _Queue:
+    def __init__(self):
+        self.requests = []
+
+    def prepend_request(self, request):
+        self.requests.insert(0, request)
+
+
+class _Request(SimpleNamespace):
+    __hash__ = object.__hash__
+
+
+def _request(
+    request_id="handoff",
+    *,
+    in_flight=4,
+    placeholders=4,
+    stale=0,
+    drop_stale=False,
+):
+    return _Request(
+        request_id=request_id,
+        status=backport.RequestStatus.RUNNING,
+        num_in_flight_tokens=in_flight,
+        num_output_placeholders=placeholders,
+        num_stale_output_tokens=stale,
+        drop_stale_output=drop_stale,
+        spec_token_ids=[],
+        num_computed_tokens=32,
+        num_preemptions=0,
+    )
+
+
+def _scheduler():
+    scheduler = backport.BalanceScheduler.__new__(backport.BalanceScheduler)
+    scheduler._free_request_blocks = lambda _request: None
+    scheduler.encoder_cache_manager = SimpleNamespace(free=lambda _request: None)
+    scheduler._inflight_prefills = set()
+    scheduler.log_stats = False
+    scheduler.waiting = _Queue()
+    scheduler.reset_preempted_req_ids = set()
+    return scheduler
+
+
+@pytest.mark.parametrize(
+    ("is_producer", "expected"),
+    [(True, True), (False, False)],
+)
+def test_requires_kv_delivery_defaults_to_producer_role(is_producer, expected):
+    connector = SimpleNamespace(_kv_transfer_config=SimpleNamespace(is_kv_producer=is_producer))
+
+    assert backport._requires_kv_delivery(connector) is expected
+
+
+@pytest.mark.parametrize(
+    ("children", "expected"),
+    [
+        ([False, False], False),
+        ([False, True], True),
+        ([True, False], True),
+        ([], False),
+    ],
+)
+def test_multi_connector_aggregates_child_delivery_requirements(children, expected):
+    connector = SimpleNamespace(_connectors=[SimpleNamespace(requires_kv_delivery=value) for value in children])
+
+    assert backport._multi_requires_kv_delivery(connector) is expected
+
+
+def test_best_effort_offload_cache_does_not_require_delivery():
+    assert backport._best_effort_cache_requires_kv_delivery(SimpleNamespace()) is False
+
+
+def test_request_stale_output_state_has_neutral_defaults():
+    assert backport.Request.num_stale_output_tokens == 0
+    assert backport.Request.drop_stale_output is False
+
+
+def test_async_scheduler_inherits_existing_balance_patch():
+    assert issubclass(async_backport.AsyncScheduler, backport.BalanceScheduler)
+
+
+@pytest.mark.parametrize("drop_stale_output", [False, True])
+def test_preempt_marks_all_inflight_output_stale(drop_stale_output):
+    """#48245 records the complete in-flight token share, not one frame."""
+    request = _request(in_flight=12, placeholders=4)
+    scheduler = _scheduler()
+
+    backport.BalanceScheduler._preempt_request(
+        scheduler,
+        request,
+        0.0,
+        drop_stale_output=drop_stale_output,
+    )
+
+    assert request.status == backport.RequestStatus.PREEMPTED
+    assert request.num_stale_output_tokens == 12
+    assert request.num_output_placeholders == 0
+    assert request.drop_stale_output is drop_stale_output
+    assert scheduler.waiting.requests == [request]
+
+
+def test_repreempt_keeps_an_undrained_drop_share_dropped():
+    request = _request(in_flight=8, stale=4, drop_stale=True)
+    scheduler = _scheduler()
+
+    backport.BalanceScheduler._preempt_request(scheduler, request, 0.0)
+
+    assert request.num_stale_output_tokens == 8
+    assert request.drop_stale_output is True
+
+
+def test_reset_prefix_cache_uses_drop_mode_for_same_step_resume():
+    """#48245 replaces v0.26's placeholder-count discard bookkeeping."""
+    request = _request(in_flight=8, placeholders=4)
+    scheduler = _scheduler()
+    scheduler.running = [request]
+    scheduler.prev_step_scheduled_req_ids = {request.request_id}
+    scheduler.kv_cache_manager = SimpleNamespace(reset_prefix_cache=lambda: True)
+    scheduler.connector = None
+
+    assert backport.BalanceScheduler.reset_prefix_cache(scheduler, reset_running_requests=True) is True
+
+    assert request.num_stale_output_tokens == 8
+    assert request.num_output_placeholders == 0
+    assert request.drop_stale_output is True
+    assert scheduler.prev_step_scheduled_req_ids == set()
+
+
+@pytest.mark.parametrize(
+    ("is_stale", "expected_placeholders"),
+    [(True, 0), (False, 2)],
+)
+def test_async_placeholder_update_skips_stale_delivery(monkeypatch, is_stale, expected_placeholders):
+    """#48245 stale output is delivered without decrementing reset counters."""
+    request = _request(placeholders=4)
+    scheduler = async_backport.AsyncScheduler.__new__(async_backport.AsyncScheduler)
+    scheduler.kv_cache_manager = SimpleNamespace(cache_blocks=lambda *_args: None)
+
+    def update_original(_scheduler, _request, new_token_ids, is_stale=False):
+        return new_token_ids, False
+
+    monkeypatch.setattr(backport.BalanceScheduler, "_update_request_with_output", update_original)
+    if is_stale:
+        request.num_output_placeholders = 0
+
+    new_token_ids, stopped = async_backport._update_request_with_output(
+        scheduler,
+        request,
+        [10, 11],
+        is_stale=is_stale,
+    )
+
+    assert new_token_ids == [10, 11]
+    assert stopped is False
+    assert request.num_output_placeholders == expected_placeholders
+
+
+def test_50297_pressure_preemption_uses_connector_delivery_requirement():
+    """The existing schedule copy carries #50297 at the original call site."""
+    source = inspect.getsource(backport.BalanceScheduler.schedule)
+
+    assert "drop_stale_output=self.requires_kv_delivery" in source
+
+
+def test_48245_waits_for_deliverable_stale_output_before_resume():
+    source = inspect.getsource(backport.BalanceScheduler.schedule)
+
+    assert "request.num_stale_output_tokens > 0" in source
+    assert "not request.drop_stale_output" in source
+
+
+def test_48245_preemption_keeps_upstream_operation_order():
+    source = inspect.getsource(backport.BalanceScheduler._preempt_request)
+    markers = [
+        "self._free_request_blocks(request)",
+        "request.num_computed_tokens = 0",
+        "request.drop_stale_output =",
+        "request.num_stale_output_tokens = request.num_in_flight_tokens",
+        "request.num_preemptions += 1",
+        "self.waiting.prepend_request(request)",
+    ]
+
+    assert [source.index(marker) for marker in markers] == sorted(source.index(marker) for marker in markers)
+
+
+def test_48245_output_update_keeps_stale_checks_in_original_flow():
+    source = inspect.getsource(inspect.unwrap(backport.BalanceScheduler.update_from_output))
+    markers = [
+        "request.num_in_flight_tokens -= num_tokens_scheduled",
+        "request.num_stale_output_tokens -= num_tokens_scheduled",
+        "if failed_kv_load_req_ids",
+        "if request is None or request.is_finished()",
+        "if output_is_stale and request.drop_stale_output",
+        "req_index = model_runner_output.req_id_to_index[req_id]",
+    ]
+
+    assert [source.index(marker) for marker in markers] == sorted(source.index(marker) for marker in markers)

tests/ut/patch/platform/test_patch_balance_schedule.py

@@ -20,11 +20,11 @@
 * the module-level class swaps actually took effect;
 * the upstream Scheduler/DPEngineCoreProc methods the patch calls/super-calls
   still exist;
-* the 4 platform deltas remain present in ``schedule()`` (intent lock);
+* the 5 platform deltas remain present in ``schedule()`` (intent lock);
 * the copied ``schedule()`` body stays a verbatim copy of the ``schedule()``
   at vllm-ascend's pinned vLLM release tag (read from
   ``.github/vllm-release-tag.commit`` -- the same file CI uses), modulo exactly
-  those 4 deltas. Reading the tag from the pin file means a pin advance
+  those 5 deltas. Reading the tag from the pin file means a pin advance
   auto-flips this guard to the new tag until the copy is re-synced.
 
 What is NOT guarded here (structurally unreachable without a real engine):
@@ -249,13 +249,13 @@ def fake_scheduler_init(self, *args, **kwargs):
 
 
 def _schedule_body_ast(source: str) -> str:
-    """Canonical AST dump of a ``schedule`` method body with the 4 platform
+    """Canonical AST dump of a ``schedule`` method body with the 5 platform
     deltas stripped, so the remainder can be compared verbatim against the
     pinned release tag's ``schedule()``. AST-based on purpose: it is blind to
     comments and whitespace, so the only differences that surface are real code
     drift (not the escape-quoting of a comment or reformatting).
 
-    The 3 deltas removed:
+    The 5 deltas removed:
       * delta 1 -- the disabled-path early return (``if not
         self._balance_enabled: ... super().schedule(...)``);
       * delta 2 -- the ``balance_flag`` gate (``max(t.item() for t in
@@ -265,6 +265,8 @@ def _schedule_body_ast(source: str) -> str:
         ``assert request_queue is not None``. Both are stripped so the two
         bodies align.
       * delta 4 -- the consumer-only partial-group cache lookup gate.
+      * delta 5 -- the #48245/#50297 stale-output admission gate and
+        reliable-delivery preemption flag.
     """
     tree = ast.parse(textwrap.dedent(source))
     func = next(
@@ -274,6 +276,12 @@ def _schedule_body_ast(source: str) -> str:
     assert func is not None, "no schedule() in source"
 
     class _BalanceDeltaStripper(ast.NodeTransformer):
+        def visit_Call(self, node: ast.Call):  # noqa: N802
+            self.generic_visit(node)
+            if isinstance(node.func, ast.Attribute) and node.func.attr == "_preempt_request":
+                node.keywords = [kw for kw in node.keywords if kw.arg != "drop_stale_output"]
+            return node
+
         def visit_BoolOp(self, node: ast.BoolOp):  # noqa: N802
             self.generic_visit(node)
             node.values = [value for value in node.values if "_use_consumer_partial_group_hits" not in ast.dump(value)]
@@ -290,6 +298,9 @@ def visit_If(self, node: ast.If):  # noqa: N802
             # delta 3 (ours): if request_queue is None: break.
             if "request_queue" in test and "None" in test and "Is" in test:
                 return None
+            # delta 5: wait for deliverable stale output before resuming.
+            if "num_stale_output_tokens" in test and "drop_stale_output" in test:
+                return None
             return self.generic_visit(node)
 
         def visit_Assert(self, node: ast.Assert):  # noqa: N802
@@ -366,13 +377,13 @@ def test_schedule_signature_matches_installed_vllm():
 
 
 # ---------------------------------------------------------------------------
-# 1b. the 3 balance deltas remain present in schedule() (intent lock)
+# 1b. the platform deltas remain present in schedule() (intent lock)
 # ---------------------------------------------------------------------------
 
 
 def test_balance_deltas_present_in_schedule():
     """The whole point of copying schedule() is to inject the balance logic.
-    If a future re-sync against the pinned tag drops any of the 3 deltas,
+    If a future re-sync against the pinned tag drops any of the 5 deltas,
     balance silently stops working -- this locks their presence in the source."""
     src = inspect.getsource(BalanceScheduler.schedule)
 
@@ -390,6 +401,10 @@ def test_balance_deltas_present_in_schedule():
     # delta 4: producer schedulers never use the consumer-only per-group lookup.
     assert "self._use_consumer_partial_group_hits" in src
 
+    # delta 5: stale output is drained or dropped according to connector need.
+    assert "request.num_stale_output_tokens" in src
+    assert "drop_stale_output=self.requires_kv_delivery" in src
+
 
 # ---------------------------------------------------------------------------
 # 1c. copied schedule() body stays verbatim with the pinned release tag
@@ -399,7 +414,7 @@ def test_balance_deltas_present_in_schedule():
 def test_schedule_body_matches_pinned_release_tag():
     """The copied ``schedule()`` body must stay a verbatim copy of the
     ``schedule()`` at vllm-ascend's pinned vLLM release tag, modulo exactly the
-    3 balance deltas.
+    5 platform deltas.
 
     The tag is read dynamically from ``.github/vllm-release-tag.commit`` -- the
     same file CI uses to pick the tag, NOT a hardcoded string or a design doc
@@ -422,9 +437,10 @@ def test_schedule_body_matches_pinned_release_tag():
     theirs = _schedule_body_ast(pinned_src)
     assert ours == theirs, (
         f"BalanceScheduler.schedule body drifted from the pinned release tag "
-        f"({tag}) beyond the 3 balance deltas. Re-sync the copy against "
+        f"({tag}) beyond the 5 platform deltas. Re-sync the copy against "
         f"{tag} and re-apply only: (1) disabled-path early return, "
-        f"(2) balance_flag gate, (3) if request_queue is None: break."
+        f"(2) balance_flag gate, (3) if request_queue is None: break, "
+        f"(4) producer partial-group lookup gate, (5) stale-output handling."
     )
 
 
@@ -513,6 +529,8 @@ def test_module_level_swaps_take_effect():
 _SCHEDULER_METHOD_SEAMS = [
     "schedule",  # super().schedule() on the disabled path
     "_preempt_request",
+    "update_from_output",
+    "reset_prefix_cache",
     "_try_schedule_encoder_inputs",
     "_mamba_block_aligned_split",
     "_select_waiting_queue_for_scheduling",

vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py

@@ -1497,6 +1497,7 @@ def add_new_req(
 class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
     def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole, kv_cache_config: KVCacheConfig | None = None):
         assert vllm_config.kv_transfer_config is not None
+        self._kv_transfer_config = vllm_config.kv_transfer_config
         self.engine_id = vllm_config.kv_transfer_config.engine_id
         self._connector_metadata = MooncakeConnectorMetadata()
 

vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py

@@ -1081,6 +1081,7 @@ def add_new_req(
 class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
     def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole, kv_cache_config: KVCacheConfig | None = None):
         assert vllm_config.kv_transfer_config is not None
+        self._kv_transfer_config = vllm_config.kv_transfer_config
         self.engine_id = vllm_config.kv_transfer_config.engine_id
         self._connector_metadata = MooncakeConnectorMetadata()
 

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/ascend_store_connector.py

@@ -78,6 +78,11 @@ def __repr__(self) -> str:
 
 
 class AscendStoreConnector(KVConnectorBase_V1, SupportsHMA):
+    @property
+    def requires_kv_delivery(self) -> bool:
+        # AscendStore is a best-effort cache: a dropped save is a future miss.
+        return False
+
     @classmethod
     def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
         """

vllm_ascend/patch/__init__.py

@@ -32,7 +32,7 @@
 # =================
 # Entries are listed in alphabetical order by file name.
 #
-# ** 1. File: platform/patch_balance_schedule.py**
+# ** 1. Files: platform/patch_async_scheduler.py, platform/patch_balance_schedule.py**
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #   1. `vllm.v1.engine.core.EngineCoreProc.run_engine_core`
 #      `vllm.v1.core.sched.scheduler.Scheduler`
@@ -48,6 +48,18 @@
 #    Future Plan:
 #       Remove this patch when vLLM merge the PR.
 #
+#   2. `vllm.v1.core.sched.async_scheduler.AsyncScheduler._update_request_with_output`
+#    Why:
+#       vLLM #48245 adds lossless stale-output handling for async scheduling.
+#       AsyncScheduler owns placeholder accounting, so this method cannot
+#       inherit the implementation from the scheduler patch.
+#    How:
+#       Replace only `_update_request_with_output` with the #48245 version.
+#    Related PR (if no, explain why):
+#       https://github.com/vllm-project/vllm/pull/48245
+#    Future Plan:
+#       Remove this patch when the supported vLLM version includes PR #48245.
+#
 # ** 2. File: platform/patch_camem_allocator.py**
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #   1. `vllm.config.model.is_cumem_allocator_available`