[PD] Coalesced transfers, UCX error-handling mode, larger default budget

tests/test_pd.py

@@ -89,6 +89,32 @@ def test_segment_ranges_cover_only_valid_tokens_and_all_layers_depths():
     parts = list(partition_segments(segments, segments, 333, 3))
     assert sum(n for _, _, n in parts) == len(expected)
     assert all(n <= 333 and len(src) <= 3 and src == dst for src, dst, n in parts)
+    covered = set()
+    for src, _, _ in parts:
+        for address, size, _ in src:
+            covered.update(range(address, address + size))
+    assert covered == expected
+
+
+def test_partition_merges_segments_contiguous_on_both_sides():
+    # Three adjacent 100-byte blocks on both pools collapse into one range; a target
+    # that is not adjacent at the same boundary keeps its own descriptor.
+    source = [(1000, 100, 0), (1100, 100, 0), (1200, 100, 0), (5000, 50, 0)]
+    target = [(7000, 100, 1), (7100, 100, 1), (9000, 100, 1), (9100, 50, 1)]
+    parts = list(partition_segments(source, target, 10**6, 256))
+    assert parts == [
+        (
+            [(1000, 200, 0), (1200, 100, 0), (5000, 50, 0)],
+            [(7000, 200, 1), (9000, 100, 1), (9100, 50, 1)],
+            350,
+        )
+    ]
+    # Byte limits still split a merged range, descriptor limits still bound a chunk.
+    parts = list(partition_segments(source, target, 120, 256))
+    assert [n for _, _, n in parts] == [120, 120, 110]
+    assert sum(len(src) for src, _, _ in parts) == 5
+    with pytest.raises(ValueError):
+        list(partition_segments(source, target[:-1], 10**6, 256))
 
 
 @pytest.mark.parametrize(
@@ -105,6 +131,16 @@ def test_pd_configuration_validation(options):
         PDConfig(**options)
 
 
+def test_pd_default_budget_keeps_a_whole_long_prompt_in_flight():
+    """Chunks are submitted and acknowledged only between engine steps, so every budget
+    round a request needs adds steps to its commit tail. The default budget holds the
+    KV of a 2048-token Ouro-1.4B prompt (24 layers x 4 depths x K/V x 2048 bf16 values)."""
+    config = PDConfig()
+    ouro_kv_bytes_per_token = 24 * 4 * 2 * 2048 * 2
+    assert config.max_inflight_bytes >= 2048 * ouro_kv_bytes_per_token
+    assert config.max_inflight_bytes % config.transfer_chunk_bytes == 0
+
+
 def test_pd_rejects_huginn_before_spawning(monkeypatch):
     from unittest.mock import Mock
 
@@ -440,9 +476,48 @@ def test_pd_multiple_workers_and_peer_failure():
         peer = next(p for p in e.peers.values() if p.role == "prefill")
         peer.client.process.kill()
         peer.client.process.join()
+        started = time.monotonic()
         with pytest.raises(RuntimeError):
             e.step()
+        # Liveness, not a missing transfer error, ends the deployment: UCX runs
+        # without a peer failure handler and the control channel detects the exit.
+        assert time.monotonic() - started < 30
         assert e.closed and all(not p.client.process.is_alive() for p in e.peers.values())
+        assert all(p.client.closed for p in e.peers.values())
+
+
+def test_control_channel_detects_a_worker_exit_and_reaps_every_peer():
+    from vllm_rlt.engine.llm_engine import StageClient
+    from vllm_rlt.engine.stage_config import StageConfig
+    from vllm_rlt.pd.engine import PDEngine, Peer
+
+    engine = PDEngine.__new__(PDEngine)
+    engine._profile_session = None
+    engine.failure = None
+    engine.peers = {}
+    for role in ("prefill", "decode"):
+        client = StageClient(
+            tiny_ouro_config(head_dim=64),
+            device="cpu",
+            stage_config=StageConfig(stage_id=role, role=role),
+        )
+        engine.peers[role] = Peer(role, role, client)
+    try:
+        engine.peers["prefill"].client.process.kill()
+        engine.peers["prefill"].client.process.join()
+        started = time.monotonic()
+        with pytest.raises(RuntimeError) as error:
+            engine._poll()
+        assert "prefill" in str(error.value) or "control channel" in str(error.value)
+        assert time.monotonic() - started < 5
+        engine._terminate()
+    finally:
+        for peer in engine.peers.values():
+            if not peer.client.closed:
+                peer.client.join(force=True)
+    assert engine.closed
+    assert all(p.client.closed for p in engine.peers.values())
+    assert all(not p.client.process.is_alive() for p in engine.peers.values())
 
 
 @pytest.mark.gpu
@@ -713,3 +788,61 @@ def test_pd_reservations_are_reported_before_worker_acknowledgments():
     assert engine.stats() == dict(requests=1, reserved_blocks=20)
     assert len(messages) == 1 and messages[0][0] == ("decode", "reserve")
     assert not hasattr(engine, "cache_manager") and not hasattr(engine, "scheduler")
+
+
+def fake_nixl(monkeypatch):
+    """Install a fake NIXL agent package and record how the backend was created."""
+    import sys
+    import types
+
+    created = {}
+
+    class Agent:
+        def __init__(self, name, config=None):
+            created["agent_config"] = config
+
+        def create_backend(self, backend, params):
+            created["backend"] = (backend, dict(params))
+
+        def register_memory(self, tensor):
+            return tensor
+
+        def get_agent_metadata(self):
+            return b"metadata"
+
+    def agent_config(**kwargs):
+        return kwargs
+
+    package = types.ModuleType("nixl")
+    package.__path__ = []
+    api = types.ModuleType("nixl._api")
+    api.nixl_agent = Agent
+    api.nixl_agent_config = agent_config
+    monkeypatch.setitem(sys.modules, "nixl", package)
+    monkeypatch.setitem(sys.modules, "nixl._api", api)
+    return created
+
+
+def test_ucx_backend_is_created_with_error_handling_none(monkeypatch):
+    from vllm_rlt.pd.transport import NixlConnector
+
+    created = fake_nixl(monkeypatch)
+    # No hidden switch: an obsolete environment variable cannot change the mode.
+    monkeypatch.setenv("VLLM_RLT_NIXL_UCX_ERROR_HANDLING", "peer")
+    cache = KVCacheManager(3, 2, 8, 24, 4, 4)
+    connector = NixlConnector("stage", cache, torch.zeros(2, 16), PDConfig())
+    assert created["agent_config"] == {"backends": []}
+    assert created["backend"] == ("UCX", {"ucx_error_handling_mode": "none"})
+    assert connector.info["hidden_bytes"] == 64
+    element = cache.key_cache.element_size()
+    assert connector.info["block_bytes"] == cache.key_cache.stride(0) * element
+
+
+def test_non_ucx_backend_is_not_created_with_ucx_parameters(monkeypatch):
+    from vllm_rlt.pd.transport import NixlConnector
+
+    created = fake_nixl(monkeypatch)
+    cache = KVCacheManager(3, 2, 8, 24, 4, 4)
+    NixlConnector("stage", cache, torch.zeros(2, 16), PDConfig(backend="GDS"))
+    assert created["agent_config"] == {"backends": ["GDS"]}
+    assert "backend" not in created

vllm_rlt/entrypoints/pd_serve.py

@@ -120,8 +120,10 @@ def main():
     parser.add_argument("--pd-max-receiving-requests", type=int, default=32)
     parser.add_argument("--pd-max-draining-requests", type=int, default=8)
     parser.add_argument("--nixl-backend", default="UCX")
-    parser.add_argument("--pd-transfer-chunk-bytes", type=int, default=64 * 1024**2)
-    parser.add_argument("--pd-max-inflight-bytes", type=int, default=256 * 1024**2)
+    parser.add_argument(
+        "--pd-transfer-chunk-bytes", type=int, default=PDConfig.transfer_chunk_bytes
+    )
+    parser.add_argument("--pd-max-inflight-bytes", type=int, default=PDConfig.max_inflight_bytes)
     parser.add_argument("--pd-max-transfer-descriptors", type=int, default=256)
     parser.add_argument("--pd-startup-timeout", type=float, default=300)
     parser.add_argument("--host", default="127.0.0.1")

vllm_rlt/pd/config.py

@@ -7,8 +7,14 @@ class PDConfig:
     prefill_devices: tuple[int, ...] = (0,)
     decode_devices: tuple[int, ...] = (1,)
     max_pending_requests: int = 256
-    transfer_chunk_bytes: int = 64 * 1024**2
-    max_inflight_bytes: int = 256 * 1024**2
+    # KV chunks are submitted, and their completion observed, only between engine
+    # steps, so a request whose KV exceeds the in-flight budget waits extra steps
+    # after its prefill (the commit tail). 2 GiB covers one 2048-token Ouro-1.4B
+    # request (1.6 GB): on a GH200 node the tail fell from ~1 s to 17 ms and 1P1D
+    # throughput rose 1.7x. In-flight bytes are not a buffer: NIXL writes straight
+    # from the registered prefill pool into the decode pool.
+    transfer_chunk_bytes: int = 256 * 1024**2
+    max_inflight_bytes: int = 2 * 1024**3
     max_transfer_descriptors: int = 256
     max_control_messages: int = 32
     startup_timeout: float = 300.0

vllm_rlt/pd/connector.py

@@ -30,6 +30,8 @@ class RequestTransfer:
     transfer_boundary_first: tuple[int, int] | None = None
     transfer_boundary_last: tuple[int, int] | None = None
     transfer_overlaps: int = 0
+    # A chunk of this transfer went on the wire since the last prefill submission.
+    submitted_since_compute: bool = False
     received: set = field(default_factory=set)
     next_sequence: int = 0
     expected: int | None = None
@@ -191,7 +193,12 @@ def prefill_submitted(self, request, *, phase, start, progress, event, stream):
             phase == RequestPhase.PREFILLING_RECURRENT
             and self.engine.scheduler.config.wavefront_chunked_prefill
         ):
