diff --git a/tests/v1/kv_connector/unit/test_mooncake_arrival_doorbell.py b/tests/v1/kv_connector/unit/test_mooncake_arrival_doorbell.py new file mode 100644 index 000000000000..3097ba14308d --- /dev/null +++ b/tests/v1/kv_connector/unit/test_mooncake_arrival_doorbell.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the SGLang-style rid-based arrival doorbell. + +D-side ``ArrivalDoorbell`` allocates a per-request slot in a GPU buffer +registered with Mooncake; the slot is zeroed and a random nonce is held +in process memory as the expected value. P writes the nonce into the +slot via RDMA as the LAST descriptor of the batch, then ZMQ-replies +``ok_reqs``. D reads its own local slot and verifies the nonce -- a +match proves both the doorbell and all preceding KV descriptors in the +same RDMA batch landed at the remote. + +P-side ``NoncePad`` stages a copy of each request's nonce in a small +device-resident scratch buffer (registered with Mooncake) so the +producer can supply the source pointers for those extra descriptors. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.arrival_doorbell import ( + ARRIVAL_SLOT_BYTES, + ArrivalDoorbell, + NoncePad, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestArrivalDoorbellAllocate: + def _make(self, capacity: int = 8) -> tuple[MagicMock, ArrivalDoorbell]: + engine = MagicMock() + engine.batch_register_memory = MagicMock(return_value=0) + return engine, ArrivalDoorbell(engine, torch.device("cuda"), capacity=capacity) + + def test_allocate_returns_unique_addrs(self): + _, db = self._make(capacity=4) + a = db.allocate("r1") + b = db.allocate("r2") + assert a.slot_addr != b.slot_addr + assert a.expected_nonce != b.expected_nonce + + def test_allocate_zeros_slot_on_device(self): + _, db = self._make() + # Pre-poison the buffer. + db._buffer.fill_(0xAB) + a = db.allocate("r1") + slot_offset = a.slot_addr - db._buffer.data_ptr() + slot = db._buffer[slot_offset : slot_offset + ARRIVAL_SLOT_BYTES] + assert torch.all(slot == 0).item() + + def test_allocate_raises_when_exhausted(self): + _, db = self._make(capacity=2) + db.allocate("r1") + db.allocate("r2") + with pytest.raises(RuntimeError, match="exhausted"): + db.allocate("r3") + + def test_release_lets_capacity_recover(self): + _, db = self._make(capacity=1) + a = db.allocate("r1") + db.release("r1") + b = db.allocate("r2") + assert b.slot_addr == a.slot_addr # reused + + def test_double_release_is_safe(self): + _, db = self._make() + db.allocate("r1") + db.release("r1") + db.release("r1") # must not raise + + def test_registers_buffer_with_engine_once(self): + engine, db = self._make() + engine.batch_register_memory.assert_called_once() + ptrs, lens = engine.batch_register_memory.call_args[0] + assert ptrs == [db._buffer.data_ptr()] + assert lens == [db._buffer.numel()] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestArrivalDoorbellVerify: + def _make(self): + engine = MagicMock() + engine.batch_register_memory = MagicMock(return_value=0) + return ArrivalDoorbell(engine, torch.device("cuda"), capacity=4) + + def test_verify_matches_after_remote_write(self): + db = self._make() + h = db.allocate("r1") + # Simulate the producer's remote RDMA write landing on the slot. + slot_offset = h.slot_addr - db._buffer.data_ptr() + nonce_bytes = h.expected_nonce.to_bytes(8, "little") + payload = torch.tensor( + list(nonce_bytes) + [0] * (ARRIVAL_SLOT_BYTES - 8), + dtype=torch.uint8, + device=db._buffer.device, + ) + db._buffer[slot_offset : slot_offset + ARRIVAL_SLOT_BYTES] = payload + assert db.verify("r1") is True + + def test_verify_fails_when_slot_still_zero(self): + db = self._make() + db.allocate("r1") + # Slot was zeroed at allocation; producer never wrote anything. + assert db.verify("r1") is False + + def test_verify_fails_on_wrong_nonce(self): + db = self._make() + h = db.allocate("r1") + slot_offset = h.slot_addr - db._buffer.data_ptr() + bad = (h.expected_nonce ^ 0xDEADBEEF).to_bytes(8, "little") + payload = torch.tensor( + list(bad) + [0] * (ARRIVAL_SLOT_BYTES - 8), + dtype=torch.uint8, + device=db._buffer.device, + ) + db._buffer[slot_offset : slot_offset + ARRIVAL_SLOT_BYTES] = payload + assert db.verify("r1") is False + + def test_verify_unknown_req_returns_false(self): + db = self._make() + assert db.verify("never_allocated") is False + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +class TestNoncePad: + def _make(self, capacity: int = 8) -> tuple[MagicMock, NoncePad]: + engine = MagicMock() + engine.batch_register_memory = MagicMock(return_value=0) + return engine, NoncePad(engine, torch.device("cuda"), capacity=capacity) + + def test_stage_writes_nonces_and_returns_src_ptrs(self): + _, pad = self._make(capacity=4) + nonces = [0x1111111111111111, 0x2222222222222222] + src_ptrs = pad.stage(nonces) + assert len(src_ptrs) == 2 + assert src_ptrs[0] != src_ptrs[1] + # Verify the nonces actually landed on the device. + for nonce, ptr in zip(nonces, src_ptrs): + offset = ptr - pad._buffer.data_ptr() + chunk = pad._buffer[offset : offset + 8].cpu().numpy().tobytes() + assert int.from_bytes(chunk, "little") == nonce + + def test_stage_round_robins_slots(self): + _, pad = self._make(capacity=2) + p1 = pad.stage([1])[0] + _ = pad.stage([2])[0] + p3 = pad.stage([3])[0] # wraps + assert p3 == p1 + + def test_registers_with_engine(self): + engine, pad = self._make() + engine.batch_register_memory.assert_called_once() + ptrs, lens = engine.batch_register_memory.call_args[0] + assert ptrs == [pad._buffer.data_ptr()] + assert lens == [pad._buffer.numel()] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/arrival_doorbell.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/arrival_doorbell.py new file mode 100644 index 000000000000..830ad7ece2ce --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/arrival_doorbell.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SGLang-style rid-based arrival doorbell for MooncakeConnector. + +Background +---------- +SGLang's ``MooncakeKVManager`` proves a transfer landed at the decoder +by tagging each request with a ``bootstrap_room`` (the rid) and routing +the per-request status through a separate, ordered channel +(``sync_status_to_decode_endpoint`` over ZMQ, plus the ``AUX_DATA`` +RDMA push that carries the same room). The decoder treats the rid as +identity: only when the rid arrives back does the consumer mark the +request as truly received. + +vLLM's MooncakeConnector already routes per-request status via +``MooncakeXferResponse.ok_reqs`` / ``err_reqs`` (req_id == rid). That +channel tells D *which* request the producer believes succeeded, but +it does not by itself prove the RDMA bytes for that request landed in +the decoder's HBM -- ``batch_transfer_sync_write`` returning zero only +means the local NIC has been informed, not that the remote memory was +updated. The original Qwen3-Omni KV cache corruption surfaced exactly +this gap (engine returned 0, decoder saw head/tail zeros). + +Design +------ +This module implements a small "arrival doorbell" mechanism aligned to +SGLang's rid pattern: + +* D allocates an 8-byte slot in a GPU buffer registered with Mooncake + and generates a random 64-bit nonce per request. The slot is zeroed. +* D ships ``(slot_addr, expected_nonce)`` to P inside + ``MooncakeXferMetadata``. +* P stages the nonces in its own GPU scratch buffer (``NoncePad``) and + appends one tiny descriptor per request to the end of its + ``batch_transfer_sync_write`` call. RDMA WRITE ordering on a session + guarantees the doorbell write retires after the preceding KV writes + on the same destination memory region. +* After ZMQ ``ok_reqs`` arrives, D reads its own local slot and checks + ``actual_nonce == expected_nonce``. A match proves the rid arrived, + and -- by RDMA ordering -- that the preceding KV bytes did too. A + mismatch demotes the request to ``finished_recving`` failure so the + scheduler can retry rather than feed corrupted KV cache into decode. + +The slot is intentionally tiny (16 bytes including 8-byte padding) so +the per-request overhead is one descriptor and 16 bytes of RDMA write. +""" + +from __future__ import annotations + +import secrets +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from mooncake.engine import TransferEngine + +logger = init_logger(__name__) + +# Size of one doorbell slot in bytes. The first 8 bytes hold the +# little-endian nonce, the remaining 8 are reserved (zero) padding so +# slots stay aligned and we have room to grow the protocol later +# without breaking layout compatibility. +ARRIVAL_SLOT_BYTES = 16 + + +@dataclass(frozen=True) +class DoorbellHandle: + """Per-request handle returned by ``ArrivalDoorbell.allocate``. + + ``slot_addr`` is shipped to the producer; ``expected_nonce`` is the + value the consumer expects to read back from the slot after the + producer's RDMA write completes. + """ + + req_id: str + slot_addr: int + expected_nonce: int + + +class ArrivalDoorbell: + """Consumer-side pool of arrival-doorbell slots on GPU memory. + + The whole buffer is registered with the Mooncake transfer engine + once at construction so that producers can RDMA-write directly + into individual slots. + """ + + def __init__( + self, + engine: TransferEngine, + device: torch.device, + capacity: int = 4096, + ) -> None: + if capacity <= 0: + raise ValueError("capacity must be positive") + self._capacity = capacity + self._device = device + total_bytes = capacity * ARRIVAL_SLOT_BYTES + self._buffer = torch.zeros(total_bytes, dtype=torch.uint8, device=device) + ret = engine.batch_register_memory( + [self._buffer.data_ptr()], [self._buffer.numel()] + ) + if ret != 0: + raise RuntimeError( + f"Failed to register arrival-doorbell buffer with Mooncake (ret={ret})" + ) + self._lock = threading.Lock() + self._free: list[int] = list(range(capacity)) + # Maps req_id -> (slot_idx, expected_nonce). + self._in_use: dict[str, tuple[int, int]] = {} + + def _slot_addr(self, slot_idx: int) -> int: + return self._buffer.data_ptr() + slot_idx * ARRIVAL_SLOT_BYTES + + def allocate(self, req_id: str) -> DoorbellHandle: + """Reserve a slot for ``req_id`` and return its handle.""" + with self._lock: + if not self._free: + raise RuntimeError( + "ArrivalDoorbell pool exhausted " + f"(capacity={self._capacity}, in_use={len(self._in_use)}). " + "Increase capacity or drain pending transfers." + ) + if req_id in self._in_use: + raise RuntimeError( + f"req_id {req_id} already has an allocated doorbell slot" + ) + slot_idx = self._free.pop() + # 8-byte nonce; the high bit is masked to keep the value + # inside Python's signed-int safe range when round-tripping + # through some msgspec encoders later. + nonce = secrets.token_bytes(8) + nonce_int = int.from_bytes(nonce, "little") & ((1 << 63) - 1) + self._in_use[req_id] = (slot_idx, nonce_int) + # Zero the slot on device; this happens outside the lock since + # the slice is exclusive to this owner. + start = slot_idx * ARRIVAL_SLOT_BYTES + self._buffer[start : start + ARRIVAL_SLOT_BYTES].zero_() + return DoorbellHandle( + req_id=req_id, + slot_addr=self._slot_addr(slot_idx), + expected_nonce=nonce_int, + ) + + def has_slot(self, req_id: str) -> bool: + """Return ``True`` iff ``req_id`` currently owns a slot.""" + with self._lock: + return req_id in self._in_use + + def verify(self, req_id: str) -> bool: + """Return ``True`` iff the consumer-side slot holds ``expected_nonce``. + + Returns ``False`` if the slot is unknown or the nonce does not + match. This reads the device slot via a tiny D2H copy. Callers + must only invoke verify *after* the producer has signalled + success through the ZMQ status channel; otherwise the slot is + racy. + """ + with self._lock: + entry = self._in_use.get(req_id) + if entry is None: + return False + slot_idx, expected = entry + start = slot_idx * ARRIVAL_SLOT_BYTES + slot_bytes = bytes(self._buffer[start : start + 8].cpu().tolist()) + actual = int.from_bytes(slot_bytes, "little") + if actual != expected: + logger.warning( + "Arrival doorbell mismatch for req=%s: expected=%x actual=%x", + req_id, + expected, + actual, + ) + return False + return True + + def release(self, req_id: str) -> None: + """Return the slot to the free pool. Safe to call twice.""" + with self._lock: + entry = self._in_use.pop(req_id, None) + if entry is not None: + slot_idx, _ = entry + self._free.append(slot_idx) + + @property + def capacity(self) -> int: + return self._capacity + + +class NoncePad: + """Producer-side device-resident scratch buffer for staging nonces. + + Each call to ``stage`` reserves the next N slots (wrapping around + the capacity), writes the supplied nonces to those slots, and + returns the source pointers so the caller can append RDMA write + descriptors targeting the matching ``ArrivalDoorbell`` slots on + the consumer. + + Slot reuse is by ring buffer: the pad assumes the previous use of + each slot has fully retired by the time the producer wraps. With a + capacity an order of magnitude larger than the producer's + concurrent transfer fan-out this is safe in practice (the pad is + on the *send* side of an already-completed batch by the time we + reuse the slot). + """ + + def __init__( + self, + engine: TransferEngine, + device: torch.device, + capacity: int = 4096, + ) -> None: + if capacity <= 0: + raise ValueError("capacity must be positive") + self._capacity = capacity + total_bytes = capacity * ARRIVAL_SLOT_BYTES + self._buffer = torch.zeros(total_bytes, dtype=torch.uint8, device=device) + ret = engine.batch_register_memory( + [self._buffer.data_ptr()], [self._buffer.numel()] + ) + if ret != 0: + raise RuntimeError( + f"Failed to register nonce pad with Mooncake (ret={ret})" + ) + self._lock = threading.Lock() + self._cursor = 0 + + def _slot_addr(self, slot_idx: int) -> int: + return self._buffer.data_ptr() + slot_idx * ARRIVAL_SLOT_BYTES + + def stage(self, nonces: list[int]) -> list[int]: + """Write ``nonces`` to consecutive slots and return src pointers.""" + if not nonces: + return [] + n = len(nonces) + with self._lock: + slots = [(self._cursor + i) % self._capacity for i in range(n)] + self._cursor = (self._cursor + n) % self._capacity + # Build a tiny host tensor and one D2D-style copy per slot. The + # slots may wrap so we just iterate; the bytes-per-stage is + # ``ARRIVAL_SLOT_BYTES * n`` which is negligible compared to KV + # transfer sizes. + for slot_idx, nonce in zip(slots, nonces): + payload = nonce.to_bytes(8, "little") + b"\x00" * (ARRIVAL_SLOT_BYTES - 8) + staged = torch.tensor( + list(payload), dtype=torch.uint8, device=self._buffer.device + ) + start = slot_idx * ARRIVAL_SLOT_BYTES + self._buffer[start : start + ARRIVAL_SLOT_BYTES] = staged + return [self._slot_addr(s) for s in slots] + + @property + def capacity(self) -> int: + return self._capacity diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 1bc23cead5b4..790799d71a77 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -32,6 +32,11 @@ SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.arrival_doorbell import ( + ARRIVAL_SLOT_BYTES, + ArrivalDoorbell, + NoncePad, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import ( MooncakeBootstrapServer, RegisterWorkerPayload, @@ -334,6 +339,16 @@ class MooncakeXferMetadata( block_lens: list[int] registered_layer_names: list[str] = msgspec.field(default_factory=list) registered_layer_indices: list[int] = msgspec.field(default_factory=list) + # SGLang-style rid-based arrival doorbell, keyed by D-side req_id. + # Each value is ``(slot_addr_on_consumer_gpu, expected_nonce)``. The + # producer stages the nonce into a local scratch buffer and appends + # a tiny RDMA-write descriptor that targets ``slot_addr`` at the + # tail of the batch. The consumer reads its own slot after ``ok_reqs`` + # arrives -- a matching nonce proves both the doorbell and (by RDMA + # WRITE ordering on the same session/region) the preceding KV + # descriptors landed in remote HBM. ``None`` keeps backward wire- + # compatibility with peers that don't implement the doorbell yet. + doorbells: dict[ReqId, tuple[int, int]] | None = None class MooncakeXferResponseStatus(IntEnum): @@ -869,6 +884,24 @@ def __init__( self.device_kv_caches: dict[str, torch.Tensor] = {} self.reqs_need_send: dict[TransferId, SendBlockMeta] = {} + # SGLang-style per-request arrival doorbell (rid -> nonce). The + # consumer owns the slot pool; the producer owns the staging pad. + # Roles with ``kv_both`` participate on both sides. + doorbell_capacity = int( + kv_transfer_config.kv_connector_extra_config.get("doorbell_capacity", 4096) + ) + self._device = torch.device("cuda", self.device_id) + self._arrival_doorbell: ArrivalDoorbell | None = None + self._nonce_pad: NoncePad | None = None + if not self.is_kv_producer: + self._arrival_doorbell = ArrivalDoorbell( + self.engine, self._device, capacity=doorbell_capacity + ) + if not self.is_kv_consumer: + self._nonce_pad = NoncePad( + self.engine, self._device, capacity=doorbell_capacity + ) + # For kv_both, we will act both prefiller and decoder. if not self.is_kv_consumer: # Background threads for sending kvcaches to D. @@ -1212,6 +1245,30 @@ async def wait_and_ret( if d_req_id not in err_req_set ] + # Append rid-tagged arrival-doorbell descriptors to the batch. + # Each one is a tiny RDMA write that the consumer reads back + # after ``ok_reqs`` arrives -- a matching nonce proves the + # preceding KV writes on this session/region landed in HBM. + if ( + src_ptrs + and ok_ready_reqs + and self._nonce_pad is not None + and meta.doorbells + ): + doorbell_pairs: list[tuple[ReqId, tuple[int, int]]] = [ + (d_req_id, meta.doorbells[d_req_id]) + for d_req_id, _ in ok_ready_reqs + if d_req_id in meta.doorbells + ] + if doorbell_pairs: + nonces = [pair[1][1] for pair in doorbell_pairs] + dst_slot_addrs = [pair[1][0] for pair in doorbell_pairs] + src_slot_addrs = self._nonce_pad.stage(nonces) + for src, dst in zip(src_slot_addrs, dst_slot_addrs): + src_ptrs.append(src) + dst_ptrs.append(dst) + lengths.append(ARRIVAL_SLOT_BYTES) + if src_ptrs: remote_session = f"{meta.remote_hostname}:{meta.remote_port}" ret_value = await self.sender_loop.run_in_executor( @@ -1451,6 +1508,14 @@ def _send_blocks( dst_ptrs: list[int], lengths: list[int], ) -> int: + # The batch ``src_ptrs``/``dst_ptrs``/``lengths`` already include + # the per-req arrival-doorbell descriptors appended by + # ``send_kv_to_decode`` (when the consumer supplied + # ``MooncakeXferMetadata.doorbells``). Their rid-tagged nonces + # are validated on the consumer in ``process_pulling_result``; + # this mirrors SGLang's ``bootstrap_room`` arrival proof while + # keeping vLLM's existing ``ok_reqs``/``err_reqs`` ZMQ channel + # as the carrier for per-request status. start_time = time.perf_counter() ret_value = self.engine.batch_transfer_sync_write( remote_session, src_ptrs, dst_ptrs, lengths @@ -1631,6 +1696,37 @@ async def receive_kv_from_single_worker( pull_metas: dict[ReqId, PullReqMeta], ): req_ids = set(pull_metas) + + # Allocate per-(worker, req) arrival-doorbell slots. ``worker_addr`` + # may exhaust the pool if there are too many in-flight transfers; + # in that case we fall back to no-doorbell semantics (slightly + # weaker arrival proof) rather than failing the request. + doorbells: dict[ReqId, tuple[int, int]] | None = None + doorbell_req_ids: list[ReqId] = [] + if self._arrival_doorbell is not None: + doorbells = {} + for req_id in pull_metas: + # Tag the slot key with the producer worker addr so the + # same req can be in-flight to multiple P workers + # concurrently without colliding on the same slot. + slot_key = f"{req_id}@{worker_addr}" + try: + handle = self._arrival_doorbell.allocate(slot_key) + except RuntimeError as e: + logger.warning( + "Arrival doorbell pool exhausted for %s: %s -- " + "falling back to ok_reqs-only signalling.", + slot_key, + e, + ) + for k in doorbell_req_ids: + self._arrival_doorbell.release(f"{k}@{worker_addr}") + doorbell_req_ids = [] + doorbells = None + break + doorbell_req_ids.append(req_id) + doorbells[req_id] = (handle.slot_addr, handle.expected_nonce) + metadata = MooncakeXferMetadata( remote_hostname=self.hostname, remote_port=self.rpc_port, @@ -1644,6 +1740,7 @@ async def receive_kv_from_single_worker( block_lens=self.block_len_per_layer, registered_layer_names=self.registered_layer_names, registered_layer_indices=self.registered_layer_indices, + doorbells=doorbells, ) encoded_data = self._encoder.encode(metadata) @@ -1654,7 +1751,6 @@ async def receive_kv_from_single_worker( "Sending kv transfer request for %s on path: %s", req_ids, worker_addr ) - # Send query for the request. try: with make_zmq_socket( self.async_zmq_ctx, worker_addr, zmq.DEALER, bind=False, linger=0 @@ -1675,7 +1771,9 @@ async def receive_kv_from_single_worker( ) self.xfer_stats.record_failed_recv() return - self.process_pulling_result(response, pull_metas) + self.process_pulling_result( + response, pull_metas, worker_addr=worker_addr + ) if response.status == MooncakeXferResponseStatus.FINISH: break except zmq.ContextTerminated: @@ -1684,13 +1782,38 @@ async def receive_kv_from_single_worker( logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e) self.xfer_stats.record_failed_recv() return + finally: + if self._arrival_doorbell is not None: + for req_id in doorbell_req_ids: + self._arrival_doorbell.release(f"{req_id}@{worker_addr}") def process_pulling_result( self, response: MooncakeXferResponse, pull_metas: dict[ReqId, PullReqMeta], + worker_addr: str | None = None, ): - ok_reqs: list[ReqId] = response.ok_reqs or [] + raw_ok_reqs: list[ReqId] = response.ok_reqs or [] + + # Filter ok_reqs by the arrival doorbell. A missing doorbell entry + # (consumer-side pool exhausted, or peer running an older + # protocol) is treated as "trust ok_reqs" so we degrade + # gracefully; a present-but-mismatched slot demotes the req so + # the scheduler retries instead of running with corrupt KV. + ok_reqs: list[ReqId] = [] + doorbell_failures: list[ReqId] = [] + for req_id in raw_ok_reqs: + if self._arrival_doorbell is None or worker_addr is None: + ok_reqs.append(req_id) + continue + slot_key = f"{req_id}@{worker_addr}" + if not self._arrival_doorbell.has_slot(slot_key): + ok_reqs.append(req_id) + continue + if self._arrival_doorbell.verify(slot_key): + ok_reqs.append(req_id) + else: + doorbell_failures.append(req_id) for req_id in ok_reqs: pull_meta = pull_metas[req_id] @@ -1702,6 +1825,17 @@ def process_pulling_result( if ok_reqs: logger.debug("pulling kv_caches for %s finished", ok_reqs) + if doorbell_failures: + logger.error( + "Arrival doorbell mismatch from %s for %s -- the producer " + "reported success but the rid-tagged nonce did not land " + "in the consumer's KV region. Marking as recv failure so " + "the scheduler can retry.", + worker_addr, + doorbell_failures, + ) + self.xfer_stats.record_failed_recv() + if response.err_reqs: logger.error( "pulling kv_caches for %s failed: %s",