[bug-fix] Fix KV block leak with fault tolerance enabled in P/D disaggregation (#56430)

Signed-off-by: tiankaijin <tiankaijin1@h-partners.com>
Signed-off-by: tiankaijin <1669619752@qq.com>
Co-authored-by: tiankaijin <tiankaijin1@h-partners.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Pioneer-Weirdo <80094616+Pioneer-Weirdo@users.noreply.github.com>

.buildkite/test_areas/fault_tolerance.yaml

@@ -14,7 +14,9 @@ steps:
   - vllm/entrypoints/serve/fault_tolerance/
   - vllm/distributed/elastic_ep/
   - vllm/distributed/device_communicators/
+  - vllm/distributed/kv_transfer/
   - vllm/v1/engine/
+  - vllm/v1/executor/
   - vllm/v1/worker/
   - tests/v1/fault_tolerance/
   - tests/v1/distributed/test_external_lb_dp.py
@@ -23,6 +25,7 @@ steps:
   - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
+  - pytest -v -s v1/fault_tolerance/test_engine_core_sentinel.py
   - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
   mirror:
     amd:

tests/v1/executor/test_executor.py

@@ -3,25 +3,31 @@
 
 import asyncio
 import os
+from collections import deque
 from collections.abc import Callable
 from concurrent.futures import Future
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
 
 from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
-from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
+from vllm.distributed.kv_transfer.kv_connector.utils import (
+    KVConnectorOutput,
+    KVOutputAggregator,
+)
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.llm_engine import LLMEngine
 from vllm.v1.executor import multiproc_executor as multiproc_executor_module
 from vllm.v1.executor.abstract import Executor
-from vllm.v1.executor.multiproc_executor import MultiprocExecutor
+from vllm.v1.executor.multiproc_executor import MultiprocExecutor, WorkerProc
 from vllm.v1.executor.uniproc_executor import (
     ExecutorWithExternalLauncher,
     UniProcExecutor,
 )
+from vllm.v1.outputs import ModelRunnerOutput
 
 
 class Mock: ...
@@ -205,3 +211,113 @@ def test_custom_executor_async(distributed_executor_backend, tmp_path):
         assert os.path.exists(".marker")
     finally:
         os.chdir(cwd)
+
+
+class _FakeResponseMQ:
+    """Stands in for a worker response MessageQueue."""
+
+    def __init__(self, responses: list[tuple[Any, Any]]):
+        self._responses = list(responses)
+
+    def dequeue(self, timeout: float | None = None):
+        return self._responses.pop(0)
+
+
+def _make_failed_step_executor(enable_ft: bool) -> MultiprocExecutor:
+    executor = MultiprocExecutor.__new__(MultiprocExecutor)
+    executor.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(enable_fault_tolerance=enable_ft)
+    )
+    executor.is_failed = False
+    executor.rpc_broadcast_mq = SimpleNamespace(enqueue=lambda *_: None)
+    executor.response_mqs = [
+        _FakeResponseMQ(
+            [
+                (
+                    WorkerProc.ResponseStatus.FAILURE_WITH_KV_OUTPUT,
+                    KVConnectorOutput(finished_recving={"req1"}, invalid_block_ids={7}),
+                )
+            ]
+        ),
+        _FakeResponseMQ(
+            [
+                (
+                    WorkerProc.ResponseStatus.SUCCESS,
+                    ModelRunnerOutput(
+                        req_ids=[],
+                        req_id_to_index={},
+                        kv_connector_output=KVConnectorOutput(
+                            finished_recving={"req1"}
+                        ),
+                    ),
+                )
+            ]
+        ),
+    ]
+    executor.futures_queue = deque()
+    return executor
+
+
+@pytest.mark.parametrize("enable_ft", [True, False])
+def test_collective_rpc_worker_failure_merges_kv_output(enable_ft):
+    """A worker failing with KV output must not lose KV progress (FT only)."""
+    executor = _make_failed_step_executor(enable_ft)
+    aggregator = KVOutputAggregator(expected_finished_count=2)
+
+    future = executor.collective_rpc(
+        "sample_tokens", non_block=True, kv_output_aggregator=aggregator
+    )
+    with pytest.raises(RuntimeError, match="[Ww]orker"):
+        future.result()
+
+    if enable_ft:
+        assert len(aggregator._salvaged_outputs) == 2
+        aggregated = aggregator.aggregate(
+            [
+                ModelRunnerOutput(req_ids=[], req_id_to_index={}),
+                ModelRunnerOutput(req_ids=[], req_id_to_index={}),
+            ]
+        )
+        assert aggregated is not None
+        kv = aggregated.kv_connector_output
+        assert kv is not None
+        assert kv.finished_recving == {"req1"}
+        assert kv.invalid_block_ids == {7}
+    else:
+        assert aggregator._salvaged_outputs == []
+
+
+@pytest.mark.parametrize("non_block", [False, True])
+@pytest.mark.parametrize("connector", [None, "ec", "kv"])
+@pytest.mark.parametrize("second_fails", [False, True])
+def test_ft_rpc_failure_keeps_next_rpc_responses_aligned(
+    non_block, connector, second_fails
+):
+    """FT retries must consume fresh replies, including RPCs without KV output."""
+    executor = _make_failed_step_executor(enable_ft=True)
+    success = WorkerProc.ResponseStatus.SUCCESS
+    failure = WorkerProc.ResponseStatus.FAILURE
+    executor.response_mqs = [
+        _FakeResponseMQ([(failure, "device error"), (success, "recovered0")]),
+        _FakeResponseMQ(
+            [
+                (failure if second_fails else success, "previous step"),
+                (success, "recovered1"),
+            ]
+        ),
+    ]
+    kwargs = {}
+    if connector == "ec":
+        kwargs["ec_output_aggregator"] = ECOutputAggregator()
+    elif connector == "kv":
+        kwargs["kv_output_aggregator"] = KVOutputAggregator(expected_finished_count=2)
+
+    with pytest.raises(RuntimeError, match="[Ww]orker"):
+        ret = executor.collective_rpc("check_health", non_block=non_block, **kwargs)
+        if non_block:
+            ret.result()
+
+    assert executor.collective_rpc("handle_ft_command") == [
+        "recovered0",
+        "recovered1",
+    ]