-            transfer.transfer_overlaps += self.transport.busy(transfer.tid)
+            # Overlap: the transfer ran during this compute step, either still on the
+            # wire now or sent and already drained (coalesced descriptors drain fast).
+            transfer.transfer_overlaps += bool(
+                self.transport.busy(transfer.tid) or transfer.submitted_since_compute
+            )
+        transfer.submitted_since_compute = False
         if not progress.chunk_complete:
             transfer.last_event = event
             return
@@ -388,23 +395,9 @@ def progress(self):
                         ),
                         overlaps=transfer.transfer_overlaps,
                     )
-                # Round robin: at most one bounded transfer per request per tick.
-                if transfer.chunks:
-                    event, seq, local, remote, size = transfer.chunks[0]
-                    if event.query() and self.transport.submit(
-                        transfer.tid, seq, transfer.peer, local, remote, size
-                    ):
-                        transfer.chunks.popleft()
-                if transfer.compute_done and not transfer.chunks and not transfer.committed:
-                    transfer.committed = True
-                    self.notify("commit", tid=transfer.tid, count=transfer.next_sequence)
-                if (
-                    transfer.acknowledged
-                    and transfer.committed
-                    and transfer.request.compute_released
-                    and not self.transport.busy(transfer.tid)
-                ):
-                    self._remove(transfer)
+                # Round robin: one bounded transfer per request per pass; the
+                # in-flight budget is filled with further passes below.
+                self._submit_chunk(transfer)
             elif transfer.cancelled:
                 if transfer.ready or transfer.cancel_confirmed:
                     self._remove(transfer)
