diff --git a/tests/test_pd.py b/tests/test_pd.py index 8cf39195..2dd42979 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 @@ -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 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..092bb47f 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: @@ -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