tests/v1/fault_tolerance/test_engine_core_sentinel.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for EngineCoreSentinel fault handling."""
+
+from collections import deque
+from concurrent.futures import Future
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from vllm.distributed.kv_transfer.kv_connector.utils import KVConnectorOutput
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.fault_tolerance.engine_core_sentinel import EngineCoreSentinel
+from vllm.v1.outputs import ModelRunnerOutput
+
+
+def _model_output(kv_connector_output: KVConnectorOutput | None) -> ModelRunnerOutput:
+    output = ModelRunnerOutput(req_ids=[], req_id_to_index={})
+    output.kv_connector_output = kv_connector_output
+    return output
+
+
+def _done_future(result=None, exception=None) -> Future:
+    future: Future = Future()
+    if exception is not None:
+        future.set_exception(exception)
+    else:
+        future.set_result(result)
+    return future
+
+
+def _make_sentinel(batch_queue: deque) -> tuple[EngineCoreSentinel, MagicMock]:
+    scheduler = MagicMock()
+    sentinel = EngineCoreSentinel.__new__(EngineCoreSentinel)
+    sentinel.engine = SimpleNamespace(batch_queue=batch_queue, scheduler=scheduler)
+    return sentinel, scheduler
+
+
+def test_drain_batch_queue_surfaces_kv_connector_output():
+    """Pre-fault KV output must reach the scheduler, not be lost on clear."""
+    kv_output = KVConnectorOutput(finished_sending={"req1"})
+    batch_queue = deque([(_done_future(_model_output(kv_output)), None, None)])
+    sentinel, scheduler = _make_sentinel(batch_queue)
+
+    sentinel._drain_batch_queue()
+
+    assert not batch_queue
+    scheduler.update_from_output.assert_called_once()
+    sched_output, model_output = scheduler.update_from_output.call_args[0]
+    assert isinstance(sched_output, SchedulerOutput)
+    assert sched_output.num_scheduled_tokens == {}
+    assert model_output.kv_connector_output is kv_output
+
+
+def test_drain_batch_queue_drops_failures_and_empty_outputs():
+    """Failed futures and empty KV outputs are dropped without scheduler calls."""
+    batch_queue = deque(
+        [
+            (_done_future(exception=RuntimeError("boom")), None, None),
+            (_done_future(_model_output(None)), None, None),
+            (_done_future(_model_output(KVConnectorOutput())), None, None),
+        ]
+    )
+    sentinel, scheduler = _make_sentinel(batch_queue)
+
+    sentinel._drain_batch_queue()
+
+    assert not batch_queue
+    scheduler.update_from_output.assert_not_called()

tests/v1/kv_connector/unit/test_output_aggregator.py

@@ -3,12 +3,64 @@
 
 import pytest
 
+from vllm.distributed.kv_events import (
+    BlockStored,
+    KVCacheEvent,
+    KVConnectorKVEvents,
+)
 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorWorkerMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import (
+    LMCacheKVEvents,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector import (
+    MooncakeStoreKVEvents,
+)
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
 pytestmark = pytest.mark.cpu_test
 
 
+class DummyWorkerMeta(KVConnectorWorkerMetadata):
+    def __init__(self, tags: set[str]):
+        self.tags = set(tags)
+
+    def aggregate(self, other: "DummyWorkerMeta") -> "DummyWorkerMeta":
+        return DummyWorkerMeta(self.tags | other.tags)
+
+
+def _block_stored(*token_ids: int) -> BlockStored:
+    return BlockStored(
+        block_hashes=[b"\xab" * 32],
+        parent_block_hash=None,
+        token_ids=list(token_ids),
+        block_size=4,
+        lora_id=None,
+        medium="CPU",
+        lora_name=None,
+    )
+
+
+def make_lmcache_events(*token_ids: int) -> LMCacheKVEvents:
+    events = LMCacheKVEvents(num_workers=1)
+    events.add_events([_block_stored(*token_ids)])
+    return events
+
+
+def make_mooncake_events(*token_ids: int) -> MooncakeStoreKVEvents:
+    events = MooncakeStoreKVEvents(num_workers=1)
+    events.add_events([_block_stored(*token_ids)])
+    return events
+
+
+def lmcache_published(events: LMCacheKVEvents) -> list[KVCacheEvent]:
+    """Mirror LMCacheConnectorV1.take_events: aggregate, then collect."""
+    events.aggregate()
+    return events.get_all_events()
+
+
 class DummyModelRunnerOutput(ModelRunnerOutput):
     kv_connector_output: KVConnectorOutput
 
@@ -19,6 +71,8 @@ class DummyModelRunnerOutput(ModelRunnerOutput):
         invalid_block_ids: set[int] | None = None,
         failed_recving: set[str] | None = None,
         expected_finished_count: int = 0,
+        kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None,
+        kv_cache_events: KVConnectorKVEvents | None = None,
     ):
         self.kv_connector_output = KVConnectorOutput(
             finished_sending=finished_sending,
@@ -26,6 +80,8 @@ class DummyModelRunnerOutput(ModelRunnerOutput):
             invalid_block_ids=invalid_block_ids or set(),
             failed_recving=failed_recving or set(),
             expected_finished_count=expected_finished_count,
+            kv_connector_worker_meta=kv_connector_worker_meta,
+            kv_cache_events=kv_cache_events,
         )
 
     def __repr__(self):
@@ -138,3 +194,157 @@ def test_aggregate_workers_output_with_expected_finished_count():
     # NOTE: This is to showcase dynamic update. Workers are responsible for
     # ensuring "req1" termination in this case
     assert aggregator._send_remaining_count["req1"] == 2
+
+
+def test_replay_preserves_get_and_clear_fields():
+    """Salvaged outputs must be replayed and their fields surfaced later."""
+    aggregator = KVOutputAggregator(expected_finished_count=4)
+
+    # Failed step: worker 0's output is salvaged, worker 1 is lost.
+    salvaged = KVConnectorOutput(
+        finished_sending={"req1"},
+        finished_recving={"req2"},
+        failed_recving={"req3"},
+        invalid_block_ids={7},
+        kv_connector_worker_meta=DummyWorkerMeta({"salvaged"}),
+        kv_cache_events=make_lmcache_events(1, 2, 3, 4),
+        expected_finished_count=2,
+    )
+    aggregator.ft_defer_kv_connector_outputs([salvaged, None])
+
+    # Recovered step: both workers succeed.
+    output0 = DummyModelRunnerOutput(
+        finished_recving={"req3"},
+        kv_connector_worker_meta=DummyWorkerMeta({"recovered"}),
+        kv_cache_events=make_lmcache_events(9, 10, 11, 12),
+    )
+    output1 = DummyModelRunnerOutput(
+        finished_sending={"req1"},
+        finished_recving={"req2", "req3"},
+    )
+    aggregated = aggregator.aggregate([output0, output1])
+    kv = aggregated.kv_connector_output
+
+    assert kv.expected_finished_count == 2
+    assert kv.finished_sending == {"req1"}
+    assert kv.finished_recving == {"req2", "req3"}
+    assert kv.failed_recving == {"req3"}
+    assert kv.invalid_block_ids == {7}
+    assert kv.kv_connector_worker_meta.tags == {"salvaged", "recovered"}
+    assert kv.kv_cache_events is not None
+    all_token_ids = sorted(
+        tuple(e.token_ids) for e in kv.kv_cache_events.get_all_events()
+    )
+    assert all_token_ids == [(9, 10, 11, 12)]
+
+
+def test_replay_completes_finished_recving_votes():
+    """Salvaged recv votes must complete via replay, or the request waits
+    on WAITING_FOR_REMOTE_KVS forever."""
+    aggregator = KVOutputAggregator(expected_finished_count=2)
+
+    # Failed step: req1's recv votes completed (2/2), flagged as failed.
+    aggregator.ft_defer_kv_connector_outputs(
+        [
+            KVConnectorOutput(finished_recving={"req1"}),
+            KVConnectorOutput(finished_recving={"req1"}, failed_recving={"req1"}),
+        ]
+    )
+
+    # Recovered step: nobody reports req1 anymore.
+    aggregated = aggregator.aggregate(
+        [DummyModelRunnerOutput(), DummyModelRunnerOutput()]
+    )
+    kv = aggregated.kv_connector_output
+    assert kv.finished_recving == {"req1"}
+    assert kv.failed_recving == {"req1"}
+
+
+def test_salvaged_kv_cache_events_do_not_raise_quorum():
+    """Salvaged events are extra votes only; quorum stays at live workers."""
+    aggregator = KVOutputAggregator(expected_finished_count=2)
+
+    # Failed step: only worker 0 reported event E0.
+    aggregator.ft_defer_kv_connector_outputs(
+        [KVConnectorOutput(kv_cache_events=make_lmcache_events(1, 2, 3, 4)), None]
+    )
+
+    # Recovered step: both workers report event E1.
+    output0 = DummyModelRunnerOutput(kv_cache_events=make_lmcache_events(5, 6, 7, 8))
+    output1 = DummyModelRunnerOutput(kv_cache_events=make_lmcache_events(5, 6, 7, 8))
+    aggregated = aggregator.aggregate([output0, output1])
+
+    events = aggregated.kv_connector_output.kv_cache_events
+    assert events is not None
+    assert events.get_number_of_workers() == 2
+    # Mirror the LMCache consumer's take_events.
+    assert [e.token_ids for e in lmcache_published(events)] == [[5, 6, 7, 8]]
+
+
+def test_salvaged_votes_dropped_live_votes_unaffected_lmcache_exact_quorum():
+    """A dropped salvaged vote must not take a live event down with it
+    (LMCache exact-match quorum)."""
+    aggregator = KVOutputAggregator(expected_finished_count=2)
+
+    # Failed step: worker 0's vote for block E is salvaged.
+    aggregator.ft_defer_kv_connector_outputs(
+        [KVConnectorOutput(kv_cache_events=make_lmcache_events(1, 2, 3, 4)), None]
+    )
+
+    # Recovered step: the surviving worker reports the same block E while
+    # the other live worker stays quiet.
+    output0 = DummyModelRunnerOutput()
+    output1 = DummyModelRunnerOutput(kv_cache_events=make_lmcache_events(1, 2, 3, 4))
+    aggregated = aggregator.aggregate([output0, output1])
+
+    events = aggregated.kv_connector_output.kv_cache_events
+    assert events is not None
+    assert len(events.get_all_events()) == 1
+    assert [e.token_ids for e in lmcache_published(events)] == [[1, 2, 3, 4]]
+
+
+def test_salvaged_votes_dropped_live_votes_unaffected_mooncake_quorum():
+    """Same as the LMCache test, under MooncakeStore's >= quorum."""
+    aggregator = KVOutputAggregator(expected_finished_count=2)
+
+    # Failed step: worker 0's vote for block E is salvaged.
+    aggregator.ft_defer_kv_connector_outputs(
+        [KVConnectorOutput(kv_cache_events=make_mooncake_events(1, 2, 3, 4)), None]
+    )
+
+    # Recovered step: the surviving worker reports the same block E.
+    output0 = DummyModelRunnerOutput()
+    output1 = DummyModelRunnerOutput(kv_cache_events=make_mooncake_events(1, 2, 3, 4))
+    aggregated = aggregator.aggregate([output0, output1])
+
+    events = aggregated.kv_connector_output.kv_cache_events
+    assert events is not None
+    # Mirror the MooncakeStore consumer (take_events -> pop_common_events).
+    published = events.pop_common_events()
+    assert [e.token_ids for e in published] == [[1, 2, 3, 4]]
+
+
+def test_salvaged_event_votes_are_one_shot():
+    """Salvaged votes must not leak into a later step's quorum."""
+    aggregator = KVOutputAggregator(expected_finished_count=2)
+
+    # Failed step: worker 0's vote for block E0 is salvaged.
+    aggregator.ft_defer_kv_connector_outputs(
+        [KVConnectorOutput(kv_cache_events=make_lmcache_events(1, 2, 3, 4)), None]
+    )
+
+    # Recovered step is quiet: no live events, salvaged votes discarded.
+    aggregated = aggregator.aggregate(
+        [DummyModelRunnerOutput(), DummyModelRunnerOutput()]
+    )
+    assert aggregated.kv_connector_output.kv_cache_events is None
+
+    # A later step reports a different event: the salvaged votes must not
+    # leak into this step's quorum.
+    output0 = DummyModelRunnerOutput(kv_cache_events=make_lmcache_events(5, 6, 7, 8))
+    output1 = DummyModelRunnerOutput(kv_cache_events=make_lmcache_events(5, 6, 7, 8))
+    aggregated = aggregator.aggregate([output0, output1])
+
+    events = aggregated.kv_connector_output.kv_cache_events
+    assert events is not None
+    assert [e.token_ids for e in lmcache_published(events)] == [[5, 6, 7, 8]]

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -4,6 +4,7 @@
 
 from collections.abc import Iterator
 from dataclasses import dataclass