@@ -428,3 +421,35 @@ def progress(self):
                         raise RuntimeError("stale import activation")
                     transfer.ready = True
                     self.notify("received", tid=transfer.tid)
+        # Keep the link busy: as long as the in-flight budget admits another chunk,
+        # take one more pass over the outgoing transfers (fairness is preserved by
+        # the pass order). Then commit transfers whose last chunk has gone out.
+        outgoing = [t for t in transfers if t.outgoing and not t.cancelled]
+        while outgoing and any([self._submit_chunk(t) for t in outgoing]):
+            outgoing = [t for t in outgoing if t.chunks]
+        for transfer in transfers:
+            if not transfer.outgoing or transfer.cancelled:
+                continue
+            if transfer.compute_done and not transfer.chunks and not transfer.committed:
+                transfer.committed = True
+                self.notify("commit", tid=transfer.tid, count=transfer.next_sequence)
+            if (
+                transfer.acknowledged
+                and transfer.committed
+                and transfer.request.compute_released
+                and not self.transport.busy(transfer.tid)
+            ):
+                self._remove(transfer)
+
+    def _submit_chunk(self, transfer):
+        """Submit the transfer's next ready chunk; True when one went out."""
+        if not transfer.chunks:
+            return False
+        event, seq, local, remote, size = transfer.chunks[0]
+        if event.query() and self.transport.submit(
+            transfer.tid, seq, transfer.peer, local, remote, size
+        ):
+            transfer.chunks.popleft()
+            transfer.submitted_since_compute = True
+            return True
+        return False

