diff --git a/tests/test_pd.py b/tests/test_pd.py index 8cf39195..521d01a3 100644 --- a/tests/test_pd.py +++ b/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") diff --git a/tests/test_stage_engine.py b/tests/test_stage_engine.py index 6290592a..ae10da7b 100644 --- a/tests/test_stage_engine.py +++ b/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() diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 40a6c9a8..97e6d8e4 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/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: diff --git a/vllm_rlt/entrypoints/pd_serve.py b/vllm_rlt/entrypoints/pd_serve.py index 92c3ef48..34977287 100644 --- a/vllm_rlt/entrypoints/pd_serve.py +++ b/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") diff --git a/vllm_rlt/pd/config.py b/vllm_rlt/pd/config.py index 4d0c5f14..ade43375 100644 --- a/vllm_rlt/pd/config.py +++ b/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 diff --git a/vllm_rlt/pd/connector.py b/vllm_rlt/pd/connector.py index 121a6538..9b5ed676 100644 --- a/vllm_rlt/pd/connector.py +++ b/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 diff --git a/vllm_rlt/pd/transport.py b/vllm_rlt/pd/transport.py index fca044a1..4fc97f5c 100644 --- a/vllm_rlt/pd/transport.py +++ b/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