[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
 
@@ -713,3 +749,16 @@ 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 test_nixl_backend_params_from_environment(monkeypatch):
+    from vllm_rlt.pd.transport import nixl_backend_params
+
+    monkeypatch.delenv("VLLM_RLT_NIXL_UCX_ERROR_HANDLING", raising=False)
+    assert nixl_backend_params("UCX") == {}
+    monkeypatch.setenv("VLLM_RLT_NIXL_UCX_ERROR_HANDLING", "none")
+    assert nixl_backend_params("UCX") == {"ucx_error_handling_mode": "none"}
+    assert nixl_backend_params("GDS") == {}
+    monkeypatch.setenv("VLLM_RLT_NIXL_UCX_ERROR_HANDLING", "off")
+    with pytest.raises(ValueError, match="none"):
+        nixl_backend_params("UCX")

tests/test_stage_engine.py

@@ -356,7 +356,8 @@ def test_stage_client_owns_process_and_isolates_cancelled_generations():
     finally:
         client.close()
     assert not client.process.is_alive()
-    assert client.stats() == dict(requests=0, used_blocks=0)
+    stats = client.stats()
+    assert (stats["requests"], stats["used_blocks"]) == (0, 0)
 
 
 def test_stage_client_reports_worker_failure():
@@ -533,7 +534,10 @@ def test_stage_client_controls_can_be_submitted_before_collecting():
                 if reply is not None:
                     replies[control_id] = reply
             time.sleep(0.001)
-        assert all(r["result"] == dict(requests=0, used_blocks=0) for r in replies.values())
+        assert all(
+            r["result"]["requests"] == 0 and r["result"]["used_blocks"] == 0
+            for r in replies.values()
+        )
         assert not client.pending_controls and not client.replies
     finally:
         client.close()

vllm_rlt/engine/stage_engine.py

@@ -346,8 +346,13 @@ def has_request(self, request_id):
         return request is not None and request.phase != RequestPhase.FINISHED
 
     def stats(self):
+        cache = self.cache_manager
         return dict(
-            requests=len(self.scheduler.requests), used_blocks=self.cache_manager.num_used_blocks
+            requests=len(self.scheduler.requests),
+            used_blocks=cache.num_used_blocks,
+            num_blocks=cache.num_blocks,
+            prefix_hit_tokens=cache.prefix_hits,
+            prefix_queries=cache.prefix_queries,
         )
 
     def has_unfinished_requests(self) -> bool:

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:
@@ -68,12 +95,38 @@ def partition_segments(source, target, byte_limit, descriptor_limit):
         yield local, remote, size
 
 
+def nixl_backend_params(backend):
+    """Optional NIXL backend init parameters from the environment.
+
+    VLLM_RLT_NIXL_UCX_ERROR_HANDLING=none|peer sets the UCX plugin's
+    ``ucx_error_handling_mode``. With the default (peer) UCX cannot use shared
+    memory or cuda_ipc lanes, which have no peer failure handler, and emulates
+    intra-node RMA with active messages over TCP (measured 1.3-5 GB/s on a GH200
+    node). ``none`` lets UCX select cuda_ipc for device-memory WRITEs, but a peer
+    failure is then not reported as a transfer error.
+    """
+    import os
+
+    mode = os.environ.get("VLLM_RLT_NIXL_UCX_ERROR_HANDLING")
+    if backend != "UCX" or not mode:
+        return {}
+    if mode not in ("none", "peer"):
+        raise ValueError("VLLM_RLT_NIXL_UCX_ERROR_HANDLING must be 'none' or 'peer'")
+    return {"ucx_error_handling_mode": mode}
+
+
 class NixlConnector:
     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]))
+        params = nixl_backend_params(config.backend)
+        if params:
+            # nixl_agent_config cannot pass backend init parameters; create it explicitly.
+            self.agent = nixl_agent(name, nixl_agent_config(backends=[]))
+            self.agent.create_backend(config.backend, params)
+        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