vllm_rlt/pd/transport.py

@@ -46,19 +46,46 @@ def kv_segments(info, tables, start, end):
                         )
 
 
-def partition_segments(source, target, byte_limit, descriptor_limit):
-    """Bound both transfer size and descriptor count without an intermediate KV copy."""
+def coalesce_segments(source, target):
+    """Merge consecutive segment pairs that are contiguous on both the source and target.
+
+    Physical blocks handed out in order are adjacent in both pools, so a prompt's
+    per-block descriptors collapse into a few long ranges. Merging never moves a
+    byte: each merged pair covers exactly the bytes of its members, in order.
+    """
     from itertools import zip_longest
 
-    local, remote, size = [], [], 0
+    current = None
     for src, dst in zip_longest(source, target):
         if src is None or dst is None or src[1] != dst[1]:
             raise ValueError("incompatible source/target KV geometry")
+        if (
+            current is not None
+            and current[0] + current[2] == src[0]
+            and current[1] + current[2] == dst[0]
+            and current[3] == src[2]
+            and current[4] == dst[2]
+        ):
+            current[2] += src[1]
+            continue
+        if current is not None:
+            yield tuple(current)
+        current = [src[0], dst[0], src[1], src[2], dst[2]]
+    if current is not None:
+        yield tuple(current)
+
+
+def partition_segments(source, target, byte_limit, descriptor_limit):
+    """Bound both transfer size and descriptor count without an intermediate KV copy."""
+    local, remote, size = [], [], 0
+    for src_address, dst_address, length, src_device, dst_device in coalesce_segments(
+        source, target
+    ):
         offset = 0
-        while offset < src[1]:
-            amount = min(src[1] - offset, byte_limit - size)
-            local.append((src[0] + offset, amount, src[2]))
-            remote.append((dst[0] + offset, amount, dst[2]))
+        while offset < length:
+            amount = min(length - offset, byte_limit - size)
+            local.append((src_address + offset, amount, src_device))
+            remote.append((dst_address + offset, amount, dst_device))
             size += amount
             offset += amount
             if size == byte_limit or len(local) == descriptor_limit:
@@ -73,7 +100,17 @@ def __init__(self, name, cache, hidden, config):
         from nixl._api import nixl_agent, nixl_agent_config
 
         self.config = config
-        self.agent = nixl_agent(name, nixl_agent_config(backends=[config.backend]))
+        if config.backend == "UCX":
+            # NIXL's default (peer) error handling requires a peer failure handler,
+            # which only the network transports provide, so UCX leaves the
+            # shared-memory lanes out of intra-node RMA. A lost worker is reported
+            # by the control channel, which owns worker liveness, so transfers do
+            # not need a handler of their own. nixl_agent_config cannot pass
+            # backend init parameters, so the backend is created explicitly.
+            self.agent = nixl_agent(name, nixl_agent_config(backends=[]))
+            self.agent.create_backend(config.backend, {"ucx_error_handling_mode": "none"})
+        else:
+            self.agent = nixl_agent(name, nixl_agent_config(backends=[config.backend]))
         self.storage = (cache.key_cache, cache.value_cache, hidden)
         self.registrations = [self.agent.register_memory(t) for t in self.storage]
         k = cache.key_cache