Stage mixed HiSparse host imports through GPU

Route same-node host-bound HiSparse regions through the generic NIXL host stager while preserving direct GPU transfers for device-resident regions. Complete mixed receives only after both transfer paths finish and preserve failure reporting across sibling cancellation.

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Author: MatthewBonanni

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -6,6 +6,7 @@
 import os
 import tempfile
 import textwrap
+import threading
 import time
 import uuid
 from collections import defaultdict
@@ -1906,6 +1907,67 @@ def test_host_stager_is_only_initialized_for_same_host_reads(monkeypatch):
         stager_factory.assert_called_once()
 
 
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("staging_finishes_first", [False, True])
+def test_mixed_host_staging_waits_for_both_receive_parts(staging_finishes_first):
+    """A mixed host/device receive completes only after both paths finish."""
+    worker = object.__new__(NixlConnectorWorker)
+    worker._recving_metadata = {"request": MagicMock()}
+    worker._recving_transfers = {"request": [7]}
+    worker._failed_recv_lock = threading.Lock()
+    worker._failed_recv_pending = set()
+    worker._failed_recv_reported = set()
+    worker._pending_recv_notifs = {"request": [("agent", b"notification")]}
+    worker._send_pending_recv_notifs = MagicMock()
+    worker._report_failed_recv = MagicMock()
+    worker.xfer_stats = MagicMock()
+    worker.nixl_wrapper = MagicMock()
+    worker.nixl_wrapper.check_xfer_state.return_value = "DONE"
+
+    stager = MagicMock()
+    stager.active_req_ids = {"request"}
+
+    def finish_staging():
+        stager.active_req_ids = set()
+        return {"request"}, set()
+
+    stager.get_finished.side_effect = finish_staging
+    worker._host_stager = stager
+
+    if staging_finishes_first:
+        assert worker._get_finished_host_staging() == set()
+        assert worker._pop_done_transfers(worker._recving_transfers, is_recv=True) == {
+            "request"
+        }
+    else:
+        assert (
+            worker._pop_done_transfers(worker._recving_transfers, is_recv=True) == set()
+        )
+        assert worker._get_finished_host_staging() == {"request"}
+
+    worker._send_pending_recv_notifs.assert_called_once_with("request")
+
+
+@pytest.mark.cpu_test
+def test_mixed_host_staging_reports_failure_after_aborted_sibling_drains():
+    worker = object.__new__(NixlConnectorWorker)
+    worker._recving_metadata = {"request": MagicMock()}
+    worker._recving_transfers = {}
+    worker._failed_recv_lock = threading.Lock()
+    worker._failed_recv_pending = {"request"}
+    worker._pending_recv_notifs = {}
+    worker._report_failed_recv = MagicMock()
+
+    stager = MagicMock()
+    stager.active_req_ids = set()
+    stager.get_finished.return_value = set(), set()
+    worker._host_stager = stager
+
+    assert worker._get_finished_host_staging() == set()
+    assert worker._failed_recv_pending == set()
+    worker._report_failed_recv.assert_called_once_with("request")
+
+
 def _run_abort_timeout_test(llm: LLM, timeout: int):
     """Helper function to run the abort timeout test logic."""
     remote_prefill_opts = {

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -3,18 +3,25 @@
 """Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill."""
 
 import gc
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import msgspec
+import numpy as np
 import pytest
 import torch
 
 from tests.v1.attention.utils import MockMambaBuilder
 from vllm import LLM, SamplingParams
 from vllm.config import KVTransferConfig, set_current_vllm_config
+from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.nixl import (
+    HiSparseNixlDestination,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
     NixlAgentMetadata,
+    RemoteMeta,
+    ReqMeta,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
     NixlConnectorWorker,
@@ -1971,6 +1978,83 @@ def test_hisparse_host_import_keeps_host_blocks_out_of_gpu_regions():
     assert request_metadata.local_block_ids == ([10, 11], [20])
 
 
+@pytest.mark.cpu_test
+def test_hisparse_same_host_import_stages_host_regions():
+    """Host regions use GPU staging while device-only regions transfer directly."""
+    destination = object.__new__(HiSparseNixlDestination)
+    destination.host_regions = [(1_000, 16), None]
+    destination._descriptor_offsets = [0, None]
+    destination._xfer_handle = 100
+    destination._host_desc_lens = np.full(4, 16, dtype=np.int64)
+    destination._host_addrs = np.arange(1_000, 1_064, 16, dtype=np.uint64)
+    host_buffer = torch.empty(64, dtype=torch.uint8)
+    destination._host_pools = {host_buffer.data_ptr(): (host_buffer, [])}
+
+    worker = MagicMock()
+    worker._has_mamba = False
+    worker.use_mla = True
+    worker.transfer_topo = MagicMock()
+    worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
+        remote_physical_blocks_per_logical=1,
+        remote_block_size=64,
+    )
+    worker.transfer_topo.tp_ratio.return_value = 1
+    worker._block_ids_by_region.side_effect = [
+        [[1, 2], [1, 2]],
+        [[3, 4], [5, 6]],
+    ]
+    worker._apply_prefix_caching_by_region.side_effect = lambda local, remote: (
+        local,
+        remote,
+    )
+    worker._compute_desc_ids.side_effect = [
+        np.array([10, 11, 20, 21]),
+        np.array([30, 31]),
+    ]
+    worker.dst_num_blocks = {"remote": 8, "local": 8}
+    worker.dst_region_num_blocks = {"remote": [8, 8]}
+    worker.region_num_blocks = [8, 8]
+    worker.region_group_ids = [0, 1]
+    worker.num_regions = 2
+    worker.engine_id = "local"
+    worker._physical_blocks_per_logical_kv_block = 1
+    worker.src_xfer_handles_by_block_size = {64: 200}
+    worker.dst_xfer_side_handles = {"remote": {0: 300}}
+    worker._remote_agents = {"remote": {(0, 0): "agent"}}
+    worker._pending_recv_notifs = {}
+    worker._recving_transfers = {}
+    worker.nixl_wrapper.make_prepped_xfer.return_value = 400
+    stager = MagicMock()
+    worker._maybe_init_host_stager_for_buffers.return_value = stager
+
+    meta = ReqMeta(
+        local_block_ids=([3, 4], [5, 6]),
+        local_physical_block_ids=([3, 4], [5, 6]),
+        tp_size=1,
+        remote=RemoteMeta(
+            block_ids=([1, 2], [1, 2]),
+            host="local-host",
+            port=1234,
+            engine_id="remote",
+            request_id="remote-request",
+        ),
+        hisparse_host_block_ids=[0, 1],
+    )
+    plan = MagicMock(all_source_ranks=(0,), local_consumers=1)
+
+    destination.read_host_blocks(worker, "request", meta, plan, [0, 1])
+
+    stager.submit.assert_called_once()
+    submit_args = stager.submit.call_args.args
+    assert submit_args[0] == "request"
+    np.testing.assert_array_equal(submit_args[1], [10, 11])
+    np.testing.assert_array_equal(submit_args[2], [0, 1])
+    assert submit_args[3] == 300
+    worker.nixl_wrapper.make_prepped_xfer.assert_called_once()
+    worker.nixl_wrapper.transfer.assert_called_once_with(400)
+    assert worker._recving_transfers == {"request": [400]}
+
+
 @pytest.mark.cpu_test
 def test_register_kv_caches_hybrid_mla_dual_purpose_regions():
     """Hybrid MLA+KDA registration: HMA tensors shared by both layer types

vllm/distributed/kv_transfer/kv_connector/v1/hisparse/nixl.py

@@ -33,12 +33,16 @@ def __init__(self, kv_cache_config: KVCacheConfig, vllm_config: VllmConfig) -> N
         self._forward_context = vllm_config.compilation_config.static_forward_context
         self.host_regions: list[tuple[int, int] | None] = []
         self._host_pools: dict[int, tuple[torch.Tensor, list[HiSparseRuntime]]] = {}
+        self._host_desc_lens = np.empty(0, dtype=np.int64)
+        self._host_addrs = np.empty(0, dtype=np.uint64)
         self._descriptor_offsets: list[int | None] = []
         self._xfer_handle: int | None = None
 
     def reset_regions(self) -> None:
         self.host_regions.clear()
         self._host_pools.clear()
+        self._host_desc_lens = np.empty(0, dtype=np.int64)
+        self._host_addrs = np.empty(0, dtype=np.uint64)
         self._descriptor_offsets.clear()
         self._xfer_handle = None
 
@@ -140,6 +144,10 @@ def prepare_host_descriptors(
                 for block_id in range(self.host_num_blocks)
             )
         descs = worker.nixl_wrapper.get_xfer_descs(blocks, "DRAM")
+        self._host_addrs = np.asarray([block[0] for block in blocks], dtype=np.uint64)
+        self._host_desc_lens = np.asarray(
+            [block[1] for block in blocks], dtype=np.int64
+        )
         self._xfer_handle = worker.nixl_wrapper.prep_xfer_dlist(
             "NIXL_INIT_AGENT", descs, backends=worker.nixl_backends
         )
@@ -242,18 +250,25 @@ def read_host_blocks(
         local_device_handle = worker.src_xfer_handles_by_block_size[
             remote_info.remote_block_size
         ]
+        host_local_ids_array = np.asarray(host_local_ids, dtype=np.int64)
+        host_remote_ids_array = np.asarray(host_remote_ids, dtype=np.int64)
+        host_stager = worker._maybe_init_host_stager_for_buffers(
+            meta.remote.host,
+            self._host_desc_lens,
+            self._host_addrs,
+            [pool for pool, _ in self._host_pools.values()],
+        )
         transfer_specs = [
-            (
-                self._xfer_handle,
-                np.asarray(host_local_ids, dtype=np.int64),
-                np.asarray(host_remote_ids, dtype=np.int64),
-            ),
             (
                 local_device_handle,
                 device_local_ids,
                 np.asarray(device_remote_ids, dtype=np.int64),
-            ),
+            )
         ]
+        if host_stager is None:
+            transfer_specs.insert(
+                0, (self._xfer_handle, host_local_ids_array, host_remote_ids_array)
+            )
 
         notif_id = f"{meta.remote.request_id}:{plan.local_consumers}".encode()
         agents = worker._remote_agents[engine_id]
@@ -278,6 +293,13 @@ def read_host_blocks(
                         selected_remote_ids,
                     )
                 )
+            if host_stager is not None and len(host_local_ids_array):
+                host_stager.submit(
+                    request_id,
+                    host_remote_ids_array,
+                    host_local_ids_array,
+                    remote_handle,
+                )
             worker._recving_transfers.setdefault(request_id, [])
             for handle in handles:
                 worker.nixl_wrapper.transfer(handle)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -2814,25 +2814,42 @@ def _handle_heartbeat(self, payload: str) -> None:
 
     def _maybe_init_host_stager(self, remote_host: str) -> HostWriteStager | None:
         """Enable device staging for same-host host-buffer reads."""
-        if remote_host != envs.VLLM_NIXL_SIDE_CHANNEL_HOST:
-            return None
-        if self._host_stager is not None or self._host_stager_init_attempted:
-            return self._host_stager
-        stage_bytes = envs.VLLM_NIXL_HOST_STAGE_BYTES
-        if stage_bytes <= 0 or not self.use_host_buffer:
+        if not self.use_host_buffer:
             return None
-        self._host_stager_init_attempted = True
         desc_lens = np.array(
             [int(entry[1]) for entry in self.src_blocks_data], dtype=np.int64
         )
         host_addrs = np.array(
             [int(entry[0]) for entry in self.src_blocks_data], dtype=np.uint64
         )
+        return self._maybe_init_host_stager_for_buffers(
+            remote_host,
+            desc_lens,
+            host_addrs,
+            list(self.host_xfer_buffers.values()),
+        )
+
+    def _maybe_init_host_stager_for_buffers(
+        self,
+        remote_host: str,
+        desc_lens: np.ndarray,
+        host_addrs: np.ndarray,
+        host_buffers: list[torch.Tensor],
+    ) -> HostWriteStager | None:
+        """Enable device staging for same-host reads into host buffers."""
+        if remote_host != envs.VLLM_NIXL_SIDE_CHANNEL_HOST:
+            return None
+        if self._host_stager is not None or self._host_stager_init_attempted:
+            return self._host_stager
+        stage_bytes = envs.VLLM_NIXL_HOST_STAGE_BYTES
+        if stage_bytes <= 0:
+            return None
+        self._host_stager_init_attempted = True
         try:
             self._host_stager = HostWriteStager(
                 desc_lens=desc_lens,
                 host_addrs=host_addrs,
-                host_buffers=list(self.host_xfer_buffers.values()),
+                host_buffers=host_buffers,
                 device=torch.device(f"cuda:{self.device_id}"),
                 nixl_wrapper=self.nixl_wrapper,
                 memory_type=self.nixl_memory_type,
@@ -2853,8 +2870,7 @@ def _get_finished_host_staging(self) -> set[str]:
         if self._host_stager is None:
             return set()
         done, failed = self._host_stager.get_finished()
-        for req_id in done:
-            self._send_pending_recv_notifs(req_id)
+        done = {req_id for req_id in done if self._finish_recv_component(req_id)}
         for req_id in failed:
             self._log_failure(
                 failure_type="transfer_failed",
@@ -2863,8 +2879,38 @@ def _get_finished_host_staging(self) -> set[str]:
             )
             self._pending_recv_notifs.pop(req_id, None)
             self._handle_failed_transfer(req_id, None)
+        with self._failed_recv_lock:
+            drained_failures = {
+                req_id
+                for req_id in self._failed_recv_pending
+                if not self._recving_transfers.get(req_id)
+                and not self._host_staging_active(req_id)
+            }
+        for req_id in drained_failures:
+            self._finish_recv_component(req_id)
         return done
 
+    def _host_staging_active(self, req_id: str) -> bool:
+        return (
+            self._host_stager is not None and req_id in self._host_stager.active_req_ids
+        )
+
+    def _finish_recv_component(self, req_id: str) -> bool:
+        """Complete a receive once its direct and staged parts are terminal."""
+        if self._recving_transfers.get(req_id) or self._host_staging_active(req_id):
+            return False
+        with self._failed_recv_lock:
+            failed = req_id in self._failed_recv_pending
+            if failed:
+                self._failed_recv_pending.discard(req_id)
+        if failed:
+            self._report_failed_recv(req_id)
+            return False
+        if req_id not in self._recving_metadata:
+            return False
+        self._send_pending_recv_notifs(req_id)
+        return True
+
     def _pop_done_transfers(
         self, transfers: dict[str, list[int]], *, is_recv: bool
     ) -> set[str]:
@@ -2922,18 +2968,9 @@ def _pop_done_transfers(
                 transfers[req_id] = in_progress
                 continue
             del transfers[req_id]
-            if is_recv:
-                with self._failed_recv_lock:
-                    failed = req_id in self._failed_recv_pending
-                    self._failed_recv_pending.discard(req_id)
-                if failed:
-                    self._report_failed_recv(req_id)
-                    continue
-                if req_id not in self._recving_metadata:
-                    continue
+            if is_recv and not self._finish_recv_component(req_id):
+                continue
             done_req_ids.add(req_id)
-            if is_recv:
-                self._send_pending_recv_notifs(req_id)
         return done_req_ids
 
     def _send_pending_recv_notifs(self, req_id: str) -> None:
@@ -2960,7 +2997,7 @@ def _handle_failed_transfer(self, req_id: str, handle: int | None):
         self.xfer_stats.record_failed_transfer()
         if self._host_stager is not None:
             self._host_stager.abort(req_id)
-            if req_id in self._host_stager.active_req_ids:
+            if self._host_staging_active(req_id):
                 with self._failed_recv_lock:
                     self._failed_recv_pending.add(req_id)
                 return