+from itertools import chain
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 import torch
@@ -72,6 +73,8 @@ class KVOutputAggregator:
         self._failed_recving_pending = set[str]()
         self._expected_finished_count = expected_finished_count
 
+        self._salvaged_outputs: list[KVConnectorOutput] = []
+
     @classmethod
     def from_connector(cls, connector: "KVConnectorBase", world_size: int):
         return cls(connector.get_finished_count() or world_size)
@@ -104,9 +107,17 @@ class KVOutputAggregator:
         aggregated_kv_connector_worker_meta = None
         combined_kv_cache_events = None
         invalid_block_ids = set[int]()
-        for model_runner_output in outputs:
-            assert model_runner_output is not None
-            kv_output = model_runner_output.kv_connector_output
+
+        salvaged, self._salvaged_outputs = self._salvaged_outputs, []
+
+        kv_outputs: list[KVConnectorOutput | None] = []
+        for o in outputs:
+            assert o is not None
+            kv_outputs.append(o.kv_connector_output)
+
+        # Salvaged outputs from failed workers join this round so their
+        # transfer progress isn't lost.
+        for kv_output in chain(salvaged, kv_outputs):
             if not kv_output:
                 continue
             # Allow the worker to dynamically update the expected number of
@@ -188,6 +199,20 @@ class KVOutputAggregator:
 
         return output
 
