From d8fece2faa4911b8b0b81f9c010971cb8bc65cab Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 25 Oct 2026 17:30:42 +0100 Subject: [PATCH 1/3] [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 From 38a0d6cd5bb156697c1d22360cfa6da10fae44aa Mon Sep 17 00:00:00 2001 From: 0z5a Date: Thu, 5 Nov 2026 13:52:35 +0100 Subject: [PATCH 2/3] [PD] Create the UCX backend with error handling off, no switch Review follow-up: the error-handling mode is not a user-facing option. When the configured backend is UCX, create it with ucx_error_handling_mode="none" and drop nixl_backend_params() and VLLM_RLT_NIXL_UCX_ERROR_HANDLING. NIXL's default (peer) mode requires a peer failure handler, which only the network transports provide, so UCX leaves the shared-memory lanes out of intra-node RMA. Peer loss is instead reported by the P/D control channel, which owns worker liveness: StageClient.poll() notices an exited worker and raises, and PDEngine._terminate() stops and joins both roles and closes both channels. New tests: the UCX backend is created with the mode for UCX and left alone for other backends, and the control channel detects an exited worker and reaps every peer. On GH200, killing either worker with transfers outstanding returns control to the coordinator in ~0.5 s (the transfer timeout is 90 s) and leaves no child, /dev/shm or GPU process behind. Tests on GH200 (NIXL 1.4.1, UCX 1.22, torch 2.14.1+cu130): CPU suite 867 passed, 342 skipped; ruff passes. The PD GPU suite needs one prefill and one decode device and the GH200 nodes here expose one GPU, so the PD path was run with both engines on GPU 0: outputs matched a local engine (86,208 bytes over 24 transfers). Signed-off-by: 0z5a --- tests/test_pd.py | 104 +++++++++++++++++++++++++++++++++++---- vllm_rlt/pd/transport.py | 32 +++--------- 2 files changed, 102 insertions(+), 34 deletions(-) diff --git a/tests/test_pd.py b/tests/test_pd.py index 521d01a3..2dd42979 100644 --- a/tests/test_pd.py +++ b/tests/test_pd.py @@ -476,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 @@ -751,14 +790,59 @@ def test_pd_reservations_are_reported_before_worker_acknowledgments(): 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 +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 - 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") + 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/pd/transport.py b/vllm_rlt/pd/transport.py index 4fc97f5c..092bb47f 100644 --- a/vllm_rlt/pd/transport.py +++ b/vllm_rlt/pd/transport.py @@ -95,36 +95,20 @@ 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 - params = nixl_backend_params(config.backend) - if params: - # nixl_agent_config cannot pass backend init parameters; create it explicitly. + 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, params) + 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) From 699b6add796588febf4f0e589dc795ee2f89ea16 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Thu, 5 Nov 2026 14:40:19 +0100 Subject: [PATCH 3/3] [PD] Keep stats() focused on requests and used blocks Review follow-up: num_blocks, prefix_hit_tokens and prefix_queries are unrelated to the transfer changes. Restore the previous stats() payload, and the two test assertions in tests/test_stage_engine.py that compare it verbatim. The prefix-cache counters themselves stay in KVCacheManager; they are not part of this PR. Tests on GH200: ruff passes; CPU suite 867 passed, 342 skipped. The affected test files pass on their own (118 passed, 9 skipped). Signed-off-by: 0z5a --- tests/test_stage_engine.py | 8 ++------ vllm_rlt/engine/stage_engine.py | 7 +------ 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/tests/test_stage_engine.py b/tests/test_stage_engine.py index ae10da7b..6290592a 100644 --- a/tests/test_stage_engine.py +++ b/tests/test_stage_engine.py @@ -356,8 +356,7 @@ def test_stage_client_owns_process_and_isolates_cancelled_generations(): finally: client.close() assert not client.process.is_alive() - stats = client.stats() - assert (stats["requests"], stats["used_blocks"]) == (0, 0) + assert client.stats() == dict(requests=0, used_blocks=0) def test_stage_client_reports_worker_failure(): @@ -534,10 +533,7 @@ 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"]["requests"] == 0 and r["result"]["used_blocks"] == 0 - for r in replies.values() - ) + assert all(r["result"] == dict(requests=0, 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 97e6d8e4..40a6c9a8 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -346,13 +346,8 @@ 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=cache.num_used_blocks, - num_blocks=cache.num_blocks, - prefix_hit_tokens=cache.prefix_hits, - prefix_queries=cache.prefix_queries, + requests=len(self.scheduler.requests), used_blocks=self.cache_manager.num_used_blocks ) def has_unfinished_requests(self) -> bool: