diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 49ca683de7..6155932ea7 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -670,16 +670,25 @@ def test_moriio_wrapper_rejects_invalid_messages(role, payload, match): wrapper._handle_message(payload) -def test_block_id_length_mismatch_raises_value_error(): +def test_local_block_ids_longer_than_remote_raises_value_error(): cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) worker = _worker({"layer": cache}, {"layer": _full_spec()}) - with pytest.raises(ValueError, match="must have the same length"): + with pytest.raises(ValueError, match="longer than remote_block_ids"): moriio_layout.compute_block_transfer_offsets( "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks ) +def test_empty_local_block_ids_is_free_only_noop(): + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + assert moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [], [4, 5], _remote_meta().num_blocks + ) == ([], [], []) + + def test_registration_regions_do_not_split_interleaved_or_mla_cache(): separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16) interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py index 9abc9b957c..6488135690 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py @@ -253,6 +253,7 @@ def test_worker_get_finished_counts_structured_release_fan_in(): worker.transfer_id_to_request_id = {"tx-fanin": "req-fanin"} worker._consumer_notification_counts = {} worker._completed_consumer_notifications = set() + worker._pending_unmapped_acks = [] assert worker.get_finished() == (set(), set()) assert worker._consumer_notification_counts == {"tx-fanin": 1} diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 3d0d216d99..ea41adf437 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1030,6 +1030,16 @@ class MoRIIOConnectorWorker: use_mla=self.use_mla, ) self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + # READ-mode producer: a decode release-ACK can arrive BEFORE + # start_load_kv populates transfer_id_to_request_id (the notify races + # ahead of the scheduler->worker sync). Buffer such ACKs and retry them + # next get_finished tick instead of dropping them -- dropping loses the + # completion, so the request is never marked done_sending, its KV blocks + # leak, and the prefill KV cache wedges at high concurrency. Buffered + # BEFORE resolve_moriio_transfer_ack, so each ACK is counted exactly once + # (on the tick its mapping exists) -- the heterogeneous-TP ack-counting + # is preserved. + self._pending_unmapped_acks: list = [] # TODO: consider the integration of flashinfer or other backends. self.backend_name = backend.get_name() @@ -1542,16 +1552,22 @@ class MoRIIOConnectorWorker: # pop_finished_req_ids returns release ACKs sent by decode. Keep # duplicate ACKs because heterogeneous TP can fan multiple decode # ranks into one prefill rank for the same transfer_id. - finished_acks = self.moriio_wrapper.pop_finished_req_ids() + # Combine freshly-arrived ACKs with any buffered from prior ticks + # whose transfer_id wasn't mapped yet (notify raced ahead of + # start_load_kv); retry the lookup every tick. Buffered before + # resolve_moriio_transfer_ack so each ACK is counted exactly once. + finished_acks = self._pending_unmapped_acks + list( + self.moriio_wrapper.pop_finished_req_ids() + ) + self._pending_unmapped_acks = [] resolved_transfer_ids: set[TransferId] = set() for ack in finished_acks: transfer_id = ack if isinstance(ack, str) else ack.transfer_id if transfer_id not in self.transfer_id_to_request_id: - logger.warning( - "Could not find %s in transfer_id_to_request_id " - "lookup table. This could lead to a possible hang.", - transfer_id, - ) + # Mapping not populated yet -- buffer and retry next tick, + # do NOT drop (dropping leaks producer KV at high conc and + # wedges the prefill). + self._pending_unmapped_acks.append(ack) continue resolved_transfer_id = resolve_moriio_transfer_ack( ack, @@ -1933,6 +1949,23 @@ class MoRIIOConnectorWorker: ), ) + @staticmethod + def _is_sq_full_status(status) -> bool: + """True if a MoRIIO transfer status is a transient RDMA send-queue-full + rejection (retryable backpressure), not a terminal failure. + + read_remote_data posts the RDMA READ synchronously (the mori executor + joins its worker before returning and marks the status on the calling + thread), so a send-queue-full rejection is a Failed() status the moment + the call returns. mori surfaces it as a generic ERR_RDMA_OP carrying + "SQ full" in the message (no distinct code), so we match the message. + Only meaningful once status.Failed() is True. + """ + try: + return bool(status.Failed()) and "SQ full" in (status.Message() or "") + except Exception: + return False + def _read_blocks( self, local_block_ids: list[int], @@ -1950,6 +1983,8 @@ class MoRIIOConnectorWorker: dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0) sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id) + # SQ-full backpressure deadline, shared across this request's layers. + _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout for layer_name in self.layer_name_to_local_kv_cache_metadata: sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( layer_name @@ -1962,9 +1997,35 @@ class MoRIIOConnectorWorker: remote_tp_size=remote_tp_size, ) # TODO : apply multi-session batch-read when moriio support it - transfer_status = self.moriio_wrapper.read_remote_data( - offs[2], offs[0], offs[1], sessions[sess_idx] - ) + # + # SQ-full backpressure: read_remote_data posts the RDMA READ + # SYNCHRONOUSLY, so a send-queue-full rejection (per-QP HW cap) comes + # back as a Failed() status right here. A SEPARATE CQ-poll thread + # drains completions and frees SQ depth, so back off and RE-POST + # rather than let a transient rejection abort the request. No + # self-deadlock (the drain is off-thread); the reserve is + # all-or-nothing (nothing posted on a rejected attempt). Bounded by + # transfer_timeout; on sustained overload store the failed status and + # let get_finished handle it non-fatally (notify prefill + drop). + _backoff = 0.001 + while True: + transfer_status = self.moriio_wrapper.read_remote_data( + offs[2], offs[0], offs[1], sessions[sess_idx] + ) + if not self._is_sq_full_status(transfer_status): + break + if time.monotonic() > _sq_deadline: + logger.warning( + "MoRIIO READ send queue stayed full past " + "transfer_timeout for req %s layer %s; storing failed " + "status (get_finished notifies prefill and drops the " + "request). Raise qp_per_transfer if frequent.", + request_id, + layer_name, + ) + break + time.sleep(_backoff) + _backoff = min(_backoff * 2, 0.05) with self.moriio_wrapper.lock: self._recving_transfers[request_id].append(transfer_status) self._recving_transfers_callback_addr[request_id] = ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index d80541ade6..90aa5f823a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -312,10 +312,16 @@ def compute_block_transfer_offsets( [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] ] = merge_contiguous_offsets, ) -> tuple[list[int], list[int], list[int]]: - if len(local_block_ids) != len(remote_block_ids): + # A shorter (or empty) local list is the READ-mode "drop the transfer, just + # free the prefill blocks" case (full-prefix-hit / aborted-before-scheduled): + # decode pulls fewer blocks than the prefill holds. The zip loop below pairs + # local[i]<->remote[i] and sizes by len(local), so a short local transfers + # only what decode allocated and an empty local is a no-op. A longer local + # list is a genuine bug and still fails loudly. + if len(local_block_ids) > len(remote_block_ids): raise ValueError( - "local_block_ids and remote_block_ids must have the same length: " - f"{len(local_block_ids)} != {len(remote_block_ids)}" + "local_block_ids longer than remote_block_ids: " + f"{len(local_block_ids)} > {len(remote_block_ids)}" ) geometry = get_layer_transfer_geometry( layer_name, kv_cache, layer_to_spec, remote_num_blocks