+    def ft_defer_kv_connector_outputs(
+        self, kv_connector_outputs: list[KVConnectorOutput | None]
+    ) -> None:
+        """Stash KV connector outputs from failed workers for replay at the
+        next aggregate().
+
+        kv_cache_events are dropped: replaying them could break the live
+        workers' event quorum, and a missed event only costs a cache miss.
+        """
+        for o in kv_connector_outputs:
+            if o is not None:
+                o.kv_cache_events = None
+        self._salvaged_outputs.extend(o for o in kv_connector_outputs if o)
+
 
 def _make_src_and_dst_indices(
     src_block_ids: list[int],

vllm/v1/executor/multiproc_executor.py

@@ -70,12 +70,25 @@ from vllm.utils.watch_dog import get_watch_dog, start_watch_dog
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.executor.abstract import Executor, FailureCallback
 from vllm.v1.executor.vllm_net_devices import set_worker_net_device
-from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
+from vllm.v1.outputs import (
+    AsyncModelRunnerOutput,
+    DraftTokenIds,
+    KVConnectorOutput,
+    ModelRunnerOutput,
+)
 from vllm.v1.worker.worker_base import WorkerWrapperBase
 
 logger = init_logger(__name__)
 
 
+class ExceptionWithKVConnectorOutput(Exception):
+    """Carries a worker's KV connector output along with its failure."""
+
+    def __init__(self, error_msg: str, kv_connector_output: KVConnectorOutput | None):
+        super().__init__(error_msg)
+        self.kv_connector_output = kv_connector_output
+
+
 class FutureWrapper(Future):
     def __init__(
         self,
@@ -433,6 +446,8 @@ class MultiprocExecutor(Executor):
 
         def get_response():
             responses = []
+            has_failure = False
+
             for mq in response_mqs:
                 dequeue_timeout = (
                     None if deadline is None else max(0.0, deadline - time.monotonic())
@@ -441,12 +456,46 @@ class MultiprocExecutor(Executor):
                     status, result = mq.dequeue(timeout=dequeue_timeout)
                 except TimeoutError as e:
                     raise TimeoutError(f"RPC call to {method} timed out.") from e
-                if status != WorkerProc.ResponseStatus.SUCCESS:
-                    raise RuntimeError(
-                        f"Worker failed with error '{result}', please check the"
-                        " stack trace above for the root cause"
+
+                if status == WorkerProc.ResponseStatus.FAILURE_WITH_KV_OUTPUT:
+                    if (
+                        not self.vllm_config.parallel_config.enable_fault_tolerance
+                        or self.is_failed
+                    ):
+                        raise RuntimeError("Worker failed with KV connector output")
+                    has_failure = True
+                    assert isinstance(result, KVConnectorOutput)
+                    responses.append(result)
+                elif status != WorkerProc.ResponseStatus.SUCCESS:
+                    if (
+                        not self.vllm_config.parallel_config.enable_fault_tolerance
+                        or self.is_failed
+                    ):
+                        raise RuntimeError(
+                            f"Worker failed with error '{result}', please check the"
+                            " stack trace above for the root cause"
+                        )
+                    has_failure = True
+                else:
+                    responses.append(result)
+
+            if has_failure:
+                kv_connector_outputs: list[KVConnectorOutput | None] = []
+                for resp in responses:
+                    if isinstance(resp, KVConnectorOutput):
+                        kv_connector_outputs.append(resp)
+                    elif resp is not None and hasattr(resp, "kv_connector_output"):
+                        kv_output = resp.kv_connector_output
+                        if kv_output is not None:
+                            kv_connector_outputs.append(kv_output)
+
+                if kv_connector_outputs and kv_output_aggregator is not None:
+                    kv_output_aggregator.ft_defer_kv_connector_outputs(
+                        kv_connector_outputs
                     )
-                responses.append(result)
+
+                raise RuntimeError("One or more workers failed")
+
             return responses[0] if output_rank is not None else responses
 
         future = FutureWrapper(
@@ -991,6 +1040,7 @@ class WorkerProc:
     class ResponseStatus(Enum):
         SUCCESS = auto()
         FAILURE = auto()
+        FAILURE_WITH_KV_OUTPUT = auto()
 
     def enqueue_output(self, output: Any):
         """Prepares output from the worker and enqueues it to the
@@ -1002,9 +1052,21 @@ class WorkerProc:
                 output = output.get_output()
             except Exception as e:
                 logger.exception("Error getting async model runner output")
-                output = e
+                kv_connector_output = output.get_kv_connector_output()
+                if (
+                    kv_connector_output is not None
+                    and not kv_connector_output.is_empty()
+                ):
+                    output = ExceptionWithKVConnectorOutput(str(e), kv_connector_output)
+                else:
+                    output = e
 
-        if isinstance(output, Exception):
+        if isinstance(output, ExceptionWithKVConnectorOutput):
+            result: tuple[WorkerProc.ResponseStatus, Any] = (
+                WorkerProc.ResponseStatus.FAILURE_WITH_KV_OUTPUT,
+                output.kv_connector_output,
+            )
+        elif isinstance(output, Exception):
             result = (WorkerProc.ResponseStatus.FAILURE, str(output))
         else:
             result = (WorkerProc.ResponseStatus.SUCCESS, output)

vllm/v1/fault_tolerance/engine_core_sentinel.py

@@ -5,7 +5,7 @@
 import json
 import threading
 from collections.abc import Callable
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import msgspec
 
@@ -13,6 +13,8 @@ from vllm.distributed import stateless_destroy_torch_distributed_process_group
 from vllm.distributed.utils import stateless_init_torch_distributed_process_group
 from vllm.logger import init_logger
 from vllm.utils.network_utils import get_open_port
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.engine import (
     FT_STATUS_CALL_ID,
     EngineCoreOutputs,
@@ -20,6 +22,7 @@ from vllm.v1.engine import (
     UtilityOutput,
 )
 from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult
+from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import RequestStatus
 from vllm.v1.serial_utils import UtilityResult, run_method
 
@@ -86,8 +89,7 @@ class EngineCoreSentinel:
         engine = self.engine
         aborted = engine.scheduler.finish_requests(None, RequestStatus.FINISHED_ABORTED)
         engine._send_abort_outputs(aborted)
-        if engine.batch_queue is not None:
-            engine.batch_queue.clear()
+
         if (
             hasattr(engine.model_executor, "is_failed")
             and engine.model_executor.is_failed
@@ -104,6 +106,39 @@ class EngineCoreSentinel:
         )
         self._push_status()
 
+    def _drain_batch_queue(self) -> None:
+        """Drain the batch queue, feeding leftover kv_connector_outputs to
+        the scheduler via update_from_output() with an empty SchedulerOutput.
+        """
+        engine = self.engine
+        if engine.batch_queue is None:
+            return
+        scheduler = cast(Scheduler, engine.scheduler)
+
+        while engine.batch_queue:
+            future, _, _ = engine.batch_queue.pop()
+            try:
+                model_output = future.result()
+            except Exception as e:
+                # Safe to drop: the KV output was already salvaged by the
+                # executor.
+                logger.warning(
+                    "[FT] Dropping exception from batch queue during fault "
+                    "handling: %s",
+                    e,
+                )
+                continue
+            kv_connector_output = (
+                model_output.kv_connector_output if model_output is not None else None
+            )
+            if kv_connector_output is None or kv_connector_output.is_empty():
+                continue
+            scheduler.update_from_output(
+                SchedulerOutput.make_empty(),
+                ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output),
+            )
+        engine.batch_queue.clear()
+
     def _push_status(self):
         """Push current health to the client so it can refresh its cache."""
         payload = {"id": self.engine_index, "status": self.status_type.name.lower()}
@@ -127,7 +162,7 @@ class EngineCoreSentinel:
             engine.step_counter = 0
 
         executor.collective_rpc("handle_ft_command", args=(ft_request,))
-
+        self._drain_batch_queue()
         self.status_type = EngineStatusType.HEALTHY
         logger.info("[FT] Engine %d status -> HEALTHY", self.engine_index)
         self.resumed.set()

vllm/v1/outputs.py

@@ -362,6 +362,15 @@ class AsyncModelRunnerOutput(ABC):
         """
         pass
 
+    @abstractmethod
+    def get_kv_connector_output(self) -> KVConnectorOutput | None:
+        """Get the underlying kv_connector_output.
+
+        When get_output() fails, this is used to salvage the step's KV
+        output instead of losing it with the failure.
+        """
+        pass
+
 
 @dataclass
 class DraftTokenIds:

vllm/v1/worker/gpu/async_utils.py

@@ -12,6 +12,7 @@ import vllm.envs as envs
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
 from vllm.v1.outputs import (
     AsyncModelRunnerOutput,
+    KVConnectorOutput,
     LogprobsTensors,
     ModelRunnerOutput,
     PoolerOutput,
@@ -216,6 +217,10 @@ class AsyncOutput(AsyncModelRunnerOutput):
 
         return self.model_runner_output
 
+    def get_kv_connector_output(self) -> KVConnectorOutput | None:
+        """Get the KV connector output for this async output."""
+        return self.model_runner_output.kv_connector_output
+
 
 class AsyncPoolingOutput(AsyncModelRunnerOutput):
     def __init__(
@@ -260,6 +265,10 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):
         self.model_runner_output.pooler_output = pooler_output
         return self.model_runner_output
 
+    def get_kv_connector_output(self) -> KVConnectorOutput | None:
+        """Get the KV connector output for this async output."""
+        return self.model_runner_output.kv_connector_output
+
 
 def async_copy_to_np(x: torch.Tensor) -> np.ndarray:
     return x.to("cpu", non_blocking=True, copy=True).numpy()

vllm/v1/worker/gpu_model_runner.py

@@ -363,6 +363,10 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
 
         return output
 
+    def get_kv_connector_output(self) -> KVConnectorOutput | None:
+        """Get the KV connector output for this async output."""
+        return self._model_runner_output.kv_connector_output
+
 
 def _copy_pooler_output_to_cpu(
     raw_pooler_output: PoolerOutput, finished_mask: list[bool]
@@ -447,6 +451,10 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
         del self._raw_pooler_output
         return self._model_runner_output
 
+    def get_kv_connector_output(self) -> KVConnectorOutput | None:
+        """Get the KV connector output for this async output."""
+        return self._model_runner_output.kv_connector_output
+
 
 class ExecuteModelState(NamedTuple):
     """Ephemeral cached state transferred between execute_model() and