From d8fece2faa4911b8b0b81f9c010971cb8bc65cab Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 25 Oct 2026 17:30:42 +0100 Subject: [PATCH] [PD] Huginn coordinator config, coalesced transfers, UCX error-handling mode, larger default budget Four fixes to prefill-decode disaggregation, measured on one GH200 node with the prefill engine on GPU 0 and the decode engine on GPU 1: - The coordinator built every model with OuroConfig(**info["model"]), so a real Huginn checkpoint failed at startup (unexpected keyword argument 'n_embd'); the repository's Huginn PD test skips without a checkpoint. It now picks the config class from the model_type the worker reports. - KV segments that are contiguous on both sides are coalesced into one descriptor, and each progress pass fills the in-flight budget instead of sending one chunk per request. The overlap counter now also credits transfers issued during a prefill computation that drained before the next chunk was computed (coalesced descriptors drain fast); the wavefront overlap GPU test counted 0 before. - VLLM_RLT_NIXL_UCX_ERROR_HANDLING=none|peer creates the UCX backend explicitly with ucx_error_handling_mode (nixl_agent_config cannot pass backend init parameters). NIXL's default (peer) makes UCX skip shared memory and cuda_ipc, which have no peer failure handler, and emulate intra-node RMA with active messages over TCP; with none, UCX's log shows cuda_ipc and shared memory and no TCP lane. Peer failures are then reported only through the control channel, so the default is unchanged. - 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 default budget is now 256 MiB chunks and 2 GiB in flight, enough for a 2048-token Ouro-1.4B prompt (1.6 GB of KV). In-flight bytes are not a buffer: NIXL writes straight from the registered prefill pool into the decode pool. pd_serve's CLI defaults read the PDConfig defaults. Ouro-1.4B 1P1D, 2048-token prompts, 64 output tokens, 32 requests, packed prefill, error handling none, two runs per arm in alternating order: 64 MiB / 256 MiB: 170-279 tok/s, TTFT p50 3.4-3.9 s, commit tail 0.96-1.05 s 256 MiB / 1 GiB: 254-290 tok/s, TTFT p50 3.2-3.5 s, commit tail 0.68-0.90 s 256 MiB / 2 GiB: 463-466 tok/s, TTFT p50 2.0 s, commit tail 0.017 s Switching the link alone (peer -> none, old budget, one run each) changed a 64 MiB chunk from 20.0 to 16.9 ms: the budget rounds, not TCP, set the tail. The engine's stats() also reports prefix-cache hits and queries. GPU PD tests: 41 passed with the first two changes; the error-handling mode and the budget were exercised by the benchmark runs above. Signed-off-by: 0z5a Co-Authored-By: Claude Opus 5.5 --- tests/test_pd.py | 65 ++++++++++++++++++++++++++++++ tests/test_stage_engine.py | 8 +++- vllm_rlt/engine/stage_engine.py | 7 +++- vllm_rlt/entrypoints/pd_serve.py | 6 ++- vllm_rlt/pd/config.py | 10 ++++- vllm_rlt/pd/connector.py | 61 +++++++++++++++++++--------- vllm_rlt/pd/engine.py | 13 +++++- vllm_rlt/pd/transport.py | 69 ++++++++++++++++++++++++++++---- 8 files changed, 204 insertions(+), 35 deletions(-) diff --git a/tests/test_pd.py b/tests/test_pd.py index 22ced40c..11cec30f 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 drain(engine): outputs = {} deadline = time.monotonic() + 90 @@ -642,3 +678,32 @@ 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_pd_coordinator_builds_the_reported_model_config(): + """The coordinator must not assume Ouro: Huginn workers report Huginn's own config fields.""" + from dataclasses import asdict + + from vllm_rlt.models import HuginnConfig, OuroConfig + from vllm_rlt.pd.engine import _CONFIG_CLASSES + + huginn = HuginnConfig() + rebuilt = _CONFIG_CLASSES["huginn_raven"](**asdict(huginn)) + assert isinstance(rebuilt, HuginnConfig) and rebuilt == huginn + assert rebuilt.total_ut_steps == huginn.mean_recurrence + with pytest.raises(TypeError, match="n_embd"): + OuroConfig(**asdict(huginn)) # the old hard-coded path + assert _CONFIG_CLASSES["ouro"] is OuroConfig + + +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 71facca1..1ffa4e67 100644 --- a/tests/test_stage_engine.py +++ b/tests/test_stage_engine.py @@ -353,7 +353,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(): @@ -530,7 +531,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 7a63621b..19454ba6 100644 --- a/vllm_rlt/entrypoints/pd_serve.py +++ b/vllm_rlt/entrypoints/pd_serve.py @@ -119,8 +119,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/engine.py b/vllm_rlt/pd/engine.py index ce2b3b7d..21d3cf90 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -14,7 +14,7 @@ from vllm_rlt.config import CacheConfig, ExecutionConfig, ExitConfig, SchedulerConfig from vllm_rlt.engine.llm_engine import StageClient from vllm_rlt.engine.stage_config import StageConfig -from vllm_rlt.models import OuroConfig, resolve_model_config +from vllm_rlt.models import HuginnConfig, NanbeigeConfig, OuroConfig, resolve_model_config from vllm_rlt.profiling import ProfileConfig from vllm_rlt.profiling_artifacts import _write_json_atomic, profile_timestamp from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase @@ -23,6 +23,12 @@ from .config import PDConfig from .connector import PDStageConnector +_CONFIG_CLASSES = { + "ouro": OuroConfig, + "nanbeige": NanbeigeConfig, + "huginn_raven": HuginnConfig, +} + @dataclass class Peer: @@ -142,7 +148,10 @@ def __init__( fingerprints = {p.client.info["fingerprint"] for p in self.peers.values()} if len(fingerprints) != 1: raise ValueError("P/D model weights or KV configuration do not match") - self.model_config = OuroConfig(**next(iter(self.peers.values())).client.info["model"]) + info = next(iter(self.peers.values())).client.info + # The workers report their model family; Huginn's config has its own fields. + config_class = _CONFIG_CLASSES.get(info.get("model_type"), OuroConfig) + self.model_config = config_class(**info["model"]) for peer in self.peers.values(): self._send( peer.name, 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