iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
50dabc9989bc9d7b5625f8256196d9935de28ade
vLLM-4x-RTX-3090/patches/0031-ple_offload-carry-the-host-mapped-pull-transport-ont.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

7872 lines
285 KiB
Diff
Raw Blame History

From 4c793b234fff0b909893dfdbcccd78b98d6fc4e0 Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Wed, 23 Sep 2026 00:01:58 -0400
Subject: [PATCH 031/103] ple_offload: carry the host-mapped pull transport
onto v0.30.0
Restores the fork's PLE offload package (not upstream) and its pull
transport: the offload process holds no CUDA context; the captured op waits on
a device-mapped host flag and copies host->GPU on the model stream; publication
is fenced; releases are synchronous across dummy/capture work; failures close
the rank. Includes the test-only fault-injection cuts for the transport.
---
probe/faultcut_gate.py | 46 +
tests/v1/ple_offload/test_correspondence.py | 197 +++
tests/v1/ple_offload/test_fatal_rpc.py | 145 ++
.../test_fault_bootstrap_disarmed.py | 80 +
.../ple_offload/test_fault_injector_basics.py | 260 ++++
.../ple_offload/test_fault_injector_gates.py | 468 ++++++
.../v1/ple_offload/test_fault_wire_frames.py | 96 ++
.../v1/ple_offload/test_ipc_socket_unlink.py | 133 ++
tests/v1/ple_offload/test_pending_rpc.py | 81 +
.../v1/ple_offload/test_progress_deadline.py | 216 +++
tests/v1/ple_offload/test_shm_unlink.py | 397 +++++
tests/v1/ple_offload/test_shutdown_order.py | 171 +++
.../ple_offload/test_spawn_bootstrap_cuda.py | 352 +++++
tests/v1/ple_offload/test_test_fault.py | 293 ++++
.../ple_offload/test_transport_bounded_io.py | 343 +++++
.../test_transport_fatal_helper.py | 387 +++++
.../layers/ple_offload_layer.py | 433 ++++++
vllm/v1/executor/multiproc_executor.py | 31 +-
vllm/v1/ple_offload/__init__.py | 2 +
vllm/v1/ple_offload/connector.py | 855 +++++++++++
vllm/v1/ple_offload/fatal.py | 153 ++
vllm/v1/ple_offload/progress.py | 314 ++++
vllm/v1/ple_offload/protocol.py | 50 +
vllm/v1/ple_offload/worker.py | 1298 +++++++++++++++++
vllm/v1/worker/gpu/cudagraph_utils.py | 140 +-
vllm/v1/worker/gpu/model_runner.py | 140 ++
vllm/v1/worker/gpu_worker.py | 145 ++
vllm_ple_offload_bootstrap/__init__.py | 210 +++
vllm_ple_offload_bootstrap/__main__.py | 7 +
29 files changed, 7429 insertions(+), 14 deletions(-)
create mode 100755 probe/faultcut_gate.py
create mode 100644 tests/v1/ple_offload/test_correspondence.py
create mode 100644 tests/v1/ple_offload/test_fatal_rpc.py
create mode 100644 tests/v1/ple_offload/test_fault_bootstrap_disarmed.py
create mode 100644 tests/v1/ple_offload/test_fault_injector_basics.py
create mode 100644 tests/v1/ple_offload/test_fault_injector_gates.py
create mode 100644 tests/v1/ple_offload/test_fault_wire_frames.py
create mode 100644 tests/v1/ple_offload/test_ipc_socket_unlink.py
create mode 100644 tests/v1/ple_offload/test_pending_rpc.py
create mode 100644 tests/v1/ple_offload/test_progress_deadline.py
create mode 100644 tests/v1/ple_offload/test_shm_unlink.py
create mode 100644 tests/v1/ple_offload/test_shutdown_order.py
create mode 100644 tests/v1/ple_offload/test_spawn_bootstrap_cuda.py
create mode 100644 tests/v1/ple_offload/test_test_fault.py
create mode 100644 tests/v1/ple_offload/test_transport_bounded_io.py
create mode 100644 tests/v1/ple_offload/test_transport_fatal_helper.py
create mode 100644 vllm/model_executor/layers/ple_offload_layer.py
create mode 100644 vllm/v1/ple_offload/__init__.py
create mode 100644 vllm/v1/ple_offload/connector.py
create mode 100644 vllm/v1/ple_offload/fatal.py
create mode 100644 vllm/v1/ple_offload/progress.py
create mode 100644 vllm/v1/ple_offload/protocol.py
create mode 100644 vllm/v1/ple_offload/worker.py
create mode 100644 vllm_ple_offload_bootstrap/__init__.py
create mode 100644 vllm_ple_offload_bootstrap/__main__.py
diff --git a/probe/faultcut_gate.py b/probe/faultcut_gate.py
new file mode 100755
index 0000000..3eb74b5
--- /dev/null
+++ b/probe/faultcut_gate.py
@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+"""CPU-testable validation for external gates excluded from GPU campaigns."""
+
+from collections.abc import Callable
+
+
+def validate_partial_send_gate(
+ *,
+ pid: int,
+ expected_start: int,
+ observed_start: int,
+ alive: bool,
+ bytes_sent: int,
+ frame_size: int,
+ completed: bool,
+ child_not_reading: bool,
+) -> tuple[bool, str]:
+ if pid <= 1 or not alive:
+ return False, "target is stale or not alive"
+ if expected_start <= 0 or observed_start != expected_start:
+ return False, "target PID was reused"
+ if frame_size <= 0 or not 0 < bytes_sent < frame_size:
+ return False, "send is zero-byte or complete"
+ if completed:
+ return False, "send already completed"
+ if not child_not_reading:
+ return False, "child non-reading state was not acknowledged"
+ return True, "valid partial send"
+
+
+def deliver_checked_signal(
+ *,
+ pid: int,
+ expected_start: int,
+ read_start: Callable[[int], int],
+ alive: Callable[[int], bool],
+ send: Callable[[int, int], bool],
+ signal_number: int,
+) -> tuple[bool, str]:
+ if pid <= 1 or not alive(pid):
+ return False, "target is stale or not alive"
+ if read_start(pid) != expected_start:
+ return False, "target PID was reused"
+ if not send(pid, signal_number):
+ return False, "signal delivery failed"
+ return True, "signal delivered"
diff --git a/tests/v1/ple_offload/test_correspondence.py b/tests/v1/ple_offload/test_correspondence.py
new file mode 100644
index 0000000..e8d8b9e
--- /dev/null
+++ b/tests/v1/ple_offload/test_correspondence.py
@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.model_executor.layers.ple_offload_layer import (
+ _record_ple_copy_geometry,
+)
+from vllm.v1.ple_offload.connector import PleOffloadConnector
+from vllm.v1.ple_offload.protocol import PleOffloadRequest
+from vllm.v1.ple_offload.worker import (
+ PleOffloadInputBuffers,
+ PleOffloadRunner,
+)
+from vllm.v1.worker.gpu.cudagraph_utils import (
+ BatchExecutionDescriptor,
+ CudaGraphManager,
+)
+from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
+ SpeculatorCudaGraphManager,
+)
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+
+class _ConnectorForValidation:
+ _input_ids_buf = torch.empty(4, dtype=torch.int32)
+ _query_start_loc_buf = torch.empty(3, dtype=torch.int32)
+ _rowdump_n = None
+
+ def _signal_dummy_outputs_impl(self, num_tokens: int) -> None:
+ raise AssertionError("must fail before dummy publication")
+
+ def _launch(self, num_reqs: int, num_tokens: int) -> None:
+ raise AssertionError("must fail before request publication")
+
+
+def test_request_consumer_count_mismatch_precedes_publication() -> None:
+ connector = _ConnectorForValidation()
+ with pytest.raises(RuntimeError, match="request/consumer"):
+ PleOffloadConnector._prepare_forward_impl(
+ connector,
+ 1,
+ 3,
+ False,
+ consumer_num_tokens=4,
+ )
+
+
+def _runner(query_start_loc: list[int]) -> PleOffloadRunner:
+ runner = PleOffloadRunner.__new__(PleOffloadRunner)
+ flag = torch.zeros(1, dtype=torch.int32)
+ target = SimpleNamespace(
+ host_output_buffer=torch.zeros(4, 2),
+ host_flag=flag,
+ )
+ runner._worker_targets = {0: {"ple": [target]}}
+ runner._scratch_bufs = {0: {"ple": torch.empty(4, 2)}}
+ query = torch.zeros(3, dtype=torch.int32)
+ query[: len(query_start_loc)] = torch.tensor(
+ query_start_loc, dtype=torch.int32
+ )
+ runner._input_bufs = {
+ 0: PleOffloadInputBuffers(
+ input_ids_buf=torch.zeros(4, dtype=torch.int32),
+ query_start_loc_buf=query,
+ ngram_context_buf=torch.zeros(2, 1, dtype=torch.int32),
+ )
+ }
+ return runner
+
+
+@pytest.mark.parametrize(
+ ("ple_request", "query_start_loc", "match"),
+ [
+ (PleOffloadRequest(dp_rank=0, num_tokens=3, num_reqs=2), [0, 2, 1],
+ "query_start_loc"),
+ (PleOffloadRequest(dp_rank=0, num_tokens=5, num_reqs=1), [0, 1],
+ "num_tokens"),
+ (PleOffloadRequest(dp_rank=0, num_tokens=1, num_reqs=3), [0, 1, 1],
+ "num_reqs"),
+ (PleOffloadRequest(dp_rank=9, num_tokens=1, num_reqs=1), [0, 1],
+ "dp_rank=9"),
+ ],
+)
+def test_invalid_request_fails_before_publication(
+ ple_request: PleOffloadRequest,
+ query_start_loc: list[int],
+ match: str,
+) -> None:
+ runner = _runner(query_start_loc)
+ flags = [
+ target.host_flag
+ for layers in runner._worker_targets.values()
+ for targets in layers.values()
+ for target in targets
+ ]
+ with pytest.raises(RuntimeError, match=match):
+ runner._validate_requests_before_publication([ple_request])
+ assert all(int(flag[0]) == 0 for flag in flags)
+
+
+class _SelectedGraphManager:
+ def ple_copy_rows(self, desc: BatchExecutionDescriptor) -> int:
+ return 12
+
+
+class _PublicationGuard:
+ called = False
+
+ def prepare_forward(self, *args: Any, **kwargs: Any) -> None:
+ self.called = True
+ raise AssertionError("must fail before connector publication")
+
+
+def test_selected_full_graph_copy_count_precedes_publication() -> None:
+ runner = GPUModelRunner.__new__(GPUModelRunner)
+ connector = _PublicationGuard()
+ runner._ple_offload_connector = connector
+ runner.cudagraph_manager = _SelectedGraphManager()
+ input_batch = SimpleNamespace(
+ num_reqs=1,
+ num_tokens_after_padding=8,
+ input_ids=torch.empty(8, dtype=torch.int32),
+ )
+ desc = BatchExecutionDescriptor(
+ cg_mode=CUDAGraphMode.FULL,
+ num_tokens=8,
+ num_reqs=1,
+ )
+
+ with pytest.raises(RuntimeError, match="selected-forward"):
+ GPUModelRunner._prepare_ple_forward(
+ runner, input_batch, desc, dummy_run=False
+ )
+ assert connector.called is False
+
+
+def test_production_capture_callback_rejects_actual_ple_copy_mismatch() -> None:
+ manager = CudaGraphManager.__new__(CudaGraphManager)
+ manager._ple_full_copy_geometry = {}
+ manager._expected_ple_copy_layers = ("ple",)
+ manager._allow_ple_full_graph = True
+ desc = BatchExecutionDescriptor(
+ cg_mode=CUDAGraphMode.FULL,
+ num_tokens=8,
+ num_reqs=1,
+ )
+
+ with pytest.raises(RuntimeError, match="actual PLE copy geometry"):
+ # This is the production capture scope. The injected observation uses
+ # the same hook called by the real wait-copy custom-op implementation.
+ with manager.capture_ple_copy_geometry(desc):
+ _record_ple_copy_geometry("ple", 12, 16)
+ assert desc not in manager._ple_full_copy_geometry
+
+
+def test_speculator_full_capture_refuses_observed_ple_op() -> None:
+ manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager)
+ manager._ple_full_copy_geometry = {}
+ manager._expected_ple_copy_layers = ()
+ manager._allow_ple_full_graph = False
+ desc = BatchExecutionDescriptor(
+ cg_mode=CUDAGraphMode.FULL,
+ num_tokens=8,
+ num_reqs=1,
+ )
+
+ with pytest.raises(RuntimeError, match="speculator/draft"):
+ with manager.capture_ple_copy_geometry(desc):
+ _record_ple_copy_geometry("unexpected-draft-ple", 8, 16)
+ assert desc not in manager._ple_full_copy_geometry
+
+
+def test_eager_correspondence_uses_actual_padded_input_view() -> None:
+ runner = GPUModelRunner.__new__(GPUModelRunner)
+ connector = _PublicationGuard()
+ runner._ple_offload_connector = connector
+ runner.cudagraph_manager = None
+ input_batch = SimpleNamespace(
+ num_reqs=1,
+ num_tokens_after_padding=8,
+ input_ids=torch.empty(12, dtype=torch.int32),
+ )
+ desc = BatchExecutionDescriptor(
+ cg_mode=CUDAGraphMode.NONE,
+ num_tokens=8,
+ num_reqs=1,
+ )
+ with pytest.raises(RuntimeError, match="selected-forward"):
+ GPUModelRunner._prepare_ple_forward(
+ runner, input_batch, desc, dummy_run=False
+ )
+ assert connector.called is False
diff --git a/tests/v1/ple_offload/test_fatal_rpc.py b/tests/v1/ple_offload/test_fatal_rpc.py
new file mode 100644
index 0000000..1c86f9b
--- /dev/null
+++ b/tests/v1/ple_offload/test_fatal_rpc.py
@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import multiprocessing
+from typing import Any
+
+import pytest
+
+from vllm.v1.ple_offload.connector import PleOffloadConnector
+from vllm.v1.ple_offload.fatal import (
+ fail_ple_execute_model,
+ ple_fail_closed,
+)
+
+
+class _Progress:
+ sequence = 7
+
+ def begin(self, **_: Any) -> int:
+ return self.sequence
+
+ def complete(self) -> None:
+ return
+
+
+class _RaisingProgress(_Progress):
+ def begin(self, **_: Any) -> int:
+ raise RuntimeError("begin injection")
+
+
+class _InjectedConnector:
+ rank = 3
+ tp_rank = 1
+ _layers = {"ple": object()}
+ _progress = _Progress()
+ _input_ids_buf = None
+ _query_start_loc_buf = None
+
+ def __init__(self, where: str) -> None:
+ self.where = where
+ if where == "begin":
+ self._progress = _RaisingProgress()
+
+ def _fail_closed(self, where: str, error: BaseException) -> None:
+ ple_fail_closed(
+ where=where,
+ rank=self.rank,
+ layer="ple",
+ sequence=self._progress.sequence,
+ error=error,
+ )
+
+ def _prepare_forward_impl(self, *args: Any, **kwargs: Any) -> None:
+ if self.where == "prepare":
+ raise RuntimeError("prepare injection")
+
+ def _release_outputs_impl(self) -> None:
+ if self.where == "release":
+ raise RuntimeError("release injection")
+
+
+class _Runner:
+ def __init__(self, where: str) -> None:
+ self._ple_offload_connector = _InjectedConnector(where)
+ self.where = where
+
+ @fail_ple_execute_model
+ def execute_model(self) -> None:
+ connector = self._ple_offload_connector
+ PleOffloadConnector.prepare_forward(
+ connector,
+ 1,
+ 1,
+ False,
+ consumer_num_tokens=1,
+ )
+ if self.where == "forward":
+ raise RuntimeError("forward injection")
+ PleOffloadConnector.release_outputs(connector)
+
+
+def _rpc_loop(where: str, successor_count: Any) -> None:
+ runner = _Runner(where)
+ try:
+ runner.execute_model()
+ except Exception:
+ # Mirrors WorkerProc._execute_worker_rpc: this must never regain control.
+ pass
+ with successor_count.get_lock():
+ successor_count.value += 1
+
+
+@pytest.mark.parametrize("where", ["prepare", "forward", "release"])
+def test_ple_exception_kills_workerproc_style_rpc(
+ where: str,
+) -> None:
+ context = multiprocessing.get_context("spawn")
+ successor_count = context.Value("i", 0)
+ process = context.Process(
+ target=_rpc_loop,
+ args=(where, successor_count),
+ )
+ process.start()
+ # Spawned children import vLLM cold (measured 4-8 s on the reference host); the bound is
+ # about start-up, not the fatal path, which exits immediately.
+ process.join(timeout=120)
+ assert not process.is_alive()
+ assert process.exitcode not in (None, 0)
+ assert successor_count.value == 0
+
+
+def _rpc_loop_connector_only(where: str, successor_count: Any) -> None:
+ # No execute_model decorator: only the connector's own fatal wrappers can
+ # stop the RPC loop from regaining control. (The decorated test above
+ # passes on the decorator alone, so it cannot vouch for these wrappers.)
+ connector = _InjectedConnector(where)
+ try:
+ if where == "begin":
+ PleOffloadConnector.begin_forward(
+ connector, dummy_run=False, num_tokens=1
+ )
+ elif where == "prepare":
+ PleOffloadConnector.prepare_forward(
+ connector, 1, 1, False, consumer_num_tokens=1
+ )
+ else:
+ PleOffloadConnector.release_outputs(connector)
+ except Exception:
+ pass
+ with successor_count.get_lock():
+ successor_count.value += 1
+
+
+@pytest.mark.parametrize("where", ["begin", "prepare", "release"])
+def test_connector_wrappers_are_fatal_without_decorator(where: str) -> None:
+ context = multiprocessing.get_context("spawn")
+ successor_count = context.Value("i", 0)
+ process = context.Process(
+ target=_rpc_loop_connector_only,
+ args=(where, successor_count),
+ )
+ process.start()
+ process.join(timeout=120)
+ assert not process.is_alive()
+ assert process.exitcode not in (None, 0)
+ assert successor_count.value == 0
diff --git a/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py b/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py
new file mode 100644
index 0000000..900851a
--- /dev/null
+++ b/tests/v1/ple_offload/test_fault_bootstrap_disarmed.py
@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Disarmed-path check: a disarmed PLE child writes nothing but the real READY."""
+
+import os
+import subprocess
+import sys
+from multiprocessing.connection import Connection
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+
+def test_partial_ready_writers_are_inert_unless_armed_for_that_cut(monkeypatch):
+ import vllm_ple_offload_bootstrap as bootstrap
+
+ sender = Mock(return_value=(8, 16))
+ monkeypatch.setattr(bootstrap, "_FAULT", None)
+ for cut in ("launch_partial_ready", "launch_partial_ready_header"):
+ assert bootstrap._write_and_ack_partial(99, cut, sender) is False
+ sender.assert_not_called()
+
+ other = SimpleNamespace(cut="exc_before_prepare", fires=Mock(return_value=True))
+ monkeypatch.setattr(bootstrap, "_FAULT", other)
+ assert bootstrap._write_and_ack_partial(99, "launch_partial_ready", sender) is False
+ sender.assert_not_called()
+ other.fires.assert_not_called()
+
+
+# The launch payload's config field is the protocol-0 encoding of None (b"N.").
+_CHILD = r"""
+import os, sys, types
+write_fd = int(sys.argv[1])
+os.environ.pop("VLLM_PLE_TEST_FAULT", None)
+import vllm_ple_offload_bootstrap as bootstrap
+assert bootstrap._FAULT is None
+bootstrap._recv_launch_payload = lambda fd: {
+ "vllm_config_payload": b"N.",
+ "num_workers": 1,
+ "ipc_addr": "ipc:///unused",
+}
+
+def proc_main(*, vllm_config, num_workers, ipc_addr, ready_pipe):
+ assert vllm_config is None
+ ready_pipe.send({"status": "READY", "layer_names": ["ple"]})
+ os._exit(0)
+
+fake = types.ModuleType("vllm.v1.ple_offload.worker")
+fake.PleOffloadWorker = types.SimpleNamespace(proc_main=proc_main)
+sys.modules["vllm.v1.ple_offload.worker"] = fake
+death_read, death_write = os.pipe()
+sys.argv = ["bootstrap", str(write_fd), str(death_read)]
+bootstrap.main()
+"""
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="inherited fds are POSIX-only")
+def test_disarmed_bootstrap_control_stream_is_exactly_one_ready_frame():
+ read_fd, write_fd = os.pipe()
+ env = dict(os.environ)
+ env.pop("VLLM_PLE_TEST_FAULT", None)
+ proc = subprocess.Popen(
+ [sys.executable, "-c", _CHILD, str(write_fd)],
+ pass_fds=(write_fd,),
+ env=env,
+ )
+ os.close(write_fd)
+ reader = Connection(read_fd)
+ try:
+ assert proc.wait(timeout=120) == 0
+ assert reader.poll(5), "no READY frame on the control stream"
+ assert reader.recv() == {"status": "READY", "layer_names": ["ple"]}
+ # Nothing may follow (or precede, which would have broken the frame).
+ with pytest.raises(EOFError):
+ reader.recv_bytes()
+ finally:
+ reader.close()
+ if proc.poll() is None:
+ proc.kill()
+ proc.wait(timeout=120)
diff --git a/tests/v1/ple_offload/test_fault_injector_basics.py b/tests/v1/ple_offload/test_fault_injector_basics.py
new file mode 100644
index 0000000..0a9fd23
--- /dev/null
+++ b/tests/v1/ple_offload/test_fault_injector_basics.py
@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import queue
+import subprocess
+import sys
+import threading
+from collections import deque
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+import vllm_ple_test_fault as test_fault
+
+
+def _spec(monkeypatch, tmp_path, value):
+ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log"))
+ spec = test_fault._parse_fault_spec(value)
+ assert spec is not None
+ return spec
+
+
+@pytest.mark.parametrize(
+ "value",
+ [
+ "cut=park_before_publish,ordinal=1,tp=0",
+ "cut=launch_silent_ready,ordinal=1,pp=0",
+ "cut=drop_request,ordinal=1,tp=1,pp=0",
+ "cut=exc_after_prepare,ordinal=1,tp=0,pp=1",
+ ],
+)
+def test_impossible_role_selectors_are_rejected(monkeypatch, tmp_path, value):
+ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log"))
+ with pytest.raises(ValueError):
+ test_fault._parse_fault_spec(value)
+
+
+def test_armed_import_requires_dedicated_log(tmp_path):
+ env = dict(os.environ)
+ env["VLLM_PLE_TEST_FAULT"] = (
+ "cut=exc_before_prepare,ordinal=1,rank=any"
+ )
+ env.pop("VLLM_PLE_TEST_FAULT_LOG", None)
+ proc = subprocess.run(
+ [sys.executable, "-c", "import vllm_ple_test_fault"],
+ env=env,
+ text=True,
+ capture_output=True,
+ timeout=120,
+ )
+ assert proc.returncode != 0
+ assert "VLLM_PLE_TEST_FAULT_LOG is required" in proc.stderr
+
+
+def test_failed_marker_write_suppresses_action(monkeypatch, tmp_path):
+ spec = _spec(
+ monkeypatch,
+ tmp_path,
+ "cut=exc_after_prepare,ordinal=1,tp=0,pp=0",
+ )
+ monkeypatch.setattr(test_fault.os, "write", Mock(side_effect=BlockingIOError))
+ assert not spec.fires(
+ "exc_after_prepare",
+ role="gpu",
+ unit="runner-forward",
+ eligible=True,
+ tp_rank=0,
+ pp_rank=0,
+ dp_rank=0,
+ transaction_id="sequence:1",
+ )
+ assert spec._state.diagnostic_failed
+ assert not spec._state.fired
+
+
+@pytest.mark.parametrize("order", ["bootstrap-first", "package-first"])
+@pytest.mark.parametrize("armed", [False, True])
+def test_fresh_interpreter_import_identity(tmp_path, order, armed):
+ env = dict(os.environ)
+ env.pop("VLLM_PLE_TEST_FAULT", None)
+ env.pop("VLLM_PLE_TEST_FAULT_LOG", None)
+ if armed:
+ env["VLLM_PLE_TEST_FAULT"] = (
+ "cut=launch_silent_ready,ordinal=1,rank=any"
+ )
+ env["VLLM_PLE_TEST_FAULT_LOG"] = str(tmp_path / f"{order}.log")
+ script = f"""
+import sys
+if {order!r} == "bootstrap-first":
+ import vllm_ple_offload_bootstrap as bootstrap
+ absent_before_package = "vllm_ple_test_fault" not in sys.modules
+ import vllm_ple_test_fault as package
+else:
+ import vllm_ple_test_fault as package
+ import vllm_ple_offload_bootstrap as bootstrap
+ absent_before_package = False
+assert bootstrap._FAULT is package._FAULT
+assert getattr(sys.modules["vllm_ple_test_fault"], "_FAULT") is package._FAULT
+if {armed!r}:
+ assert package._FAULT is not None
+else:
+ assert package._FAULT is None
+ if {order!r} == "bootstrap-first":
+ assert absent_before_package
+"""
+ subprocess.run(
+ [sys.executable, "-c", script],
+ env=env,
+ check=True,
+ timeout=120,
+ )
+
+
+class _IncompleteEvent:
+ def query(self):
+ return False
+
+
+class _FatalObserved(BaseException):
+ pass
+
+
+@pytest.mark.parametrize(
+ ("record_completion", "complete_submission", "expected"),
+ [
+ (False, False, "host-submission"),
+ (True, False, "gpu-completion"),
+ (True, True, "gpu-completion"),
+ ],
+)
+def test_real_progress_monitor_observed_state(
+ monkeypatch,
+ record_completion,
+ complete_submission,
+ expected,
+):
+ import vllm.v1.ple_offload.progress as progress_module
+
+ observed = []
+ reached = threading.Event()
+
+ def fatal(**kwargs):
+ observed.append(kwargs)
+ reached.set()
+ raise _FatalObserved
+
+ monkeypatch.setattr(progress_module, "ple_fail_closed", fatal)
+ monitor = progress_module.PleProgressMonitor(
+ rank=0,
+ layer_names=("ple",),
+ serving_deadline_s=0.02,
+ startup_deadline_s=0.02,
+ )
+ monitor.begin(startup=False, num_tokens=1)
+ if record_completion:
+ monitor.track_completion(_IncompleteEvent())
+ if complete_submission:
+ monitor.complete()
+ assert reached.wait(2)
+ assert expected in observed[0]["where"]
+ monitor.close()
+
+
+def test_drop_retains_event_and_next_pool_one_launch_fails_closed(
+ monkeypatch,
+ tmp_path,
+):
+ import vllm.v1.ple_offload.connector as connector_module
+
+ spec = _spec(
+ monkeypatch,
+ tmp_path,
+ "cut=drop_request,ordinal=1,tp=0,pp=0",
+ )
+ monkeypatch.setattr(connector_module, "_FAULT", spec)
+
+ connector = object.__new__(connector_module.PleOffloadConnector)
+ connector.tp_rank = 0
+ connector.pp_rank = 0
+ connector.dp_rank = 0
+ connector._progress = SimpleNamespace(sequence=7)
+ connector._dropped_events = []
+ connector._d2h_event_pool = queue.Queue(maxsize=1)
+ event = Mock()
+ connector._d2h_event_pool.put_nowait(event)
+ connector._request_queue = Mock()
+ connector._enqueue_cuda_inputs = Mock()
+
+ connector._launch(
+ num_reqs=1,
+ num_tokens=2,
+ serving=True,
+ transaction_id="sequence:7@begin:1.000000000",
+ correlation_id="sequence:7",
+ )
+ assert connector._dropped_events == [event]
+ connector._request_queue.put_nowait.assert_not_called()
+
+ with pytest.raises(
+ RuntimeError,
+ match="more requests than configured concurrent batches",
+ ):
+ connector._launch(
+ num_reqs=1,
+ num_tokens=2,
+ serving=True,
+ transaction_id="sequence:8@begin:2.000000000",
+ correlation_id="sequence:8",
+ )
+
+
+def test_async_output_cut_uses_real_sample_boundary(monkeypatch, tmp_path):
+ import vllm.v1.worker.gpu.model_runner as model_runner_module
+
+ spec = _spec(
+ monkeypatch,
+ tmp_path,
+ "cut=exc_async_output,ordinal=1,tp=0,pp=0",
+ )
+ monkeypatch.setattr(model_runner_module, "_FAULT", spec)
+
+ class FatalReached(Exception):
+ pass
+
+ connector = SimpleNamespace(
+ tp_rank=0,
+ pp_rank=0,
+ dp_rank=0,
+ completion_is_pending=Mock(return_value=True),
+ _fail_closed=Mock(side_effect=FatalReached),
+ )
+ runner = object.__new__(model_runner_module.GPUModelRunner)
+ runner._ple_offload_connector = connector
+ runner.is_last_pp_rank = False
+ runner.pp_handler = Mock()
+ runner.execute_model_state = SimpleNamespace(
+ input_batch=object(),
+ attn_metadata=None,
+ slot_mappings_by_layer=None,
+ hidden_states=None,
+ aux_hidden_states=None,
+ dp_sync=None,
+ finished_req_ids=[],
+ ec_connector_output=None,
+ routed_experts=None,
+ cudagraph_stats=None,
+ ple_serving=True,
+ ple_sequence=11,
+ ple_transaction="sequence:11@begin:1.000000000",
+ )
+
+ with pytest.raises(FatalReached):
+ runner.sample_tokens(None)
+ assert connector._fail_closed.call_args.args[0] == "sample_tokens"
+ error = connector._fail_closed.call_args.args[1]
+ assert isinstance(error, RuntimeError)
+ assert "exc_async_output" in str(error)
diff --git a/tests/v1/ple_offload/test_fault_injector_gates.py b/tests/v1/ple_offload/test_fault_injector_gates.py
new file mode 100644
index 0000000..70e0d3e
--- /dev/null
+++ b/tests/v1/ple_offload/test_fault_injector_gates.py
@@ -0,0 +1,468 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import ast
+import contextlib
+import importlib.util
+import os
+import pickle
+import socket
+import struct
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+import vllm_ple_test_fault as test_fault
+
+
+class _SequenceMustNotBeRead:
+ @property
+ def sequence(self):
+ raise AssertionError("disarmed path read progress.sequence")
+
+ def track_completion(self, _event):
+ pass
+
+ def complete(self):
+ pass
+
+
+def _spec(monkeypatch, tmp_path, text):
+ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log"))
+ value = test_fault._parse_fault_spec(text)
+ assert value is not None
+ return value
+
+
+def test_disarmed_sender_deadline_does_not_enter_logger(monkeypatch):
+ import vllm.v1.ple_offload.worker as worker
+
+ monkeypatch.setattr(worker, "_FAULT", None)
+ monkeypatch.setattr(
+ worker.logger,
+ "info",
+ Mock(side_effect=AssertionError("logger entered")),
+ )
+ left, right = socket.socketpair()
+ try:
+ with pytest.raises(TimeoutError):
+ worker._send_framed_payload(
+ left.fileno(),
+ b"payload",
+ time.monotonic() - 1,
+ )
+ worker.logger.info.assert_not_called()
+ finally:
+ left.close()
+ right.close()
+
+
+def test_disarmed_launch_does_not_read_progress_sequence(monkeypatch):
+ import queue
+ import vllm.v1.ple_offload.connector as connector_module
+
+ monkeypatch.setattr(connector_module, "_FAULT", None)
+ connector = object.__new__(connector_module.PleOffloadConnector)
+ connector.tp_rank = 0
+ connector.pp_rank = 0
+ connector.dp_rank = 0
+ connector._progress = _SequenceMustNotBeRead()
+ connector._d2h_event_pool = queue.Queue()
+ connector._d2h_event_pool.put_nowait(Mock())
+ connector._enqueue_cuda_inputs = Mock()
+ connector._request_queue = Mock()
+
+ connector._launch(
+ 1,
+ 2,
+ serving=True,
+ transaction_id=None,
+ correlation_id=None,
+ )
+ pending = connector._request_queue.put_nowait.call_args.args[0]
+ assert pending.transaction_id is None
+ assert pending.correlation_id is None
+
+
+def test_disarmed_release_does_not_read_progress_sequence(monkeypatch):
+ import vllm.v1.ple_offload.connector as connector_module
+
+ monkeypatch.setattr(connector_module, "_FAULT", None)
+ monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False)
+ event = Mock()
+ stream = Mock()
+ monkeypatch.setattr(connector_module.torch.cuda, "Event", lambda: event)
+ monkeypatch.setattr(
+ connector_module.torch.cuda,
+ "current_stream",
+ lambda _device: stream,
+ )
+
+ connector = object.__new__(connector_module.PleOffloadConnector)
+ connector.tp_rank = 0
+ connector.pp_rank = 0
+ connector.dp_rank = 0
+ connector.device = SimpleNamespace(index=0)
+ connector._layers = {}
+ connector._after_dummy = False
+ connector._progress = _SequenceMustNotBeRead()
+
+ connector._release_outputs_impl(serving=False)
+ event.record.assert_called_once_with(stream)
+
+
+def test_sample_tokens_connector_rank_fatalizes():
+ from vllm.v1.ple_offload.fatal import fail_ple_execute_model
+
+ class Reached(Exception):
+ pass
+
+ class Runner:
+ _ple_offload_connector = SimpleNamespace(
+ _fail_closed=Mock(side_effect=Reached)
+ )
+
+ @fail_ple_execute_model
+ def sample_tokens(self):
+ raise ValueError("sample failure")
+
+ with pytest.raises(Reached):
+ Runner().sample_tokens()
+
+
+def test_sample_tokens_non_connector_rank_rethrows():
+ from vllm.v1.ple_offload.fatal import fail_ple_execute_model
+
+ class Runner:
+ _ple_offload_connector = None
+
+ @fail_ple_execute_model
+ def sample_tokens(self):
+ raise ValueError("sample failure")
+
+ with pytest.raises(ValueError, match="sample failure"):
+ Runner().sample_tokens()
+
+
+def test_profile_is_never_eligible_even_with_phase_any(
+ monkeypatch,
+ tmp_path,
+):
+ spec = _spec(
+ monkeypatch,
+ tmp_path,
+ "cut=exc_after_prepare,ordinal=1,tp=0,pp=0,phase=any",
+ )
+ assert not spec.fires(
+ "exc_after_prepare",
+ role="gpu",
+ unit="runner-forward",
+ eligible=False,
+ tp_rank=0,
+ pp_rank=0,
+ dp_rank=0,
+ transaction_id="sequence:1@begin:1.000000000",
+ correlation_id="sequence:1",
+ )
+ assert spec._state.count == 0
+
+
+def test_offload_counts_only_serving_and_preserves_correlation(
+ monkeypatch,
+ tmp_path,
+):
+ import vllm.v1.ple_offload.worker as worker
+
+ spec = _spec(
+ monkeypatch,
+ tmp_path,
+ "cut=park_before_publish,ordinal=1,rank=any",
+ )
+ monkeypatch.setattr(worker, "_FAULT", spec)
+ runner = object.__new__(worker.PleOffloadRunner)
+ runner._validate_requests_before_publication = Mock()
+ runner._worker_targets = {0: {}}
+ runner._input_bufs = {}
+ runner._scratch_bufs = {}
+ runner._layers = {}
+ runner._clamp_input_ids = False
+ runner._step = 0
+ runner._request_index = 0
+
+ request = worker.PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1)
+ for _ in range(3):
+ runner._handle_requests_impl(
+ [worker._ReceivedPleOffloadRequest(request, False, None, None)]
+ )
+ assert runner._request_index == 0
+
+ reached = threading.Event()
+ monkeypatch.setattr(
+ worker.threading,
+ "Event",
+ lambda: SimpleNamespace(wait=lambda: reached.set()),
+ )
+ runner._handle_requests_impl(
+ [
+ worker._ReceivedPleOffloadRequest(
+ request,
+ True,
+ "sequence:17@begin:4.000000000",
+ "sequence:17",
+ )
+ ]
+ )
+ assert reached.is_set()
+ line = (tmp_path / "fault.log").read_text()
+ assert "transaction=request:1@sequence:17@begin:4.000000000" in line
+ assert "correlation=sequence:17" in line
+
+
+def test_cut6_requires_same_incomplete_completion(monkeypatch):
+ import vllm.v1.ple_offload.progress as progress
+
+ class Event:
+ complete = False
+
+ def query(self):
+ return self.complete
+
+ monitor = progress.PleProgressMonitor(
+ rank=0,
+ layer_names=("ple",),
+ serving_deadline_s=60,
+ startup_deadline_s=60,
+ )
+ event = Event()
+ try:
+ sequence = monitor.begin(startup=False, num_tokens=1)
+ monitor.track_completion(event)
+ monitor.complete()
+ assert monitor.completion_is_pending(sequence)
+ event.complete = True
+ assert not monitor.completion_is_pending(sequence)
+ assert not monitor.completion_is_pending(sequence + 1)
+ finally:
+ monitor.close()
+
+
+def test_partial_ack_requires_complete_required_prefix(monkeypatch):
+ import vllm_ple_offload_bootstrap as bootstrap
+
+ fault = SimpleNamespace(
+ cut="launch_partial_ready_header",
+ fires=Mock(return_value=True),
+ )
+ monkeypatch.setattr(bootstrap, "_FAULT", fault)
+ assert not bootstrap._write_and_ack_partial(
+ 1,
+ "launch_partial_ready_header",
+ lambda _fd: (1, 4),
+ )
+ fault.fires.assert_not_called()
+
+ assert bootstrap._write_and_ack_partial(
+ 1,
+ "launch_partial_ready_header",
+ lambda _fd: (2, 4),
+ )
+ assert fault.fires.call_args.kwargs["bytes_written"] == 2
+ assert fault.fires.call_args.kwargs["frame_size"] == 4
+
+
+@pytest.mark.timeout(180)
+@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX pass_fds")
+@pytest.mark.parametrize(
+ ("cut", "expected_transaction"),
+ [
+ ("launch_block_reconstruction", "reducer:configuration"),
+ (
+ "launch_block_worker_import",
+ "module:vllm.v1.ple_offload.worker",
+ ),
+ ],
+)
+def test_gate_is_inside_reducer_or_loader_and_parent_death_wins(
+ tmp_path,
+ cut,
+ expected_transaction,
+):
+ gate = tmp_path / "gate"
+ os.mkfifo(gate)
+ gate_keeper = os.open(gate, os.O_RDWR | os.O_NONBLOCK)
+ control_parent, control_child = socket.socketpair()
+ death_parent, death_child = socket.socketpair()
+ fault_log = tmp_path / "fault.log"
+ env = dict(os.environ)
+ env["VLLM_PLE_TEST_FAULT"] = f"cut={cut},ordinal=1,rank=any"
+ env["VLLM_PLE_TEST_FAULT_LOG"] = str(fault_log)
+ env["VLLM_PLE_TEST_FAULT_GATE"] = str(gate)
+
+ config_payload = pickle.dumps(None, protocol=pickle.HIGHEST_PROTOCOL)
+ if cut == "launch_block_reconstruction":
+ config_payload = test_fault._FaultReconstructionPayload(config_payload)
+ launch = pickle.dumps(
+ {
+ "vllm_config_payload": config_payload,
+ "num_workers": 1,
+ "ipc_addr": "ipc:///unused-fault",
+ },
+ protocol=pickle.HIGHEST_PROTOCOL,
+ )
+
+ process = subprocess.Popen(
+ [
+ sys.executable,
+ "-m",
+ "vllm_ple_offload_bootstrap",
+ str(control_child.fileno()),
+ str(death_child.fileno()),
+ ],
+ env=env,
+ pass_fds=(control_child.fileno(), death_child.fileno()),
+ )
+ control_child.close()
+ death_child.close()
+ try:
+ control_parent.sendall(struct.pack("!Q", len(launch)) + launch)
+ deadline = time.monotonic() + 30
+ while time.monotonic() < deadline:
+ if fault_log.exists() and expected_transaction in fault_log.read_text():
+ break
+ assert process.poll() is None
+ time.sleep(0.05)
+ else:
+ pytest.fail("operation-entry acknowledgement was not emitted")
+
+ assert process.poll() is None
+ death_parent.close()
+ assert process.wait(timeout=120) != 0
+ finally:
+ with contextlib.suppress(ProcessLookupError):
+ process.kill()
+ process.wait(timeout=120)
+ control_parent.close()
+ with contextlib.suppress(OSError):
+ death_parent.close()
+ os.close(gate_keeper)
+
+
+def _load_gate_helper():
+ root = Path(__file__).resolve().parents[3]
+ path = root / "probe" / "faultcut_gate.py"
+ spec = importlib.util.spec_from_file_location("faultcut_gate", path)
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+@pytest.mark.parametrize(
+ "change",
+ [
+ {"alive": False},
+ {"observed_start": 88},
+ {"bytes_sent": 0},
+ {"bytes_sent": 100},
+ {"completed": True},
+ {"child_not_reading": False},
+ ],
+)
+def test_external_partial_send_gate_rejections(change):
+ gate = _load_gate_helper()
+ values = {
+ "pid": 12,
+ "expected_start": 77,
+ "observed_start": 77,
+ "alive": True,
+ "bytes_sent": 4,
+ "frame_size": 100,
+ "completed": False,
+ "child_not_reading": True,
+ }
+ values.update(change)
+ ok, _reason = gate.validate_partial_send_gate(**values)
+ assert not ok
+
+
+def test_external_gate_rejects_failed_signal():
+ gate = _load_gate_helper()
+ ok, reason = gate.deliver_checked_signal(
+ pid=12,
+ expected_start=77,
+ read_start=lambda _pid: 77,
+ alive=lambda _pid: True,
+ send=lambda _pid, _sig: False,
+ signal_number=9,
+ )
+ assert not ok
+ assert "delivery failed" in reason
+
+
+def test_every_fault_call_has_one_boolop_guard_and_no_prior_state_read():
+ root = Path(__file__).resolve().parents[3]
+ expected = {
+ "vllm/v1/ple_offload/connector.py": 3,
+ "vllm/v1/ple_offload/worker.py": 1,
+ "vllm/v1/worker/gpu/model_runner.py": 4,
+ "vllm_ple_offload_bootstrap/__init__.py": 1,
+ }
+
+ for relative, wanted in expected.items():
+ tree = ast.parse((root / relative).read_text(encoding="utf-8"))
+ parents = {}
+ for parent in ast.walk(tree):
+ for child in ast.iter_child_nodes(parent):
+ parents[child] = parent
+
+ calls = [
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Attribute)
+ and node.func.attr == "fires"
+ and isinstance(node.func.value, ast.Name)
+ and node.func.value.id == "_FAULT"
+ ]
+ assert len(calls) == wanted, relative
+ for call in calls:
+ owner = call
+ while owner in parents and not isinstance(owner, ast.BoolOp):
+ owner = parents[owner]
+ assert isinstance(owner, ast.BoolOp), (relative, call.lineno)
+ assert isinstance(owner.op, ast.And), (relative, call.lineno)
+
+ guard_index = None
+ call_index = None
+ for index, value in enumerate(owner.values):
+ if (
+ isinstance(value, ast.Compare)
+ and isinstance(value.left, ast.Name)
+ and value.left.id == "_FAULT"
+ and len(value.ops) == 1
+ and isinstance(value.ops[0], ast.IsNot)
+ and len(value.comparators) == 1
+ and isinstance(value.comparators[0], ast.Constant)
+ and value.comparators[0].value is None
+ ):
+ guard_index = index
+ if call in ast.walk(value):
+ call_index = index
+ assert guard_index is not None
+ assert call_index is not None
+ assert guard_index < call_index
+
+ forbidden = {"sequence", "has_pending_completions",
+ "completion_is_pending", "fault_transaction"}
+ before_guard = owner.values[:guard_index]
+ assert not any(
+ isinstance(node, ast.Attribute) and node.attr in forbidden
+ for value in before_guard
+ for node in ast.walk(value)
+ )
diff --git a/tests/v1/ple_offload/test_fault_wire_frames.py b/tests/v1/ple_offload/test_fault_wire_frames.py
new file mode 100644
index 0000000..b864355
--- /dev/null
+++ b/tests/v1/ple_offload/test_fault_wire_frames.py
@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Disarmed-path checks: the disarmed PLE request wire format stays single-frame."""
+
+import contextlib
+import queue
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import msgspec
+import pytest
+
+
+def _request():
+ from vllm.v1.ple_offload.protocol import PleOffloadRequest
+
+ return PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1)
+
+
+def test_disarmed_offload_receive_uses_single_frame(monkeypatch):
+ import vllm.v1.ple_offload.worker as worker_module
+
+ monkeypatch.setattr(worker_module, "_FAULT", None)
+ request = _request()
+ socket = Mock()
+ socket.recv.return_value = msgspec.msgpack.encode(request)
+ socket.recv_multipart.side_effect = AssertionError("multipart read while disarmed")
+
+ assert worker_module._receive_offload_request(socket) == request
+ socket.recv.assert_called_once_with(0)
+
+
+def test_armed_offload_receive_requires_metadata_frame(monkeypatch):
+ import vllm.v1.ple_offload.worker as worker_module
+
+ monkeypatch.setattr(worker_module, "_FAULT", object())
+ request = _request()
+ metadata = {"serving": True, "transaction": "request:1", "correlation": "sequence:3"}
+ socket = Mock()
+ socket.recv.side_effect = AssertionError("single-frame read while armed")
+ socket.recv_multipart.return_value = [
+ msgspec.msgpack.encode(request),
+ msgspec.msgpack.encode(metadata),
+ ]
+
+ received = worker_module._receive_offload_request(socket, 1)
+ assert received.request == request
+ assert received.serving is True
+ assert received.transaction_id == "request:1"
+ assert received.correlation_id == "sequence:3"
+ socket.recv_multipart.assert_called_once_with(1)
+
+ socket.recv_multipart.return_value = [msgspec.msgpack.encode(request)]
+ with pytest.raises(RuntimeError):
+ worker_module._receive_offload_request(socket)
+
+
+@pytest.mark.parametrize("armed", [False, True])
+def test_connector_publication_frames_follow_arming(monkeypatch, armed):
+ import vllm.v1.ple_offload.connector as connector_module
+
+ fault = SimpleNamespace(fires=lambda *args, **kwargs: False) if armed else None
+ monkeypatch.setattr(connector_module, "_FAULT", fault)
+ monkeypatch.setattr(
+ connector_module.torch.accelerator,
+ "device_index",
+ lambda _index: contextlib.nullcontext(),
+ )
+ connector = object.__new__(connector_module.PleOffloadConnector)
+ connector.tp_rank = 0
+ connector.pp_rank = 0
+ connector.dp_rank = 0
+ connector.device = SimpleNamespace(index=0)
+ connector._d2h_event_pool = queue.Queue()
+ event = Mock()
+ request = _request()
+ pending = connector_module._PendingPleOffloadRequest(
+ request, event, True, "sequence:4@begin:1.000000000", "sequence:4"
+ )
+ socket = Mock()
+
+ connector._process_request(pending, socket)
+
+ event.synchronize.assert_called_once()
+ if armed:
+ socket.send.assert_not_called()
+ payload, metadata = socket.send_multipart.call_args.args[0]
+ assert msgspec.msgpack.decode(metadata) == {
+ "serving": True,
+ "transaction": "sequence:4@begin:1.000000000",
+ "correlation": "sequence:4",
+ }
+ else:
+ socket.send_multipart.assert_not_called()
+ (payload,) = socket.send.call_args.args
+ assert payload == msgspec.msgpack.encode(request)
+ assert connector._d2h_event_pool.get_nowait() is event
diff --git a/tests/v1/ple_offload/test_ipc_socket_unlink.py b/tests/v1/ple_offload/test_ipc_socket_unlink.py
new file mode 100644
index 0000000..92fbc42
--- /dev/null
+++ b/tests/v1/ple_offload/test_ipc_socket_unlink.py
@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+import os
+from pathlib import Path
+import signal
+import subprocess
+import sys
+import tempfile
+import time
+
+import pytest
+
+from vllm.v1.ple_offload.worker import PleOffloadWorkerHandle
+
+
+pytestmark = pytest.mark.skipif(
+ sys.platform != "linux", reason="requires Linux filesystem ipc sockets"
+)
+
+
+_STUB_CHILD = r"""
+import pathlib
+import signal
+import sys
+import threading
+
+import zmq
+
+from vllm.v1.ple_offload.worker import _unlink_ipc_socket
+
+ipc_addr, mode, marker = sys.argv[1:]
+stopping = threading.Event()
+if mode == "graceful":
+ signal.signal(signal.SIGTERM, lambda _signum, _frame: stopping.set())
+else:
+ signal.signal(signal.SIGTERM, signal.SIG_IGN)
+
+context = zmq.Context()
+socket = context.socket(zmq.PULL)
+socket.bind(ipc_addr)
+try:
+ stopping.wait()
+finally:
+ socket.close(linger=0)
+ context.term()
+ if mode == "graceful":
+ _unlink_ipc_socket(ipc_addr)
+ pathlib.Path(marker).touch()
+"""
+
+
+def _start_stub(ipc_addr: str, mode: str, marker: Path) -> subprocess.Popen:
+ return subprocess.Popen(
+ [sys.executable, "-c", _STUB_CHILD, ipc_addr, mode, str(marker)],
+ stdout=subprocess.DEVNULL,
+ stderr=subprocess.PIPE,
+ text=True,
+ )
+
+
+def _wait_for_socket(proc: subprocess.Popen, path: Path) -> None:
+ deadline = time.monotonic() + 30
+ while time.monotonic() < deadline:
+ if path.exists():
+ return
+ if proc.poll() is not None:
+ stderr = proc.stderr.read() if proc.stderr is not None else ""
+ pytest.fail(f"stub child exited before bind: {stderr}")
+ time.sleep(0.01)
+ pytest.fail("stub child did not bind its ipc socket")
+
+
+def _cleanup_process(proc: subprocess.Popen, socket_path: Path) -> None:
+ if proc.poll() is None:
+ proc.kill()
+ proc.wait(timeout=5)
+ with contextlib.suppress(FileNotFoundError):
+ os.unlink(socket_path)
+ if proc.stderr is not None:
+ proc.stderr.close()
+
+
+def test_handle_graceful_close_removes_child_bound_socket() -> None:
+ with tempfile.TemporaryDirectory(prefix="ple-ipc-", dir="/tmp") as temp_dir:
+ socket_path = Path(temp_dir) / "owned.sock"
+ marker = Path(temp_dir) / "child-cleaned"
+ ipc_addr = f"ipc://{socket_path}"
+ proc = _start_stub(ipc_addr, "graceful", marker)
+ try:
+ handle = PleOffloadWorkerHandle(
+ proc=proc,
+ death_writer=None,
+ ready_pipe_reader=None,
+ launch_deadline=time.monotonic() + 30,
+ ipc_addr=ipc_addr,
+ )
+ _wait_for_socket(proc, socket_path)
+
+ handle.close()
+
+ assert proc.returncode == 0
+ assert marker.exists(), "the child ran its post-context unlink"
+ assert not socket_path.exists()
+ finally:
+ _cleanup_process(proc, socket_path)
+
+
+def test_handle_close_removes_socket_after_child_sigkill() -> None:
+ with tempfile.TemporaryDirectory(prefix="ple-ipc-", dir="/tmp") as temp_dir:
+ socket_path = Path(temp_dir) / "owned.sock"
+ marker = Path(temp_dir) / "unused"
+ ipc_addr = f"ipc://{socket_path}"
+ proc = _start_stub(ipc_addr, "park", marker)
+ try:
+ handle = PleOffloadWorkerHandle(
+ proc=proc,
+ death_writer=None,
+ ready_pipe_reader=None,
+ launch_deadline=time.monotonic() + 30,
+ ipc_addr=ipc_addr,
+ )
+ _wait_for_socket(proc, socket_path)
+ proc.kill()
+ proc.wait(timeout=5)
+ assert proc.returncode == -signal.SIGKILL
+ assert socket_path.exists()
+
+ handle.close()
+
+ assert not socket_path.exists()
+ finally:
+ _cleanup_process(proc, socket_path)
diff --git a/tests/v1/ple_offload/test_pending_rpc.py b/tests/v1/ple_offload/test_pending_rpc.py
new file mode 100644
index 0000000..ef3c0fd
--- /dev/null
+++ b/tests/v1/ple_offload/test_pending_rpc.py
@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import multiprocessing
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+import vllm.v1.executor.multiproc_executor as executor_module
+from vllm.v1.executor.multiproc_executor import WorkerProc
+from vllm.v1.ple_offload.fatal import ple_fail_closed
+
+
+class _BrokenAsyncOutput:
+ def get_output(self) -> Any:
+ raise RuntimeError("injected asynchronous output failure")
+
+
+class _Progress:
+ sequence = 9
+
+ def __init__(self, pending: bool) -> None:
+ self.has_pending_completions = pending
+
+
+class _Connector:
+ rank = 2
+ _layers = {"ple": object()}
+
+ def __init__(self, pending: bool) -> None:
+ self._progress = _Progress(pending)
+
+ @property
+ def has_pending_completions(self) -> bool:
+ return self._progress.has_pending_completions
+
+ def _fail_closed(self, where: str, error: BaseException) -> None:
+ ple_fail_closed(
+ where=where,
+ rank=self.rank,
+ layer="ple",
+ sequence=self._progress.sequence,
+ error=error,
+ )
+
+
+def _async_output_failure_case(pending: bool) -> None:
+ connector = _Connector(pending)
+ worker = SimpleNamespace(
+ model_runner=SimpleNamespace(_ple_offload_connector=connector)
+ )
+ proc = WorkerProc.__new__(WorkerProc)
+ proc.worker = SimpleNamespace(worker=worker)
+ proc.worker_response_mq = None
+
+ # Exercise the real async-output exception recovery branch without needing
+ # to construct the production dataclass's CUDA-backed fields.
+ executor_module.AsyncModelRunnerOutput = _BrokenAsyncOutput
+ proc.enqueue_output(_BrokenAsyncOutput())
+
+
+@pytest.mark.parametrize(
+ ("pending", "must_die"),
+ [(True, True), (False, False)],
+)
+def test_async_output_failure_obeys_pending_ple_ownership(
+ pending: bool,
+ must_die: bool,
+) -> None:
+ context = multiprocessing.get_context("spawn")
+ process = context.Process(
+ target=_async_output_failure_case,
+ args=(pending,),
+ )
+ process.start()
+ process.join(timeout=120)
+ assert not process.is_alive()
+ if must_die:
+ assert process.exitcode not in (None, 0)
+ else:
+ assert process.exitcode == 0
diff --git a/tests/v1/ple_offload/test_progress_deadline.py b/tests/v1/ple_offload/test_progress_deadline.py
new file mode 100644
index 0000000..5e7ca19
--- /dev/null
+++ b/tests/v1/ple_offload/test_progress_deadline.py
@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import multiprocessing
+import os
+import threading
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+
+import vllm.v1.ple_offload.connector as connector_module
+from vllm.v1.ple_offload.connector import PleOffloadConnector
+from vllm.v1.ple_offload.progress import PleProgressMonitor
+
+
+class _PendingEvent:
+ def record(self, stream: Any) -> None:
+ stream.recorded = True
+
+ def query(self) -> bool:
+ return False
+
+
+class _BrokenEvent(_PendingEvent):
+ def __init__(self, query_entered: Any) -> None:
+ self.query_entered = query_entered
+
+ def query(self) -> bool:
+ self.query_entered.set()
+ raise RuntimeError("injected event query failure")
+
+
+class _FakeStream:
+ recorded = False
+
+ def synchronize(self) -> None:
+ return
+
+
+class _FakeLayer:
+ def release_offloaded_output(self, stream: Any) -> None:
+ stream.reset_written = True
+
+
+def _silent_producer_case(block_input: bool, started: Any) -> None:
+ monitor = PleProgressMonitor(
+ rank=2,
+ layer_names=("ple",),
+ serving_deadline_s=0.2,
+ startup_deadline_s=1.0,
+ )
+ monitor.begin(startup=False, num_tokens=32)
+ started.set()
+
+ if block_input:
+ # Models the request thread stuck in d2h_done_event.synchronize().
+ threading.Thread(
+ target=threading.Event().wait,
+ daemon=True,
+ ).start()
+
+ # Neither the fake producer nor the test terminates the process. The
+ # independent progress monitor must do so.
+ threading.Event().wait()
+
+
+@pytest.mark.parametrize("block_input", [False, True])
+def test_progress_deadline_terminates_alive_silent_worker(
+ block_input: bool,
+) -> None:
+ context = multiprocessing.get_context("spawn")
+ started = context.Event()
+ process = context.Process(
+ target=_silent_producer_case,
+ args=(block_input, started),
+ )
+ process.start()
+ # Cold vLLM import in the spawned child takes 4-8 s on the reference host; the deadline
+ # itself (0.2 s) is only charged after `started` is set.
+ assert started.wait(timeout=120)
+ process.join(timeout=10)
+ assert not process.is_alive()
+ assert process.exitcode not in (None, 0)
+
+
+def _pending_completion_case(use_release_wrapper: bool, started: Any) -> None:
+ monitor = PleProgressMonitor(
+ rank=2,
+ layer_names=("ple",),
+ serving_deadline_s=0.2,
+ startup_deadline_s=1.0,
+ )
+ if use_release_wrapper:
+ stream = _FakeStream()
+ connector_module.torch.cuda.current_stream = lambda device=None: stream
+ connector_module.torch.cuda.Event = _PendingEvent
+ connector_module.torch.cuda.is_current_stream_capturing = lambda: False
+ connector_module._rowdump.enabled = lambda: False
+
+ connector = PleOffloadConnector.__new__(PleOffloadConnector)
+ connector.device = torch.device("cpu")
+ connector.rank = 2
+ connector.tp_rank = 0
+ connector._layers = {"ple": _FakeLayer()}
+ connector._progress = monitor
+ connector._rowdump_n = None
+ connector._after_dummy = False
+ monitor.begin(startup=False, num_tokens=8)
+ connector.release_outputs()
+ assert stream.recorded
+ assert stream.reset_written
+ else:
+ monitor.begin(startup=False, num_tokens=8)
+ monitor.track_completion(_PendingEvent())
+ monitor.complete()
+ started.set()
+ threading.Event().wait()
+
+
+@pytest.mark.parametrize("use_release_wrapper", [False, True])
+def test_pending_gpu_completion_expires_after_release_returns(
+ use_release_wrapper: bool,
+) -> None:
+ context = multiprocessing.get_context("spawn")
+ started = context.Event()
+ process = context.Process(
+ target=_pending_completion_case,
+ args=(use_release_wrapper, started),
+ )
+ process.start()
+ assert started.wait(timeout=120)
+ process.join(timeout=10)
+ assert not process.is_alive()
+ assert process.exitcode not in (None, 0)
+
+
+@pytest.mark.parametrize("value", [float("inf"), float("nan"), -1.0, 1e12])
+def test_constructor_rejects_invalid_deadlines(value: float) -> None:
+ with pytest.raises(ValueError, match="finite"):
+ PleProgressMonitor(
+ rank=0,
+ layer_names=("ple",),
+ serving_deadline_s=value,
+ startup_deadline_s=1.0,
+ )
+ with pytest.raises(ValueError, match="finite"):
+ PleProgressMonitor(
+ rank=0,
+ layer_names=("ple",),
+ serving_deadline_s=1.0,
+ startup_deadline_s=value,
+ )
+
+
+@pytest.mark.parametrize(
+ "name",
+ ["VLLM_PLE_PROGRESS_DEADLINE_S", "VLLM_PLE_STARTUP_DEADLINE_S"],
+)
+@pytest.mark.parametrize("value", ["inf", "nan", "-1", "1e12"])
+def test_environment_rejects_invalid_deadlines(
+ monkeypatch: pytest.MonkeyPatch,
+ name: str,
+ value: str,
+) -> None:
+ monkeypatch.setenv("VLLM_PLE_PROGRESS_DEADLINE_S", "1")
+ monkeypatch.setenv("VLLM_PLE_STARTUP_DEADLINE_S", "1")
+ monkeypatch.setenv(name, value)
+ with pytest.raises(ValueError, match="finite"):
+ PleProgressMonitor.from_environment(rank=0, layer_names=("ple",))
+
+
+def _monitor_exception_case(
+ started: Any,
+ query_entered: Any,
+ diagnostic_path: str,
+) -> None:
+ diagnostic_fd = os.open(diagnostic_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+ os.dup2(diagnostic_fd, 2)
+ os.close(diagnostic_fd)
+ monitor = PleProgressMonitor(
+ rank=4,
+ layer_names=("ple",),
+ serving_deadline_s=10.0,
+ startup_deadline_s=10.0,
+ )
+ monitor.begin(startup=False, num_tokens=8)
+ # Set before arming the failing event: the monitor thread's os._exit must not
+ # race a multiprocessing Event.set() that holds the shared condition lock.
+ started.set()
+ monitor.track_completion(_BrokenEvent(query_entered))
+ monitor.complete()
+ threading.Event().wait()
+
+
+def test_monitor_thread_exception_is_fatal(tmp_path: Path) -> None:
+ context = multiprocessing.get_context("spawn")
+ started = context.Event()
+ query_entered = context.Event()
+ diagnostic_path = tmp_path / "monitor-fatal.txt"
+ process = context.Process(
+ target=_monitor_exception_case,
+ args=(started, query_entered, str(diagnostic_path)),
+ )
+ process.start()
+ assert started.wait(timeout=120)
+ assert query_entered.wait(timeout=10)
+ process.join(timeout=10)
+ assert not process.is_alive()
+ assert process.exitcode not in (None, 0)
+ diagnostic = diagnostic_path.read_text(encoding="utf-8")
+ assert "where=progress-monitor-thread" in diagnostic
+ assert (
+ "error=RuntimeError:injected event query failure"
+ in diagnostic
+ )
diff --git a/tests/v1/ple_offload/test_shm_unlink.py b/tests/v1/ple_offload/test_shm_unlink.py
new file mode 100644
index 0000000..600a7bb
--- /dev/null
+++ b/tests/v1/ple_offload/test_shm_unlink.py
@@ -0,0 +1,397 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from multiprocessing.reduction import ForkingPickler
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+import torch
+import torch.multiprocessing as torch_mp
+
+from vllm.v1.ple_offload import connector as connector_module
+from vllm.v1.ple_offload.connector import PleOffloadConnector
+from vllm.v1.ple_offload.protocol import PleOffloadRegistration
+
+
+def _require_file_system_shm() -> None:
+ if (
+ os.name != "posix"
+ or not os.path.isdir("/dev/shm")
+ or not os.access("/dev/shm", os.W_OK)
+ ):
+ pytest.skip("POSIX /dev/shm is unavailable")
+ if "file_system" not in torch_mp.get_all_sharing_strategies():
+ pytest.skip("torch file_system sharing is unavailable")
+
+
+def _shm_path(tensor: torch.Tensor) -> Path:
+ _, filename, _ = tensor.untyped_storage()._share_filename_cpu_()
+ return Path("/dev/shm") / os.fsdecode(filename).lstrip("/")
+
+
+def _make_connector_storage() -> tuple[PleOffloadConnector, list[torch.Tensor]]:
+ original_strategy = torch_mp.get_sharing_strategy()
+ torch_mp.set_sharing_strategy("file_system")
+ try:
+ host_outputs = {
+ "layer.0": torch.arange(12, dtype=torch.float32)
+ .reshape(3, 4)
+ .share_memory_(),
+ "layer.1": torch.arange(8, dtype=torch.float32)
+ .reshape(2, 4)
+ .share_memory_(),
+ }
+ host_flags = {
+ name: torch.zeros(16, dtype=torch.int32).share_memory_()
+ for name in host_outputs
+ }
+ input_ids = torch.arange(16, dtype=torch.int32).share_memory_()
+ query_start_loc = torch.arange(5, dtype=torch.int32).share_memory_()
+ ngram_context = torch.arange(12, dtype=torch.int32).reshape(4, 3)
+ ngram_context.share_memory_()
+
+ registration = PleOffloadRegistration(
+ worker_id=0,
+ tp_rank=0,
+ dp_rank=0,
+ host_output_buffers=host_outputs,
+ host_flag_tensors=host_flags,
+ input_ids_buf=input_ids,
+ query_start_loc_buf=query_start_loc,
+ ngram_context_buf=ngram_context,
+ )
+ # Match producer serialization and worker reconstruction. loads()
+ # consumes the transfer reference added by dumps(), leaving storage
+ # destruction able to reach the allocator's zero-counter branch.
+ payload = ForkingPickler.dumps(registration)
+ rebuilt_registration = ForkingPickler.loads(payload)
+ assert isinstance(rebuilt_registration, PleOffloadRegistration)
+ del rebuilt_registration
+ finally:
+ torch_mp.set_sharing_strategy(original_strategy)
+
+ connector = PleOffloadConnector.__new__(PleOffloadConnector)
+ connector._host_output_bufs = host_outputs
+ connector._host_flags = host_flags
+ connector._reserved_host_storage_filenames = set()
+ connector._input_ids_buf = input_ids
+ connector._query_start_loc_buf = query_start_loc
+ connector._ngram_context_buf = ngram_context
+ tensors = [
+ *host_outputs.values(),
+ *host_flags.values(),
+ input_ids,
+ query_start_loc,
+ ngram_context,
+ ]
+ assert len(tensors) == 7
+ return connector, tensors
+
+
+def test_release_host_storage_unlinks_names_but_keeps_mappings() -> None:
+ _require_file_system_shm()
+ connector, tensors = _make_connector_storage()
+ paths = [_shm_path(tensor) for tensor in tensors]
+ before = [tensor.clone() for tensor in tensors]
+ assert len(set(paths)) == 7
+ assert all(path.exists() for path in paths)
+
+ connector.release_host_storage()
+ connector.release_host_storage()
+
+ assert all(not path.exists() for path in paths)
+ for tensor, expected in zip(tensors, before):
+ assert torch.equal(tensor, expected)
+ tensor.add_(1)
+ assert torch.equal(tensor, expected + 1)
+
+
+def test_without_release_helper_names_remain(monkeypatch) -> None:
+ _require_file_system_shm()
+ connector, tensors = _make_connector_storage()
+ paths = [_shm_path(tensor) for tensor in tensors]
+
+ with monkeypatch.context() as patch:
+ patch.setattr(
+ connector_module,
+ "_unlink_filename_shared_storage",
+ lambda _tensor, _reservations: None,
+ )
+ connector.release_host_storage()
+
+ # Do not manually unlink these while their unreserved storages are live.
+ # Normal storage destruction removes them after the test frame is released.
+ assert all(path.exists() for path in paths)
+
+
+def test_release_host_storage_does_not_convert_fd_shared_storage() -> None:
+ if "file_descriptor" not in torch_mp.get_all_sharing_strategies():
+ pytest.skip("torch file_descriptor sharing is unavailable")
+ original_strategy = torch_mp.get_sharing_strategy()
+ torch_mp.set_sharing_strategy("file_descriptor")
+ try:
+ fd_tensor = torch.arange(8, dtype=torch.int32).share_memory_()
+ finally:
+ torch_mp.set_sharing_strategy(original_strategy)
+
+ connector = PleOffloadConnector.__new__(PleOffloadConnector)
+ connector._host_output_bufs = {"layer.0": fd_tensor}
+ connector._host_flags = {}
+ connector._reserved_host_storage_filenames = set()
+ connector._input_ids_buf = torch.arange(2, dtype=torch.int32)
+ connector._query_start_loc_buf = torch.arange(2, dtype=torch.int32)
+ connector._ngram_context_buf = None
+ expected = fd_tensor.clone()
+
+ connector.release_host_storage()
+
+ assert fd_tensor.untyped_storage()._get_shared_fd() >= 0
+ assert torch.equal(fd_tensor, expected)
+
+
+class _FakeFilenameStorage:
+ def __init__(self, filename: bytes) -> None:
+ self.device = torch.device("cpu")
+ self.filename = filename
+ self.increfs = 0
+ self.decrefs = 0
+
+ def is_shared(self) -> bool:
+ return True
+
+ def _get_shared_fd(self) -> int:
+ raise RuntimeError("filename-backed storage has no shared fd")
+
+ def _share_filename_cpu_(self) -> tuple[bytes, bytes, int]:
+ return b"manager", self.filename, 4
+
+ def _shared_incref(self) -> None:
+ self.increfs += 1
+
+ def _shared_decref(self) -> "_FakeFilenameStorage":
+ self.decrefs += 1
+ return self
+
+
+class _FakeTensor:
+ def __init__(self, storage: _FakeFilenameStorage) -> None:
+ self.storage = storage
+
+ def untyped_storage(self) -> _FakeFilenameStorage:
+ return self.storage
+
+
+def _make_fake_connector(tensor: _FakeTensor) -> PleOffloadConnector:
+ connector = PleOffloadConnector.__new__(PleOffloadConnector)
+ # Repeat the same storage in several connector positions to cover aliases.
+ connector._host_output_bufs = {"layer.0": tensor}
+ connector._host_flags = {"layer.0": tensor}
+ connector._reserved_host_storage_filenames = set()
+ connector._input_ids_buf = tensor
+ connector._query_start_loc_buf = tensor
+ connector._ngram_context_buf = None
+ return connector
+
+
+def test_release_host_storage_reserves_each_name_once(monkeypatch) -> None:
+ storage = _FakeFilenameStorage(b"/ple-test-reservation")
+ connector = _make_fake_connector(_FakeTensor(storage))
+ unlinked: list[str] = []
+ monkeypatch.setattr(connector_module.os, "unlink", unlinked.append)
+
+ connector.release_host_storage()
+ connector.release_host_storage()
+
+ assert storage.increfs == 1
+ assert storage.decrefs == 0
+ assert unlinked == ["/dev/shm/ple-test-reservation"]
+
+
+def test_release_host_storage_rolls_back_on_unlink_failure(monkeypatch) -> None:
+ storage = _FakeFilenameStorage(b"/ple-test-rollback")
+ connector = _make_fake_connector(_FakeTensor(storage))
+
+ def fail_unlink(_path: str) -> None:
+ raise PermissionError("test unlink failure")
+
+ monkeypatch.setattr(connector_module.os, "unlink", fail_unlink)
+
+ connector.release_host_storage()
+
+ # The fake connector aliases one storage in four positions. A rolled-back
+ # reservation is retried by the next alias in the same pass (by design: the
+ # name is still present), so every reservation must be balanced by a
+ # rollback and nothing may stay reserved.
+ assert storage.increfs >= 1
+ assert storage.increfs == storage.decrefs
+ assert connector._reserved_host_storage_filenames == set()
+
+
+_BALANCED_LIFETIME_SCRIPT = r"""
+import gc
+import os
+from multiprocessing.reduction import ForkingPickler
+from pathlib import Path
+
+import torch
+import torch.multiprocessing as torch_mp
+
+from vllm.v1.ple_offload.connector import (
+ PleOffloadConnector,
+ _PROCESS_LIFETIME_REGISTERED_STORAGE,
+)
+from vllm.v1.ple_offload.protocol import PleOffloadRegistration
+
+
+def shm_path(tensor):
+ _, filename, _ = tensor.untyped_storage()._share_filename_cpu_()
+ return Path("/dev/shm") / os.fsdecode(filename).lstrip("/")
+
+
+class Progress:
+ def close(self):
+ pass
+
+
+original_strategy = torch_mp.get_sharing_strategy()
+torch_mp.set_sharing_strategy("file_system")
+try:
+ host_outputs = {
+ "layer.0": torch.arange(12, dtype=torch.float32)
+ .reshape(3, 4)
+ .share_memory_(),
+ "layer.1": torch.arange(8, dtype=torch.float32)
+ .reshape(2, 4)
+ .share_memory_(),
+ }
+ host_flags = {
+ name: torch.zeros(16, dtype=torch.int32).share_memory_()
+ for name in host_outputs
+ }
+ input_ids = torch.arange(16, dtype=torch.int32).share_memory_()
+ query_start_loc = torch.arange(5, dtype=torch.int32).share_memory_()
+ ngram_context = torch.arange(12, dtype=torch.int32).reshape(4, 3)
+ ngram_context.share_memory_()
+
+ registration = PleOffloadRegistration(
+ worker_id=1,
+ tp_rank=1,
+ dp_rank=0,
+ host_output_buffers=host_outputs,
+ host_flag_tensors=host_flags,
+ input_ids_buf=input_ids,
+ query_start_loc_buf=query_start_loc,
+ ngram_context_buf=ngram_context,
+ )
+ payload = ForkingPickler.dumps(registration)
+ rebuilt_registration = ForkingPickler.loads(payload)
+finally:
+ torch_mp.set_sharing_strategy(original_strategy)
+
+connector = PleOffloadConnector.__new__(PleOffloadConnector)
+connector._host_output_bufs = host_outputs
+connector._host_flags = host_flags
+connector._reserved_host_storage_filenames = set()
+connector._input_ids_buf = input_ids
+connector._query_start_loc_buf = query_start_loc
+connector._ngram_context_buf = ngram_context
+
+# Exercise the production close ordering without CUDA or transport setup.
+connector._request_thread = None
+connector._pinned_input_buffers = []
+connector._mapped_host_buffers = []
+connector._d2h_event_pool = object()
+connector._registration_socket = None
+connector._zmq_ctx = None
+connector._progress = Progress()
+
+tensors = [
+ *host_outputs.values(),
+ *host_flags.values(),
+ input_ids,
+ query_start_loc,
+ ngram_context,
+]
+assert len(tensors) == 7
+assert len({tensor.untyped_storage()._weak_ref() for tensor in tensors}) == 7
+
+# TP != 0 creates and serializes its input buffers without CUDA registration
+# or the process-lifetime reference used by TP0.
+unpinned_tp_input = input_ids
+assert not connector._pinned_input_buffers
+assert all(
+ unpinned_tp_input is not registered
+ for registered in _PROCESS_LIFETIME_REGISTERED_STORAGE
+)
+
+paths = [shm_path(tensor) for tensor in tensors]
+before = [tensor.clone() for tensor in tensors]
+assert all(path.exists() for path in paths)
+
+connector.close()
+connector.close()
+
+assert all(not path.exists() for path in paths)
+for tensor, expected in zip(tensors, before):
+ assert torch.equal(tensor, expected)
+ tensor.add_(1)
+ assert torch.equal(tensor, expected + 1)
+del tensor, expected
+
+# Remove every original and reconstructed tensor reference. With an unreserved
+# allocator this reaches counter zero after manual unlink and aborts the child.
+connector._host_output_bufs = {}
+connector._host_flags = {}
+connector._input_ids_buf = None
+connector._query_start_loc_buf = None
+connector._ngram_context_buf = None
+del connector
+del rebuilt_registration
+del registration
+del payload
+del unpinned_tp_input
+del tensors
+del before
+del host_outputs
+del host_flags
+del input_ids
+del query_start_loc
+del ngram_context
+gc.collect()
+
+assert all(not path.exists() for path in paths)
+"""
+
+
+def test_close_balanced_filename_storage_lifetime_exits_cleanly() -> None:
+ _require_file_system_shm()
+
+ # Resolve from the connector actually imported by this pytest process. This
+ # lets the red command run this patched test against an unpatched worktree.
+ import_root = Path(connector_module.__file__).resolve().parents[3]
+ env = os.environ.copy()
+ inherited_pythonpath = env.get("PYTHONPATH")
+ env["PYTHONPATH"] = (
+ os.pathsep.join((str(import_root), inherited_pythonpath))
+ if inherited_pythonpath
+ else str(import_root)
+ )
+
+ completed = subprocess.run(
+ [sys.executable, "-c", _BALANCED_LIFETIME_SCRIPT],
+ cwd=import_root,
+ env=env,
+ capture_output=True,
+ text=True,
+ timeout=60,
+ check=False,
+ )
+
+ assert completed.returncode == 0, (
+ "balanced PLE storage destruction failed "
+ f"(exit code {completed.returncode})\n"
+ f"stdout:\n{completed.stdout}\n"
+ f"stderr:\n{completed.stderr}"
+ )
diff --git a/tests/v1/ple_offload/test_shutdown_order.py b/tests/v1/ple_offload/test_shutdown_order.py
new file mode 100644
index 0000000..101c3da
--- /dev/null
+++ b/tests/v1/ple_offload/test_shutdown_order.py
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+from types import SimpleNamespace
+from typing import Any
+
+from vllm.v1.ple_offload.connector import PleOffloadConnector
+from vllm.v1.ple_offload import worker as ple_worker
+from vllm.v1.worker import gpu_worker
+
+
+def test_shutdown_marks_ple_stopping_before_teardown(monkeypatch) -> None:
+ worker = gpu_worker.Worker.__new__(gpu_worker.Worker)
+ worker._ple_offload_stopping = False
+ order: list[tuple[str, bool]] = []
+
+ def record(name: str):
+ def call() -> None:
+ order.append((name, worker._ple_offload_stopping))
+
+ return call
+
+ monkeypatch.setattr(gpu_worker.gc, "unfreeze", record("gc"))
+ monkeypatch.setattr(
+ gpu_worker,
+ "ensure_kv_transfer_shutdown",
+ record("kv"),
+ )
+ monkeypatch.setattr(
+ gpu_worker,
+ "ensure_ec_transfer_shutdown",
+ record("ec"),
+ )
+ monkeypatch.setattr(
+ gpu_worker.current_platform,
+ "is_cuda_alike",
+ lambda: False,
+ )
+
+ worker.profiler = None
+ worker.weight_transfer_engine = None
+ worker.elastic_ep_executor = SimpleNamespace(shutdown=record("elastic"))
+ worker._ple_offload_worker_handle = SimpleNamespace(close=record("offload"))
+ worker.model_runner = SimpleNamespace(shutdown=record("runner"))
+
+ worker.shutdown()
+
+ assert order
+ assert all(stopping for _, stopping in order)
+ assert [name for name, _ in order][:2] == ["gc", "kv"]
+
+
+class _RequestThread:
+ def __init__(self, order: list[str], *, stops: bool) -> None:
+ self.order = order
+ self.stops = stops
+ self.alive = True
+
+ def is_alive(self) -> bool:
+ return self.alive
+
+ def join(self, timeout: float) -> None:
+ self.order.append("thread.join")
+ if self.stops:
+ self.alive = False
+
+
+def _record(order: list[str], name: str):
+ def call(*_args: Any, **_kwargs: Any) -> None:
+ order.append(name)
+
+ return call
+
+
+def _make_connector_for_close(
+ order: list[str], *, thread_stops: bool
+) -> PleOffloadConnector:
+ connector = PleOffloadConnector.__new__(PleOffloadConnector)
+ connector._request_thread = _RequestThread(order, stops=thread_stops)
+ connector._request_queue = SimpleNamespace(put=_record(order, "queue.stop"))
+ connector._pinned_input_buffers = []
+ connector._mapped_host_buffers = []
+ connector._d2h_event_pool = object()
+ connector._registration_socket = SimpleNamespace(
+ close=_record(order, "socket.close")
+ )
+ connector._zmq_ctx = SimpleNamespace(term=_record(order, "context.term"))
+ connector._progress = SimpleNamespace(close=_record(order, "progress.close"))
+ connector.release_host_storage = _record(order, "storage.release")
+ return connector
+
+
+def test_connector_close_releases_storage_last() -> None:
+ order: list[str] = []
+ connector = _make_connector_for_close(order, thread_stops=True)
+
+ connector.close()
+
+ assert order[-1] == "storage.release"
+ assert order.index("thread.join") < order.index("context.term")
+ assert order.index("context.term") < order.index("storage.release")
+
+
+def test_connector_close_defers_storage_release_if_thread_does_not_stop() -> None:
+ order: list[str] = []
+ connector = _make_connector_for_close(order, thread_stops=False)
+
+ connector.close()
+
+ assert "storage.release" not in order
+ assert "socket.close" not in order
+ assert "context.term" not in order
+ assert order[-1] == "progress.close"
+
+
+def test_offload_child_unlinks_owned_ipc_after_context_term(monkeypatch) -> None:
+ order: list[str] = []
+ ipc_addr = "ipc:///tmp/ple-owned-test.sock"
+ pull_socket = SimpleNamespace(
+ bind=_record(order, "socket.bind"),
+ close=_record(order, "socket.close"),
+ )
+ zmq_context = SimpleNamespace(
+ socket=lambda _kind: pull_socket,
+ term=_record(order, "context.term"),
+ )
+
+ class FakeRunner:
+ layer_names: list[str] = []
+
+ def __init__(self, _config: object) -> None:
+ pass
+
+ def accept_registrations(self, _socket: object, _count: int) -> None:
+ order.append("registrations.accept")
+
+ def busy_loop(self, _socket: object, _event: object) -> None:
+ order.append("runner.busy_loop")
+
+ ready_pipe = SimpleNamespace(
+ send_bytes=_record(order, "ready.send"),
+ close=_record(order, "ready.close"),
+ )
+ monkeypatch.setattr(ple_worker, "decorate_logs", lambda _name: None)
+ monkeypatch.setattr(ple_worker, "mark_as_offload_worker", lambda: None)
+ monkeypatch.setattr(ple_worker, "_prepare_cuda_free_process", lambda: None)
+ monkeypatch.setattr(ple_worker, "_init_offload_distributed", lambda: None)
+ monkeypatch.setattr(
+ ple_worker,
+ "set_current_vllm_config",
+ lambda _config: contextlib.nullcontext(),
+ )
+ monkeypatch.setattr(ple_worker, "PleOffloadRunner", FakeRunner)
+ monkeypatch.setattr(ple_worker.zmq, "Context", lambda: zmq_context)
+ monkeypatch.setattr(ple_worker.signal, "signal", lambda *_args: None)
+ monkeypatch.setattr(ple_worker.torch.cuda, "is_initialized", lambda: False)
+ monkeypatch.setattr(
+ ple_worker,
+ "_unlink_ipc_socket",
+ lambda address: order.append(f"ipc.unlink:{address}"),
+ )
+
+ ple_worker.PleOffloadWorker.proc_main(
+ SimpleNamespace(),
+ num_workers=1,
+ ipc_addr=ipc_addr,
+ ready_pipe=ready_pipe,
+ )
+
+ assert order.index("socket.close") < order.index("context.term")
+ assert order[-1] == f"ipc.unlink:{ipc_addr}"
diff --git a/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py
new file mode 100644
index 0000000..c709a93
--- /dev/null
+++ b/tests/v1/ple_offload/test_spawn_bootstrap_cuda.py
@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import ctypes
+import ctypes.util
+import inspect
+import json
+import multiprocessing
+import os
+import pickle
+import signal
+import socket
+import subprocess
+import sys
+import threading
+import time
+from multiprocessing.connection import Connection
+from pathlib import Path
+
+import pytest
+
+def _native_cuda_probe(output_path: str) -> None:
+ """Run during config unpickling, before the bootstrap imports vLLM."""
+ library = ctypes.util.find_library("cudart")
+ if library is None:
+ os._exit(77)
+ cudart = ctypes.CDLL(library)
+ result = int(cudart.cudaFree(ctypes.c_void_p()))
+ Path(output_path).write_text(
+ json.dumps(
+ {
+ "cuda_visible_devices": os.environ.get(
+ "CUDA_VISIBLE_DEVICES"
+ ),
+ "cuda_result": result,
+ "torch_imported": "torch" in sys.modules,
+ "vllm_imported": "vllm" in sys.modules,
+ }
+ ),
+ encoding="utf-8",
+ )
+ os._exit(0)
+
+
+class _ProbePayload:
+ def __init__(self, output_path: str) -> None:
+ self.output_path = output_path
+
+ def __reduce__(self):
+ # The bootstrap interpreter cannot import this test module, so rebuild
+ # the probe from its source text through a builtin instead of by name.
+ source = (
+ "import ctypes, ctypes.util, json, os, sys\n"
+ "from pathlib import Path\n"
+ + inspect.getsource(_native_cuda_probe)
+ + f"_native_cuda_probe({self.output_path!r})\n"
+ )
+ return exec, (source, {})
+
+
+def test_spawn_bootstrap_sets_visibility_before_native_cuda(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from vllm.v1.ple_offload.worker import PleOffloadWorker
+
+ if ctypes.util.find_library("cudart") is None:
+ pytest.skip("CUDA runtime is not installed")
+
+ original = os.environ.get("CUDA_VISIBLE_DEVICES")
+ output_path = tmp_path / "probe.json"
+
+ # Use the production Popen launcher. Config reconstruction executes the
+ # probe in the fresh bootstrap interpreter before worker imports.
+ with monkeypatch.context() as child_settings:
+ child_settings.setenv("CUDA_VISIBLE_DEVICES", "0")
+ handle = PleOffloadWorker.make_process(
+ _ProbePayload(str(output_path)), 0, "unused"
+ )
+ try:
+ returncode = handle.proc.wait(timeout=120)
+ finally:
+ handle.close()
+
+ if returncode == 77:
+ pytest.skip("spawned child could not load the CUDA runtime")
+ assert returncode == 0
+
+ result = json.loads(output_path.read_text(encoding="utf-8"))
+ assert result["cuda_visible_devices"] == "0"
+ assert result["cuda_result"] == 0
+ assert result["torch_imported"] is False
+ assert result["vllm_imported"] is False
+ assert os.environ.get("CUDA_VISIBLE_DEVICES") == original
+
+
+class _FreshInterpreterProbe:
+ def __init__(self, output_path: str) -> None:
+ self.output_path = output_path
+
+ def __reduce__(self):
+ expression = (
+ "(__import__('pathlib').Path("
+ + repr(self.output_path)
+ + ").write_text(__import__('json').dumps({"
+ "'cuda_visible_devices': __import__('os').environ.get("
+ "'CUDA_VISIBLE_DEVICES'),"
+ "'torch_imported': 'torch' in __import__('sys').modules,"
+ "'vllm_imported': 'vllm' in __import__('sys').modules"
+ "}), encoding='utf-8'), None)[1]"
+ )
+ return eval, (expression,)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only")
+def test_fresh_interpreter_has_child_environment_before_vllm_import(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from vllm.v1.ple_offload.worker import PleOffloadWorker
+
+ # This test module has unconditional top-level imports. The dedicated
+ # bootstrap interpreter must not import it as __mp_main__.
+ original = os.environ.get("CUDA_VISIBLE_DEVICES")
+ output_path = tmp_path / "fresh-interpreter.json"
+
+ with monkeypatch.context() as child_settings:
+ child_settings.setenv("CUDA_VISIBLE_DEVICES", "child-only")
+ handle = PleOffloadWorker.make_process(
+ _FreshInterpreterProbe(str(output_path)), 0, "unused"
+ )
+ try:
+ handle.proc.wait(timeout=120)
+ finally:
+ handle.close()
+
+ assert output_path.exists()
+ result = json.loads(output_path.read_text(encoding="utf-8"))
+ assert result["cuda_visible_devices"] == "child-only"
+ assert result["torch_imported"] is False
+ assert result["vllm_imported"] is False
+ assert os.environ.get("CUDA_VISIBLE_DEVICES") == original
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only")
+def test_launch_payload_send_has_absolute_deadline() -> None:
+ reader, sender = socket.socketpair()
+ sender.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024)
+ code = (
+ "import sys,time\n"
+ "from vllm.v1.ple_offload.worker import _send_framed_payload\n"
+ "payload=b'x'*(64<<20)\n"
+ "start=time.monotonic()\n"
+ "try:\n"
+ " _send_framed_payload(int(sys.argv[1]),payload,start+0.25)\n"
+ "except TimeoutError:\n"
+ " print('TIMED_OUT',time.monotonic()-start,flush=True)\n"
+ "else:\n"
+ " raise AssertionError('blocked reader unexpectedly accepted payload')\n"
+ )
+ process = subprocess.Popen(
+ [sys.executable, "-c", code, str(sender.fileno())],
+ pass_fds=(sender.fileno(),),
+ close_fds=True,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ text=True,
+ )
+ sender.close()
+ try:
+ stdout, stderr = process.communicate(timeout=120)
+ finally:
+ reader.close()
+ if process.poll() is None:
+ process.terminate()
+ process.wait(timeout=5)
+ assert process.returncode == 0, stderr
+ # The injector logs the send entry through vLLM's logger on stdout; read only
+ # the probe's own result line.
+ result = [line for line in stdout.splitlines() if line.startswith('TIMED_OUT ')]
+ assert len(result) == 1, stdout
+ marker, elapsed = result[0].split()
+ assert marker == "TIMED_OUT"
+ assert float(elapsed) < 5
+
+
+class _BlockingConfigProbe:
+ def __init__(self, output_path: str) -> None:
+ self.output_path = output_path
+
+ def __reduce__(self):
+ expression = (
+ "(__import__('pathlib').Path("
+ + repr(self.output_path)
+ + ").write_text(str(__import__('os').getpid()), encoding='utf-8'),"
+ "__import__('threading').Event().wait())[1]"
+ )
+ return eval, (expression,)
+
+
+def _launch_blocked_config(output_path: str, child_pids) -> None:
+ from vllm.v1.ple_offload.worker import PleOffloadWorker
+
+ handle = PleOffloadWorker.make_process(
+ _BlockingConfigProbe(output_path), 0, "unused"
+ )
+ child_pids.put(handle.proc.pid)
+ # The test kills this launcher parent. Do not close the handle here.
+ threading.Event().wait()
+
+
+def _pid_has_exited(pid: int) -> bool:
+ stat_path = Path(f"/proc/{pid}/stat")
+ if not stat_path.exists():
+ return True
+ try:
+ return stat_path.read_text(encoding="utf-8").split()[2] == "Z"
+ except FileNotFoundError:
+ return True
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="uses Linux process state")
+def test_parent_death_during_config_reconstruction_exits_child(
+ tmp_path: Path,
+) -> None:
+ context = multiprocessing.get_context("spawn")
+ child_pids = context.Queue()
+ marker = tmp_path / "blocked-config.pid"
+ parent = context.Process(
+ target=_launch_blocked_config,
+ args=(str(marker), child_pids),
+ )
+ parent.start()
+ child_pid = child_pids.get(timeout=120)
+ deadline = time.monotonic() + 120
+ while not marker.exists() and time.monotonic() < deadline:
+ time.sleep(0.05)
+ assert marker.exists()
+
+ os.kill(parent.pid, signal.SIGKILL)
+ parent.join(timeout=120)
+ assert not parent.is_alive()
+
+ deadline = time.monotonic() + 10
+ while not _pid_has_exited(child_pid) and time.monotonic() < deadline:
+ time.sleep(0.05)
+ if not _pid_has_exited(child_pid):
+ # Harness cleanup is performed only after the autonomous-exit assertion
+ # has failed.
+ os.kill(child_pid, signal.SIGKILL)
+ pytest.fail("PLE child survived launcher-parent death during unpickle")
+
+
+class _ExecFdProbe:
+ def __init__(
+ self,
+ output_path: str,
+ inherited: list[tuple[int, int]],
+ ) -> None:
+ self.output_path = output_path
+ self.inherited = inherited
+
+ def __reduce__(self):
+ helper = (
+ "import json,os,sys\n"
+ "from pathlib import Path\n"
+ "pairs=json.loads(sys.argv[3])\n"
+ "retained=[]\n"
+ "for fd,inode in pairs:\n"
+ " try:\n"
+ " retained.append(os.fstat(fd).st_ino==inode)\n"
+ " except OSError:\n"
+ " retained.append(False)\n"
+ "Path(sys.argv[1]).write_text(json.dumps({"
+ "'noninheritable_before_exec':sys.argv[2]=='True',"
+ "'closed_after_exec':not any(retained)}),encoding='utf-8')\n"
+ )
+ pairs = json.dumps(self.inherited)
+ expression = (
+ "__import__('os').execv(__import__('sys').executable,"
+ "[__import__('sys').executable,'-c',"
+ + repr(helper)
+ + ","
+ + repr(self.output_path)
+ + ",str(all(not __import__('os').get_inheritable(int(fd)) "
+ "for fd in __import__('sys').argv[1:])),"
+ + repr(pairs)
+ + "])"
+ )
+ return eval, (expression,)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only")
+def test_actual_bootstrap_restores_cloexec_before_helper_exec(
+ tmp_path: Path,
+) -> None:
+ from vllm.v1.ple_offload.worker import _send_framed_payload
+
+ control_parent, control_child = socket.socketpair()
+ death_parent, death_child = socket.socketpair()
+ inherited = [
+ (control_child.fileno(), os.fstat(control_child.fileno()).st_ino),
+ (death_child.fileno(), os.fstat(death_child.fileno()).st_ino),
+ ]
+ output_path = tmp_path / "bootstrap-fds.json"
+ config_payload = pickle.dumps(
+ _ExecFdProbe(str(output_path), inherited),
+ protocol=pickle.HIGHEST_PROTOCOL,
+ )
+ payload = pickle.dumps(
+ {
+ "vllm_config_payload": config_payload,
+ "num_workers": 0,
+ "ipc_addr": "unused",
+ },
+ protocol=pickle.HIGHEST_PROTOCOL,
+ )
+ control = Connection(control_parent.detach())
+ death_writer = Connection(death_parent.detach())
+ process = subprocess.Popen(
+ [
+ sys.executable,
+ "-m",
+ "vllm_ple_offload_bootstrap",
+ str(control_child.fileno()),
+ str(death_child.fileno()),
+ ],
+ close_fds=True,
+ pass_fds=(control_child.fileno(), death_child.fileno()),
+ )
+ control_child.close()
+ death_child.close()
+ try:
+ _send_framed_payload(
+ control.fileno(), payload, time.monotonic() + 10
+ )
+ assert process.wait(timeout=120) == 0
+ finally:
+ control.close()
+ death_writer.close()
+ if process.poll() is None:
+ process.terminate()
+ try:
+ process.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ process.kill()
+ process.wait(timeout=5)
+
+ result = json.loads(output_path.read_text(encoding="utf-8"))
+ assert result == {
+ "noninheritable_before_exec": True,
+ "closed_after_exec": True,
+ }
diff --git a/tests/v1/ple_offload/test_test_fault.py b/tests/v1/ple_offload/test_test_fault.py
new file mode 100644
index 0000000..3dc8470
--- /dev/null
+++ b/tests/v1/ple_offload/test_test_fault.py
@@ -0,0 +1,293 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+import os
+import queue
+import sys
+from dataclasses import FrozenInstanceError
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+
+@pytest.fixture
+def fault_module(monkeypatch, tmp_path):
+ monkeypatch.delenv("VLLM_PLE_TEST_FAULT", raising=False)
+ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log"))
+ import vllm_ple_test_fault as test_fault
+
+ return test_fault
+
+
+def test_disarmed_import_is_none(fault_module):
+ assert fault_module._FAULT is None
+
+
+@pytest.mark.parametrize(
+ "value",
+ [
+ "cut=exc_before_prepare,ordinal=60,rank=any",
+ "cut=exc_after_prepare,ordinal=2,tp=1",
+ "cut=exc_in_release,ordinal=3,pp=0",
+ "cut=park_before_publish,ordinal=4,rank=any,phase=any",
+ ],
+)
+def test_valid_specs(fault_module, value):
+ spec = fault_module._parse_fault_spec(value)
+ assert spec is not None
+ assert spec.ordinal > 0
+
+
+@pytest.mark.parametrize(
+ "value",
+ [
+ "",
+ " cut=exc_in_release,ordinal=1,rank=any",
+ "cut=unknown,ordinal=1,rank=any",
+ "cut=exc_in_release,ordinal=0,rank=any",
+ "cut=exc_in_release,ordinal=01,rank=any",
+ "cut=exc_in_release,ordinal=1",
+ "cut=exc_in_release,ordinal=1,rank=tp1",
+ "cut=exc_in_release,ordinal=1,rank=any,tp=0",
+ "cut=exc_in_release,ordinal=1,tp=-1",
+ "cut=exc_in_release,ordinal=1,tp=0,phase=startup",
+ "cut=exc_in_release,ordinal=1,tp=0,extra=x",
+ "cut=exc_in_release,cut=exc_after_release,ordinal=1,tp=0",
+ ],
+)
+def test_malformed_specs_raise(fault_module, value):
+ with pytest.raises(ValueError):
+ fault_module._parse_fault_spec(value)
+
+
+def test_one_shot_rank_phase_and_ordinal(fault_module, monkeypatch):
+ spec = fault_module._parse_fault_spec(
+ "cut=exc_after_prepare,ordinal=2,tp=1,pp=0"
+ )
+ assert spec is not None
+
+ monkeypatch.setattr(fault_module.os, "getpid", lambda: 123)
+ monkeypatch.setattr(fault_module, "_pid_start_time", lambda _pid=None: 789)
+ monkeypatch.setattr(fault_module.time, "monotonic", lambda: 12.5)
+ monkeypatch.setattr(fault_module.time, "time", lambda: 456.25)
+
+ def fires(cut, tp_rank=1, eligible=True):
+ return spec.fires(
+ cut,
+ role="gpu",
+ unit="runner-forward",
+ eligible=eligible,
+ tp_rank=tp_rank,
+ pp_rank=0,
+ dp_rank=0,
+ transaction_id="sequence:9@begin:12.000000000",
+ correlation_id="sequence:9",
+ )
+
+ assert not fires("other")
+ assert not fires("exc_after_prepare", tp_rank=0)
+ assert not fires("exc_after_prepare", eligible=False)
+ assert not fires("exc_after_prepare")
+ assert fires("exc_after_prepare")
+ assert not fires("exc_after_prepare")
+ line = Path(os.environ["VLLM_PLE_TEST_FAULT_LOG"]).read_text()
+ assert line == (
+ "PLE_TEST_FAULT v=1 cut=exc_after_prepare ordinal=2 count=2 "
+ "unit=runner-forward role=gpu tp=1 pp=0 dp=0 pid=123 "
+ "pid_start=789 transaction=sequence:9@begin:12.000000000 "
+ "correlation=sequence:9 monotonic=12.500000000 "
+ "wall=456.250000000\n"
+ )
+ with pytest.raises(FrozenInstanceError):
+ spec.ordinal = 3
+
+
+def test_disarmed_connector_sites_continue(fault_module, monkeypatch):
+ import vllm.v1.ple_offload.connector as connector_module
+
+ monkeypatch.setattr(connector_module, "_FAULT", None)
+
+ connector = object.__new__(connector_module.PleOffloadConnector)
+ connector.tp_rank = 0
+ connector.pp_rank = 0
+ connector.dp_rank = 0
+ connector.device = SimpleNamespace(index=0)
+ connector._d2h_event_pool = queue.Queue()
+ event = Mock()
+ connector._d2h_event_pool.put_nowait(event)
+ connector._request_queue = Mock()
+ connector._progress = SimpleNamespace(sequence=0)
+ connector._enqueue_cuda_inputs = Mock()
+
+ connector._launch(
+ num_reqs=1,
+ num_tokens=2,
+ serving=True,
+ transaction_id=None,
+ correlation_id=None,
+ )
+ connector._enqueue_cuda_inputs.assert_called_once()
+ connector._request_queue.put_nowait.assert_called_once()
+
+ staged = connector_module._PendingPleOffloadRequest(
+ connector_module.PleOffloadRequest(dp_rank=0, num_tokens=2, num_reqs=1),
+ event,
+ True,
+ "sequence:0",
+ )
+ socket = Mock()
+ monkeypatch.setattr(
+ connector_module.torch.accelerator,
+ "device_index",
+ lambda _index: contextlib.nullcontext(),
+ )
+ connector._process_request(staged, socket)
+ event.synchronize.assert_called_once()
+ socket.send.assert_called_once()
+
+
+def test_disarmed_release_site_continues(fault_module, monkeypatch):
+ import vllm.v1.ple_offload.connector as connector_module
+
+ monkeypatch.setattr(connector_module, "_FAULT", None)
+ monkeypatch.setattr(connector_module._rowdump, "enabled", lambda: False)
+
+ stream = Mock()
+ completion = Mock()
+ monkeypatch.setattr(
+ connector_module.torch.cuda,
+ "current_stream",
+ lambda _device: stream,
+ )
+ monkeypatch.setattr(connector_module.torch.cuda, "Event", lambda: completion)
+
+ connector = object.__new__(connector_module.PleOffloadConnector)
+ connector.tp_rank = 0
+ connector.pp_rank = 0
+ connector.device = SimpleNamespace(index=0)
+ connector._layers = {}
+ connector._after_dummy = False
+ connector._progress = Mock()
+
+ connector._release_outputs_impl()
+ completion.record.assert_called_once_with(stream)
+ connector._progress.track_completion.assert_called_once_with(completion)
+
+
+def test_disarmed_worker_park_site_continues(fault_module, monkeypatch):
+ import vllm.v1.ple_offload.worker as worker_module
+
+ monkeypatch.setattr(worker_module, "_FAULT", None)
+ runner = object.__new__(worker_module.PleOffloadRunner)
+ runner._validate_requests_before_publication = Mock()
+ runner._worker_targets = {0: {}}
+ runner._input_bufs = {}
+ runner._scratch_bufs = {}
+ runner._layers = {}
+ runner._clamp_input_ids = False
+ runner._step = 0
+ runner._request_index = 0
+ request = SimpleNamespace(dp_rank=0)
+
+ runner._handle_requests_impl([request])
+ runner._validate_requests_before_publication.assert_called_once_with([request])
+ assert runner._step == 1
+
+
+def test_all_injection_sites_use_the_disarmed_short_circuit():
+ root = Path(__file__).resolve().parents[3]
+ expected = {
+ "vllm/v1/ple_offload/connector.py": 3,
+ "vllm/v1/ple_offload/worker.py": 1,
+ "vllm/v1/worker/gpu/model_runner.py": 4,
+ "vllm_ple_offload_bootstrap/__init__.py": 1,
+ }
+ for relative, count in expected.items():
+ lines = (root / relative).read_text(encoding="utf-8").splitlines()
+ sites = [i for i, line in enumerate(lines) if "_FAULT.fires(" in line]
+ assert len(sites) == count, relative
+ for i in sites:
+ window = "\n".join(lines[max(0, i - 6): i + 1])
+ # The None check must precede fires() in the same condition, so a
+ # disarmed process never calls into the fault module.
+ assert "_FAULT is not None" in window, (relative, i + 1)
+ assert window.index("_FAULT is not None") < window.index("_FAULT.fires("), (relative, i + 1)
+
+
+class _FatalReached(Exception):
+ pass
+
+
+def test_exc_in_release_reaches_real_connector_fatal_wrapper(
+ fault_module, monkeypatch
+):
+ import vllm.v1.ple_offload.connector as connector_module
+
+ armed = fault_module._parse_fault_spec(
+ "cut=exc_in_release,ordinal=1,tp=1,pp=0"
+ )
+ assert armed is not None
+ monkeypatch.setattr(connector_module, "_FAULT", armed)
+
+ fatal = Mock(side_effect=_FatalReached)
+ monkeypatch.setattr(connector_module, "ple_fail_closed", fatal)
+
+ connector = object.__new__(connector_module.PleOffloadConnector)
+ connector.rank = 1
+ connector.tp_rank = 1
+ connector.pp_rank = 0
+ connector.dp_rank = 0
+ connector._layers = {}
+ connector._after_dummy = False
+ connector._progress = SimpleNamespace(sequence=9)
+
+ with pytest.raises(_FatalReached):
+ connector.release_outputs(
+ serving=True,
+ transaction_id="sequence:9@begin:1.000000000",
+ correlation_id="sequence:9",
+ )
+
+ assert fatal.call_args.kwargs["where"] == "release_outputs"
+ assert fatal.call_args.kwargs["rank"] == 1
+ assert isinstance(fatal.call_args.kwargs["error"], RuntimeError)
+
+
+@pytest.mark.timeout(180)
+@pytest.mark.skipif(
+ sys.platform == "win32",
+ reason="production PLE launcher requires POSIX pass_fds",
+)
+@pytest.mark.parametrize("cut", ["launch_partial_ready", "launch_partial_ready_header"])
+def test_partial_ready_times_out_and_terminates_production_child(
+ monkeypatch,
+ tmp_path,
+ cut,
+):
+ import vllm.v1.ple_offload.worker as worker_module
+
+ monkeypatch.setenv(
+ "VLLM_PLE_TEST_FAULT",
+ f"cut={cut},ordinal=1,rank=any",
+ )
+ monkeypatch.setenv("VLLM_PLE_TEST_FAULT_LOG", str(tmp_path / "fault.log"))
+ monkeypatch.setattr(
+ worker_module.envs,
+ "VLLM_PLE_OFFLOAD_READY_TIMEOUT",
+ 1.0,
+ )
+
+ handle = worker_module.PleOffloadWorker.make_process(
+ object(), # The bootstrap fault fires before config reconstruction.
+ num_workers=1,
+ ipc_addr="ipc:///unused-ple-test",
+ )
+ try:
+ with pytest.raises(TimeoutError):
+ worker_module.PleOffloadWorker.wait_for_ready(handle)
+ assert handle.proc.poll() is not None
+ finally:
+ handle.close(graceful=False)
diff --git a/tests/v1/ple_offload/test_transport_bounded_io.py b/tests/v1/ple_offload/test_transport_bounded_io.py
new file mode 100644
index 0000000..fb37122
--- /dev/null
+++ b/tests/v1/ple_offload/test_transport_bounded_io.py
@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import multiprocessing
+import os
+import signal
+import struct
+import subprocess
+import sys
+import time
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+
+class _BlockingStrError(Exception):
+ def __str__(self) -> str:
+ __import__("threading").Event().wait()
+ return "unreachable"
+
+
+def _blocking_fatal_formatter_case(marker_path: str) -> None:
+ from vllm.v1.ple_offload.fatal import ple_fail_closed
+
+ Path(marker_path).write_text("armed", encoding="utf-8")
+ try:
+ raise _BlockingStrError()
+ except BaseException as error:
+ ple_fail_closed(
+ where="blocking-formatter",
+ rank=1,
+ layer="ple",
+ sequence=3,
+ error=error,
+ )
+
+
+def test_fatal_helper_does_not_call_blocking_exception_str(
+ tmp_path: Path,
+) -> None:
+ context = multiprocessing.get_context("spawn")
+ marker = tmp_path / "fatal-formatter-armed"
+ process = context.Process(
+ target=_blocking_fatal_formatter_case,
+ args=(str(marker),),
+ )
+ process.start()
+ try:
+ deadline = time.monotonic() + 120
+ while not marker.exists() and time.monotonic() < deadline:
+ time.sleep(0.05)
+ assert marker.exists()
+
+ process.join(timeout=5)
+ assert not process.is_alive()
+ assert process.exitcode not in (None, 0)
+ finally:
+ if process.is_alive():
+ process.terminate()
+ process.join(timeout=120)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only")
+@pytest.mark.parametrize(
+ ("behavior", "ready_timeout", "expected_exception"),
+ [
+ pytest.param("partial", "1", TimeoutError, id="partial-frame"),
+ pytest.param("silent", "1", TimeoutError, id="no-response"),
+ pytest.param("ready", "120", None, id="complete-ready"),
+ pytest.param("invalid", "120", RuntimeError, id="invalid-layer-names"),
+ ],
+)
+def test_wait_for_ready_receives_complete_bounded_frame(
+ monkeypatch: pytest.MonkeyPatch,
+ behavior: str,
+ ready_timeout: str,
+ expected_exception: type[BaseException] | None,
+) -> None:
+ import vllm.v1.ple_offload.worker as worker_module
+
+ real_popen = subprocess.Popen
+ spawned: list[subprocess.Popen] = []
+ child_code = (
+ "import os,signal,struct,sys,time\n"
+ "from multiprocessing.connection import Connection\n"
+ "signal.alarm(10)\n"
+ "fd=int(sys.argv[1])\n"
+ "behavior=sys.argv[3]\n"
+ "if behavior=='partial':\n"
+ " os.write(fd,struct.pack('!i',128)+b'partial')\n"
+ "elif behavior=='ready':\n"
+ " Connection(fd).send({'status':'READY','layer_names':['ple']})\n"
+ "elif behavior=='invalid':\n"
+ " Connection(fd).send({'status':'READY','layer_names':'ple'})\n"
+ "time.sleep(120)\n"
+ )
+
+ def launch_fixture(argv: list[str], **kwargs: Any) -> subprocess.Popen:
+ process = real_popen(
+ [
+ sys.executable,
+ "-c",
+ child_code,
+ argv[-2],
+ argv[-1],
+ behavior,
+ ],
+ **kwargs,
+ )
+ spawned.append(process)
+ return process
+
+ monkeypatch.setenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", ready_timeout)
+ monkeypatch.setattr(worker_module.subprocess, "Popen", launch_fixture)
+
+ handle = worker_module.PleOffloadWorker.make_process(None, 0, "unused")
+ try:
+ started = time.monotonic()
+ if expected_exception is None:
+ worker_module.PleOffloadWorker.wait_for_ready(handle)
+ assert handle.ready_pipe_reader is None
+ else:
+ with pytest.raises(expected_exception):
+ worker_module.PleOffloadWorker.wait_for_ready(handle)
+ if expected_exception is TimeoutError:
+ assert time.monotonic() - started < 5
+ assert spawned[0].poll() is not None
+ finally:
+ handle.close(graceful=False)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only")
+def test_make_process_large_payload_send_has_absolute_deadline(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ import vllm.v1.ple_offload.worker as worker_module
+
+ real_popen = subprocess.Popen
+ spawned: list[subprocess.Popen] = []
+
+ def launch_nonreader(argv: list[str], **kwargs: Any) -> subprocess.Popen:
+ process = real_popen(
+ [
+ sys.executable,
+ "-c",
+ "import signal,time; signal.alarm(120); time.sleep(120)",
+ argv[-2],
+ argv[-1],
+ ],
+ **kwargs,
+ )
+ spawned.append(process)
+ return process
+
+ monkeypatch.setenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "5")
+ monkeypatch.setattr(worker_module.subprocess, "Popen", launch_nonreader)
+
+ started = time.monotonic()
+ with pytest.raises(TimeoutError, match="transmission"):
+ worker_module.PleOffloadWorker.make_process(
+ b"x" * (64 << 20),
+ 0,
+ "unused",
+ )
+
+ assert time.monotonic() - started < 20
+ assert spawned
+ assert spawned[0].poll() is not None
+
+
+@pytest.mark.parametrize(
+ "identities",
+ [
+ pytest.param(("shared", "shared"), id="duplicate"),
+ pytest.param(("<unbound>",), id="unbound"),
+ ],
+)
+def test_expected_ple_layers_are_bound_and_unique_before_publication(
+ monkeypatch: pytest.MonkeyPatch,
+ identities: tuple[str, ...],
+) -> None:
+ from vllm.config import CUDAGraphMode
+ from vllm.v1.worker.gpu.cudagraph_utils import (
+ CudaGraphManager,
+ ModelCudaGraphManager,
+ )
+
+ manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager)
+ manager.graphs = {}
+ manager.use_breakable_cg = False
+ manager.cudagraph_mode = CUDAGraphMode.FULL
+
+ modules = [
+ (
+ f"module_{index}",
+ SimpleNamespace(
+ _is_cpu_offloaded=True,
+ _host_output_ptr=1,
+ _ple_offload_name=identity,
+ ),
+ )
+ for index, identity in enumerate(identities)
+ ]
+ model = SimpleNamespace(named_modules=lambda: iter(modules))
+
+ def publish_graph(self: Any, *_args: Any, **_kwargs: Any) -> None:
+ self.graphs["published"] = object()
+
+ monkeypatch.setattr(CudaGraphManager, "capture", publish_graph)
+
+ with pytest.raises(RuntimeError, match="distinct, bound"):
+ ModelCudaGraphManager.capture(
+ manager,
+ model=model,
+ model_state=object(),
+ input_buffers=object(),
+ intermediate_tensors=None,
+ block_tables=object(),
+ attn_groups=[],
+ kv_cache_config=object(),
+ )
+ assert manager.graphs == {}
+
+
+def test_real_wait_copy_geometry_rejected_before_graph_publication(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ import torch
+
+ import vllm.model_executor.layers.ple_offload_layer as ple_layer_module
+ from vllm.config import CUDAGraphMode
+ from vllm.v1.worker.gpu.cudagraph_utils import (
+ BatchExecutionDescriptor,
+ CudaGraphManager,
+ )
+
+ manager = CudaGraphManager.__new__(CudaGraphManager)
+ manager.graphs = {}
+ manager._ple_full_copy_geometry = {}
+ manager._expected_ple_copy_layers = ("ple",)
+ manager._allow_ple_full_graph = True
+ desc = BatchExecutionDescriptor(
+ cg_mode=CUDAGraphMode.FULL,
+ num_tokens=8,
+ num_reqs=1,
+ )
+
+ success = SimpleNamespace(value=0)
+ monkeypatch.setattr(
+ ple_layer_module.torch.cuda,
+ "current_stream",
+ lambda: SimpleNamespace(cuda_stream=1),
+ )
+ monkeypatch.setattr(
+ ple_layer_module.cuda_driver,
+ "CUstream",
+ lambda value: value,
+ )
+ monkeypatch.setattr(
+ ple_layer_module.cuda_driver,
+ "CUdeviceptr",
+ lambda value: value,
+ )
+ monkeypatch.setattr(
+ ple_layer_module.cuda_driver,
+ "cuStreamWaitValue32",
+ lambda *_args: (success,),
+ )
+ monkeypatch.setattr(
+ ple_layer_module.cuda_driver,
+ "cuMemcpyHtoDAsync",
+ lambda *_args: (success,),
+ )
+
+ gpu_output_buffer = torch.empty(12, 4, dtype=torch.float32)
+ hidden_states = torch.empty(12, 1, dtype=torch.float32)
+
+ def capture_then_publish() -> None:
+ with manager.capture_ple_copy_geometry(desc):
+ ple_layer_module._ple_offload_wait_copy_impl(
+ gpu_output_buffer,
+ hidden_states,
+ 1,
+ 2,
+ 16,
+ 12,
+ "ple",
+ )
+ manager.graphs[desc] = object()
+
+ with pytest.raises(RuntimeError, match="actual PLE copy geometry"):
+ capture_then_publish()
+ assert desc not in manager._ple_full_copy_geometry
+ assert desc not in manager.graphs
+
+
+class _CleanupConfig:
+ def __init__(self, started_path: str, cleaned_path: str) -> None:
+ self.started_path = started_path
+ self.cleaned_path = cleaned_path
+
+ def __reduce__(self):
+ source = (
+ "import signal,threading\n"
+ "from pathlib import Path\n"
+ "from vllm.v1.ple_offload.worker import PleOffloadWorker\n"
+ f"_started=Path({self.started_path!r})\n"
+ f"_cleaned=Path({self.cleaned_path!r})\n"
+ "def _cleanup_proc_main(vllm_config,num_workers,ipc_addr,ready_pipe):\n"
+ " _stop=threading.Event()\n"
+ " signal.signal(signal.SIGTERM,lambda *_args:_stop.set())\n"
+ " _started.write_text('started',encoding='utf-8')\n"
+ " try:\n"
+ " _stop.wait()\n"
+ " finally:\n"
+ " _cleaned.write_text('cleaned',encoding='utf-8')\n"
+ "PleOffloadWorker.proc_main=staticmethod(_cleanup_proc_main)\n"
+ )
+ return exec, (source, {})
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="pass_fds is POSIX-only")
+def test_normal_close_runs_child_cleanup(tmp_path: Path) -> None:
+ from vllm.v1.ple_offload.worker import PleOffloadWorker
+
+ started = tmp_path / "cleanup-started"
+ cleaned = tmp_path / "cleanup-finished"
+ handle = PleOffloadWorker.make_process(
+ _CleanupConfig(str(started), str(cleaned)),
+ 0,
+ "unused",
+ )
+ try:
+ deadline = time.monotonic() + 120
+ while not started.exists() and time.monotonic() < deadline:
+ time.sleep(0.05)
+ assert started.exists()
+ finally:
+ handle.close()
+
+ assert handle.proc.returncode == 0
+ assert cleaned.read_text(encoding="utf-8") == "cleaned"
diff --git a/tests/v1/ple_offload/test_transport_fatal_helper.py b/tests/v1/ple_offload/test_transport_fatal_helper.py
new file mode 100644
index 0000000..4e5d5c6
--- /dev/null
+++ b/tests/v1/ple_offload/test_transport_fatal_helper.py
@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import pytest
+
+
+def _run_bounded_child(
+ source: str,
+ tmp_path: Path,
+ *arguments: str,
+ expected_exitcode: int = 0,
+) -> None:
+ marker = tmp_path / f"started-{time.monotonic_ns()}"
+ process = subprocess.Popen(
+ [sys.executable, "-c", source, str(marker), *arguments],
+ stderr=subprocess.DEVNULL,
+ start_new_session=os.name == "posix",
+ )
+ stayed_alive = False
+ try:
+ deadline = time.monotonic() + 120
+ while not marker.exists() and time.monotonic() < deadline:
+ time.sleep(0.05)
+ assert marker.exists(), (
+ "child did not finish cold startup within 120 seconds"
+ )
+ process.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ stayed_alive = True
+ finally:
+ if process.poll() is None:
+ if os.name == "posix":
+ os.killpg(process.pid, signal.SIGKILL)
+ else:
+ process.terminate()
+ process.wait(timeout=120)
+
+ assert not stayed_alive, "child stayed alive after the bounded operation"
+ assert process.returncode == expected_exitcode
+
+
+_FATAL_SOURCE = r"""
+import os
+import sys
+import threading
+from pathlib import Path
+
+from vllm.v1.ple_offload.fatal import ple_fail_closed
+
+scenario = sys.argv[2]
+
+if scenario == "blocking-name-property":
+ class Meta(type):
+ @property
+ def __name__(cls):
+ threading.Event().wait()
+
+ class Error(Exception, metaclass=Meta):
+ pass
+
+ error = Error()
+elif scenario == "blocking-str-descriptor":
+ class Descriptor:
+ def __get__(self, instance, owner):
+ threading.Event().wait()
+
+ class Error(Exception):
+ __str__ = Descriptor()
+
+ error = Error()
+elif scenario == "blocking-name-format":
+ class BlockingName(str):
+ def __format__(self, format_spec):
+ threading.Event().wait()
+
+ class Meta(type):
+ @property
+ def __name__(cls):
+ return BlockingName("Error")
+
+ class Error(Exception, metaclass=Meta):
+ pass
+
+ error = Error()
+elif scenario == "blocking-str-method":
+ class Error(Exception):
+ def __str__(self):
+ threading.Event().wait()
+
+ error = Error()
+elif scenario == "overridden-getattribute":
+ class Error(Exception):
+ def __getattribute__(self, name):
+ threading.Event().wait()
+
+ error = Error("detail")
+elif scenario == "str-subclass-argument":
+ class BlockingString(str):
+ def __format__(self, format_spec):
+ threading.Event().wait()
+
+ def __str__(self):
+ threading.Event().wait()
+
+ error = RuntimeError(BlockingString("detail"))
+elif scenario == "huge-string":
+ error = RuntimeError("x" * (32 << 20))
+else:
+ raise AssertionError(scenario)
+
+Path(sys.argv[1]).write_text("armed", encoding="utf-8")
+ple_fail_closed(
+ where="transport-fatal",
+ rank=1,
+ layer="ple",
+ sequence=3,
+ error=error,
+)
+"""
+
+
+@pytest.mark.parametrize(
+ "scenario",
+ [
+ "blocking-name-property",
+ "blocking-str-descriptor",
+ "blocking-name-format",
+ ],
+)
+def test_fatal_helper_avoids_exception_controlled_descriptors(
+ tmp_path: Path,
+ scenario: str,
+) -> None:
+ _run_bounded_child(
+ _FATAL_SOURCE,
+ tmp_path,
+ scenario,
+ expected_exitcode=1,
+ )
+
+
+@pytest.mark.parametrize(
+ "scenario",
+ [
+ "blocking-str-method",
+ "overridden-getattribute",
+ "str-subclass-argument",
+ "huge-string",
+ ],
+)
+def test_fatal_helper_hard_exit_controls(
+ tmp_path: Path,
+ scenario: str,
+) -> None:
+ _run_bounded_child(
+ _FATAL_SOURCE,
+ tmp_path,
+ scenario,
+ expected_exitcode=1,
+ )
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="process groups are POSIX-only")
+def test_wait_for_ready_rejects_non_string_status_without_comparison(
+ tmp_path: Path,
+) -> None:
+ source = r"""
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+import vllm.v1.ple_offload.worker as worker_module
+
+class Status:
+ def __eq__(self, other):
+ threading.Event().wait()
+
+ def __ne__(self, other):
+ threading.Event().wait()
+
+class Reader:
+ def close(self):
+ return None
+
+sleeper = subprocess.Popen(
+ [sys.executable, "-c", "import time; time.sleep(120)"]
+)
+handle = worker_module.PleOffloadWorkerHandle(
+ proc=sleeper,
+ death_writer=None,
+ ready_pipe_reader=Reader(),
+ launch_deadline=time.monotonic() + 120,
+)
+worker_module._recv_bounded_connection_message = (
+ lambda reader, deadline: {
+ "status": Status(),
+ "layer_names": ["ple"],
+ }
+)
+Path(sys.argv[1]).write_text("armed", encoding="utf-8")
+rejected = False
+closed = False
+reader_cleared = False
+try:
+ try:
+ worker_module.PleOffloadWorker.wait_for_ready(handle)
+ except RuntimeError:
+ rejected = True
+ # Postconditions are read before the emergency cleanup below can act.
+ # Stay inside the parent's 5 second bounded wait so a miss reports here.
+ close_deadline = time.monotonic() + 3
+ while sleeper.poll() is None and time.monotonic() < close_deadline:
+ time.sleep(0.05)
+ closed = sleeper.poll() is not None
+ reader_cleared = handle.ready_pipe_reader is None
+finally:
+ if sleeper.poll() is None:
+ sleeper.kill()
+ sleeper.wait(timeout=120)
+
+if not rejected:
+ raise AssertionError("invalid response status was not rejected")
+if not closed:
+ raise AssertionError("invalid response did not close the child")
+if not reader_cleared:
+ raise AssertionError("invalid response did not clear the ready reader")
+"""
+ _run_bounded_child(source, tmp_path)
+
+
+@pytest.mark.parametrize("identity", ["duplicate", "unbound"])
+def test_profile_breakable_initializer_follows_identity_validation(
+ tmp_path: Path,
+ identity: str,
+) -> None:
+ source = r"""
+import os
+import sys
+import threading
+from pathlib import Path
+from types import SimpleNamespace
+
+from vllm.v1.worker.gpu.cudagraph_utils import ModelCudaGraphManager
+
+identity = sys.argv[2]
+if identity == "duplicate":
+ identities = ("shared", "shared")
+else:
+ identities = ("<unbound>",)
+
+modules = [
+ (
+ "module_" + str(index),
+ SimpleNamespace(
+ _is_cpu_offloaded=True,
+ _host_output_ptr=1,
+ _ple_offload_name=value,
+ ),
+ )
+ for index, value in enumerate(identities)
+]
+model = SimpleNamespace(named_modules=lambda: iter(modules))
+manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager)
+spy = Path(sys.argv[1] + ".initializer")
+
+def blocking_initializer(model):
+ spy.write_text("called", encoding="utf-8")
+ threading.Event().wait()
+
+manager.init_breakable_cg_runner = blocking_initializer
+Path(sys.argv[1]).write_text("armed", encoding="utf-8")
+
+prepare = getattr(manager, "_prepare_ple_capture", None)
+if prepare is None:
+ # This is the pre-fix profiling order: initialize before validation.
+ manager.init_breakable_cg_runner(model)
+else:
+ try:
+ prepare(model, initialize_breakable=True)
+ except RuntimeError as error:
+ if "distinct, bound" not in str(error):
+ raise
+ else:
+ raise AssertionError("invalid identities were accepted")
+
+if spy.exists():
+ raise AssertionError("initializer ran before identity rejection")
+"""
+ _run_bounded_child(source, tmp_path, identity)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="process groups are POSIX-only")
+@pytest.mark.parametrize("status", ["READY", "FAILURE"])
+def test_oversized_startup_response_becomes_bounded_failure(
+ tmp_path: Path,
+ status: str,
+) -> None:
+ source = r"""
+import socket
+import sys
+from multiprocessing.connection import Connection
+from pathlib import Path
+
+import vllm.v1.ple_offload.worker as worker_module
+
+parent_socket, child_socket = socket.socketpair()
+raw_writer = Connection(child_socket.detach())
+writer_type = getattr(
+ worker_module,
+ "_BoundedStartupResponseWriter",
+ None,
+)
+writer = raw_writer if writer_type is None else writer_type(raw_writer)
+
+if sys.argv[2] == "READY":
+ message = {
+ "status": "READY",
+ "layer_names": ["x" * (2 << 20)],
+ }
+else:
+ message = {
+ "status": "FAILURE",
+ "error": "x" * (2 << 20),
+ }
+
+Path(sys.argv[1]).write_text("armed", encoding="utf-8")
+writer.send(message)
+"""
+ _run_bounded_child(source, tmp_path, status)
+
+
+def test_startup_response_payload_at_cap_is_accepted(tmp_path: Path) -> None:
+ source = r"""
+import pickle
+import sys
+import threading
+from pathlib import Path
+
+import vllm.v1.ple_offload.worker as worker_module
+
+writer_type = getattr(
+ worker_module,
+ "_BoundedStartupResponseWriter",
+ None,
+)
+Path(sys.argv[1]).write_text("armed", encoding="utf-8")
+if writer_type is None:
+ threading.Event().wait()
+
+cap = worker_module._PLE_READY_MAX_PAYLOAD_BYTES
+name = "x" * cap
+for unused in range(8):
+ message = {"status": "READY", "layer_names": [name]}
+ payload = pickle.dumps(message, protocol=pickle.HIGHEST_PROTOCOL)
+ if len(payload) == cap:
+ break
+ next_size = len(name) + cap - len(payload)
+ if next_size < 0:
+ raise AssertionError("cannot construct exact-cap payload")
+ name = "x" * next_size
+else:
+ raise AssertionError("exact-cap payload did not converge")
+
+class Writer:
+ def __init__(self):
+ self.payload = None
+
+ def send_bytes(self, payload):
+ self.payload = payload
+
+ def close(self):
+ return None
+
+raw_writer = Writer()
+writer_type(raw_writer).send(message)
+if raw_writer.payload != payload or len(raw_writer.payload) != cap:
+ raise AssertionError("exact-cap startup response was not preserved")
+"""
+ _run_bounded_child(source, tmp_path)
diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py
new file mode 100644
index 0000000..1b67578
--- /dev/null
+++ b/vllm/model_executor/layers/ple_offload_layer.py
@@ -0,0 +1,433 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Base class for PLE layers that support cross-process CPU offload.
+
+Synchronization between the offload process and the GPU main stream is
+handled by :class:`CpuGpuSemaphore`, which uses CUDA's
+``cuStreamWaitValue32`` / ``cuStreamWriteValue32`` stream memory operations.
+
+Decode-step timeline, host-mapped PULL design
+--------------------------------------------------------
+Offload process (CPU, no CUDA) GPU main stream (forward)
+------------------------------ ------------------------------
+forward_impl() on CPU ... other GPU ops ...
+wait host flag == 0 WaitValue32(host flag == 1) <- in graph
+rows -> shared pinned host slot cuMemcpyHtoDAsync(slot -> GPU) <- in graph
+publish fence; host flag = 1 return _gpu_output_buffer[:n]
+ ... model consumes output ...
+ WriteValue32(host flag = 0) <- release_outputs
+
+``release_outputs`` performs the sole reset. Safety requires serialized local
+forwards on one ordered model stream: a new transaction cannot begin until the
+prior release has been enqueued. If any prepared forward is abandoned or its
+correspondence becomes uncertain, the worker must terminate; resetting,
+retrying, or continuing could reuse stale rows or overwrite live input.
+"""
+
+import contextlib
+import functools
+import threading
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterator
+from typing import Any, cast
+
+import torch
+from cuda.bindings import driver as cuda_driver
+from cuda.bindings.driver import CUstreamWaitValue_flags
+from torch import nn
+
+import vllm.envs as envs
+from vllm.utils.torch_utils import direct_register_custom_op
+
+# Module-level flag set to True inside the offload subprocess.
+# Because the offload process and GPU worker processes are separate OS
+# processes (spawned via multiprocessing), each has its own memory space.
+# A plain module-level bool is sufficient -- no thread-local storage needed.
+_offload_worker_flag = False
+
+# Native FULL capture runs on one Python thread. The custom-op implementation
+# appends the exact arguments from every wait-copy invocation while that thread
+# is inside the graph manager's capture scope.
+_ple_copy_capture = threading.local()
+
+
+@contextlib.contextmanager
+def observe_ple_copy_geometry() -> Iterator[list[tuple[str, int, int]]]:
+ """Collect ``(layer, rows, bytes)`` from actual wait-copy executions."""
+ if getattr(_ple_copy_capture, "records", None) is not None:
+ raise RuntimeError("nested PLE copy-geometry observation is unsupported")
+ records: list[tuple[str, int, int]] = []
+ _ple_copy_capture.records = records
+ try:
+ yield records
+ finally:
+ del _ple_copy_capture.records
+
+
+def _record_ple_copy_geometry(
+ layer_name: str,
+ num_tokens: int,
+ row_bytes: int,
+) -> None:
+ records = getattr(_ple_copy_capture, "records", None)
+ if records is not None:
+ rows = int(num_tokens)
+ records.append((str(layer_name), rows, rows * int(row_bytes)))
+
+
+def is_offload_process() -> bool:
+ """Return True inside the dedicated PLE CPU-offload subprocess."""
+ return _offload_worker_flag
+
+
+def mark_as_offload_worker() -> None:
+ """Mark the current process as the dedicated PLE offload worker."""
+ global _offload_worker_flag
+ _offload_worker_flag = True
+
+
+def _cuda_check(result: Any, operation: str) -> Any:
+ """Check the ``(CUresult, ...)`` tuple returned by cuda-python calls."""
+ error = result[0] if isinstance(result, tuple) else result
+ if error.value != 0:
+ raise RuntimeError(f"{operation} failed: {error}")
+ return result
+
+
+class CpuGpuSemaphore:
+ """Cross-process semaphore backed by a one-element int32 CUDA tensor.
+
+ The flag is stored in a regular GPU tensor that is shared through
+ PyTorch's CUDA IPC mechanism. Both processes obtain device pointers that
+ map to the same physical GPU memory, so a stream-memory wait in the GPU
+ worker observes the write issued by the offload process.
+ """
+
+ RESET_VALUE = 0
+ DONE_VALUE = 1
+
+ def __init__(self, device: torch.device) -> None:
+ self._flag_tensor = torch.zeros(1, dtype=torch.int32, device=device)
+ self._flag_ptr: int | None = None
+
+ @classmethod
+ def from_ipc_tensor(cls, flag_tensor: torch.Tensor) -> "CpuGpuSemaphore":
+ """Construct a semaphore from a CUDA tensor received through IPC."""
+ semaphore = cls.__new__(cls)
+ semaphore._flag_tensor = flag_tensor
+ semaphore._flag_ptr = None
+ return semaphore
+
+ @classmethod
+ def from_host_mapped(
+ cls, host_flag_tensor: torch.Tensor, device_ptr: int
+ ) -> "CpuGpuSemaphore":
+ """Semaphore on a page-locked, device-mapped shared host int32.
+
+ ``host_flag_tensor`` keeps the shared storage alive; ``device_ptr`` is
+ its device address (``cuMemHostGetDevicePointer``). Stream memory ops
+ on it are executed by the GPU against host memory, so a CPU process
+ with no CUDA context can signal it with a plain store.
+ """
+ semaphore = cls.__new__(cls)
+ semaphore._flag_tensor = host_flag_tensor
+ semaphore._flag_ptr = int(device_ptr)
+ return semaphore
+
+ @property
+ def flag_tensor(self) -> torch.Tensor:
+ """Return the tensor that backs the semaphore flag."""
+ return self._flag_tensor
+
+ @property
+ def flag_ptr(self) -> int:
+ """Device address used by the stream memory operations."""
+ if self._flag_ptr is not None:
+ return self._flag_ptr
+ return self._flag_tensor.data_ptr()
+
+ def reset(self, stream: torch.cuda.Stream | None = None) -> None:
+ """Enqueue ``WriteValue32(flag=0)`` on ``stream``."""
+ if stream is None:
+ stream = torch.cuda.current_stream()
+ _cuda_check(
+ cuda_driver.cuStreamWriteValue32(
+ cuda_driver.CUstream(stream.cuda_stream),
+ cuda_driver.CUdeviceptr(self.flag_ptr),
+ self.RESET_VALUE,
+ 0,
+ ),
+ "CpuGpuSemaphore.reset",
+ )
+
+ def signal(self, stream: torch.cuda.Stream | None = None) -> None:
+ """Enqueue ``WriteValue32(flag=1)`` on ``stream``."""
+ if stream is None:
+ stream = torch.cuda.current_stream()
+ _cuda_check(
+ cuda_driver.cuStreamWriteValue32(
+ cuda_driver.CUstream(stream.cuda_stream),
+ cuda_driver.CUdeviceptr(self.flag_ptr),
+ self.DONE_VALUE,
+ 0,
+ ),
+ "CpuGpuSemaphore.signal",
+ )
+
+ def wait_reset(self, stream: torch.cuda.Stream | None = None) -> None:
+ """Enqueue ``WaitValue32(flag==0)`` on ``stream``."""
+ if stream is None:
+ stream = torch.cuda.current_stream()
+ _cuda_check(
+ cuda_driver.cuStreamWaitValue32(
+ cuda_driver.CUstream(stream.cuda_stream),
+ cuda_driver.CUdeviceptr(self.flag_ptr),
+ self.RESET_VALUE,
+ CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value,
+ ),
+ "CpuGpuSemaphore.wait_reset",
+ )
+
+
+# ---------------------------------------------------------------------------
+# Custom op: vllm::ple_offload_wait
+# ---------------------------------------------------------------------------
+# Wraps cuStreamWaitValue32 as a torch custom op so that
+# torch.compile(fullgraph=True) treats the CUDA Driver call as an opaque node.
+# The hidden_states argument creates a dynamo data-dependency edge that prevents
+# the wait from being reordered before preceding GPU work.
+# ---------------------------------------------------------------------------
+
+
+def _ple_offload_wait_impl(
+ sem_flag_tensor: torch.Tensor,
+ gpu_output_buffer: torch.Tensor,
+ hidden_states: torch.Tensor,
+) -> None:
+ """Wait for the CPU result without releasing its output buffer."""
+ stream = torch.cuda.current_stream()
+ cuda_stream = cuda_driver.CUstream(stream.cuda_stream)
+ dev_ptr = cuda_driver.CUdeviceptr(sem_flag_tensor.data_ptr())
+ _cuda_check(
+ cuda_driver.cuStreamWaitValue32(
+ cuda_stream,
+ dev_ptr,
+ CpuGpuSemaphore.DONE_VALUE,
+ CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value,
+ ),
+ "cuStreamWaitValue32(done)",
+ )
+
+
+def _ple_offload_wait_fake(
+ sem_flag_tensor: torch.Tensor,
+ gpu_output_buffer: torch.Tensor,
+ hidden_states: torch.Tensor,
+) -> None:
+ """Represent the side-effect-only wait during dynamo tracing."""
+ pass
+
+
+direct_register_custom_op(
+ op_name="ple_offload_wait",
+ op_func=_ple_offload_wait_impl,
+ mutates_args=["gpu_output_buffer"],
+ fake_impl=_ple_offload_wait_fake,
+)
+
+
+# ---------------------------------------------------------------------------
+# Custom op: vllm::ple_offload_wait_copy (host-mapped pull design)
+# ---------------------------------------------------------------------------
+# The offload process writes the result into a page-locked shared host buffer
+# and stores flag=1 in a page-locked shared host int32; it holds no CUDA
+# context. On the CURRENT stream (captured into the graph like the wait above):
+# cuStreamWaitValue32(host flag == 1) -> GPU polls host memory
+# cuMemcpyAsync(gpu buffer <- host buffer, num_tokens * row_bytes)
+# The copy engine reads host memory coherently after the flag was observed,
+# and the consumers read device memory as before. The flag is NOT reset here:
+# ``release_outputs`` resets it after the forward (unchanged semantics), so
+# capture warm-ups and dummy runs behave exactly as with ``ple_offload_wait``.
+# Pointers are passed as ints so no CPU tensor enters the compiled graph.
+# ---------------------------------------------------------------------------
+
+
+def _ple_offload_wait_copy_impl(
+ gpu_output_buffer: torch.Tensor,
+ hidden_states: torch.Tensor,
+ flag_ptr: int,
+ host_ptr: int,
+ row_bytes: int,
+ num_tokens: int,
+ layer_name: str,
+) -> None:
+ """Wait for the host-side result, then DMA it into the GPU buffer."""
+ if not 0 <= int(num_tokens) <= gpu_output_buffer.shape[0] or int(row_bytes) != (
+ gpu_output_buffer.shape[1] * gpu_output_buffer.element_size()
+ ):
+ raise RuntimeError(
+ f"PLE wait_copy out of bounds: {num_tokens} rows x {row_bytes} B vs "
+ f"GPU buffer {tuple(gpu_output_buffer.shape)} {gpu_output_buffer.dtype}"
+ )
+ nbytes = int(num_tokens) * int(row_bytes)
+ _record_ple_copy_geometry(layer_name, num_tokens, row_bytes)
+ stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream)
+ _cuda_check(
+ cuda_driver.cuStreamWaitValue32(
+ stream,
+ cuda_driver.CUdeviceptr(flag_ptr),
+ CpuGpuSemaphore.DONE_VALUE,
+ CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value,
+ ),
+ "cuStreamWaitValue32(host flag done)",
+ )
+ if nbytes > 0:
+ # host_ptr is this worker's host address of page-locked (registered)
+ # memory, so the copy is a true async DMA and is capturable.
+ _cuda_check(
+ cuda_driver.cuMemcpyHtoDAsync(
+ cuda_driver.CUdeviceptr(gpu_output_buffer.data_ptr()),
+ host_ptr,
+ nbytes,
+ stream,
+ ),
+ "cuMemcpyHtoDAsync(PLE host -> GPU)",
+ )
+
+
+def _ple_offload_wait_copy_fake(
+ gpu_output_buffer: torch.Tensor,
+ hidden_states: torch.Tensor,
+ flag_ptr: int,
+ host_ptr: int,
+ row_bytes: int,
+ num_tokens: int,
+ layer_name: str,
+) -> None:
+ """Represent the side-effect-only wait + copy during dynamo tracing."""
+ pass
+
+
+direct_register_custom_op(
+ op_name="ple_offload_wait_copy",
+ op_func=_ple_offload_wait_copy_impl,
+ mutates_args=["gpu_output_buffer"],
+ fake_impl=_ple_offload_wait_copy_fake,
+)
+
+
+class PleOffloadLayer(nn.Module, ABC):
+ """Base class for embedding-like PLE layers that can run on CPU.
+
+ Subclasses implement :meth:`forward_impl`, which contains the regular
+ on-device computation. In a GPU worker with PLE offload enabled, subclass
+ constructors are skipped so their large weights are never allocated. The
+ CPU process constructs a complete copy and owns those weights instead.
+ """
+
+ _is_cpu_offloaded = False
+ _gpu_output_buffer: torch.Tensor
+ _sem: CpuGpuSemaphore
+ # Host-mapped pull design: host address of the shared output buffer and
+ # the bytes per output row. ``None`` keeps the IPC push behaviour.
+ _host_output_ptr: int | None = None
+ _host_row_bytes: int = 0
+ _ple_offload_name: str = "<unbound>"
+
+ def __init_subclass__(cls, **kwargs: object) -> None:
+ """Skip subclass initialization in cross-process GPU-worker mode."""
+ super().__init_subclass__(**kwargs)
+ original_init_obj = cls.__dict__.get("__init__")
+ if original_init_obj is None or not callable(original_init_obj):
+ return
+ original_init = cast(Callable[..., None], original_init_obj)
+
+ @functools.wraps(original_init)
+ def guarded_init(
+ self: "PleOffloadLayer", *args: object, **kwargs: object
+ ) -> None:
+ if envs.VLLM_PLE_CPU_OFFLOAD and not is_offload_process():
+ nn.Module.__init__(self)
+ return
+ original_init(self, *args, **kwargs)
+
+ cls.__init__ = guarded_init # type: ignore[method-assign, assignment]
+
+ @classmethod
+ def get_target_device(cls) -> torch.device:
+ """Return CPU for the offload process and the active GPU otherwise."""
+ if envs.VLLM_PLE_CPU_OFFLOAD:
+ return torch.device("cpu")
+ return torch.device("cuda", torch.accelerator.current_device_index())
+
+ @abstractmethod
+ def forward_impl(
+ self,
+ hidden_states: torch.Tensor,
+ input_ids: torch.Tensor,
+ *args: object,
+ **kwargs: object,
+ ) -> torch.Tensor:
+ """Execute the actual embedding computation on the owning device."""
+ raise NotImplementedError
+
+ def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype:
+ """Return the dtype used by cross-process output buffers."""
+ return default_dtype
+
+ def get_offload_output_dim(self, default_dim: int) -> int:
+ """Return the last dimension used by cross-process output buffers."""
+ return default_dim
+
+ def initialize_dummy_offload_metadata(self, device: torch.device) -> None:
+ """Initialize metadata needed by dummy offload execution."""
+
+ def setup_cross_process_offload(
+ self,
+ gpu_output_buffer: torch.Tensor,
+ semaphore: CpuGpuSemaphore,
+ host_output_ptr: int | None = None,
+ host_row_bytes: int = 0,
+ ) -> None:
+ """Configure the GPU-worker placeholder with its offload resources."""
+ self._is_cpu_offloaded = True
+ self._gpu_output_buffer = gpu_output_buffer
+ self._sem = semaphore
+ self._host_output_ptr = host_output_ptr
+ self._host_row_bytes = int(host_row_bytes)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_ids: torch.Tensor,
+ *args: object,
+ **kwargs: object,
+ ) -> torch.Tensor:
+ """Wait for an offloaded result or delegate to ``forward_impl``."""
+ if self._is_cpu_offloaded:
+ if self._host_output_ptr is not None:
+ torch.ops.vllm.ple_offload_wait_copy(
+ self._gpu_output_buffer,
+ hidden_states,
+ self._sem.flag_ptr,
+ self._host_output_ptr,
+ self._host_row_bytes,
+ input_ids.shape[0],
+ self._ple_offload_name,
+ )
+ else:
+ torch.ops.vllm.ple_offload_wait(
+ self._sem.flag_tensor,
+ self._gpu_output_buffer,
+ hidden_states,
+ )
+ return self._gpu_output_buffer[: input_ids.shape[0]]
+ return self.forward_impl(hidden_states, input_ids, *args, **kwargs)
+
+ def release_offloaded_output(
+ self,
+ stream: torch.cuda.Stream | None = None,
+ ) -> None:
+ """Mark the cross-process output buffer reusable on ``stream``."""
+ if self._is_cpu_offloaded:
+ self._sem.reset(stream)
diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py
index 1ccd23a..e27a18b 100644
--- a/vllm/v1/executor/multiproc_executor.py
+++ b/vllm/v1/executor/multiproc_executor.py
@@ -352,9 +352,10 @@ class MultiprocExecutor(Executor):
proc = next(h.proc for h in workers if h.proc.sentinel == died[0])
logger.error(
"Worker proc %s died unexpectedly (exit code: %s), "
- "shutting down executor.",
+ "shutting down executor. executor_shutdown_monotonic=%.9f",
proc.name,
proc.exitcode,
+ time.monotonic(),
)
_self.shutdown()
callback = _self.failure_callback
@@ -1034,6 +1035,28 @@ class WorkerProc:
SUCCESS = auto()
FAILURE = auto()
+ def _fail_ple_rpc_if_pending(
+ self, where: str, error: BaseException
+ ) -> None:
+ """Fail closed if RPC recovery would abandon queued PLE GPU work."""
+ owners = (self.worker, getattr(self.worker, "worker", None))
+ for owner in owners:
+ if owner is None:
+ continue
+ runner = getattr(owner, "model_runner", None)
+ connector = getattr(runner, "_ple_offload_connector", None)
+ if connector is None:
+ continue
+ try:
+ pending = connector.has_pending_completions
+ except BaseException as state_error:
+ connector._fail_closed(
+ f"{where}:pending-state", state_error
+ )
+ if pending:
+ connector._fail_closed(where, error)
+ return
+
def enqueue_output(self, output: Any):
"""Prepares output from the worker and enqueues it to the
worker_response_mq. If the output is an Exception, it is
@@ -1044,6 +1067,9 @@ class WorkerProc:
output = output.get_output()
except Exception as e:
logger.exception("Error getting async model runner output")
+ self._fail_ple_rpc_if_pending(
+ "async-output-copy", e
+ )
output = e
if isinstance(output, Exception):
@@ -1104,6 +1130,9 @@ class WorkerProc:
if output_rank is None or self.rank == output_rank:
self.handle_output(output)
except Exception as e:
+ self._fail_ple_rpc_if_pending(
+ f"worker-rpc:{method}", e
+ )
# Notes have been introduced in python 3.11
if hasattr(e, "add_note"):
e.add_note(traceback.format_exc())
diff --git a/vllm/v1/ple_offload/__init__.py b/vllm/v1/ple_offload/__init__.py
new file mode 100644
index 0000000..208f01a
--- /dev/null
+++ b/vllm/v1/ple_offload/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py
new file mode 100644
index 0000000..52bc462
--- /dev/null
+++ b/vllm/v1/ple_offload/connector.py
@@ -0,0 +1,855 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Exchange PLE data between a GPU worker and the CPU-offload process."""
+
+import os
+import queue
+import threading
+from dataclasses import dataclass
+from multiprocessing.reduction import ForkingPickler
+from typing import Any
+
+import msgspec
+import torch
+from vllm.v1.ple_offload import rowdump as _rowdump
+from vllm.v1.ple_offload.fatal import ple_fail_closed
+from vllm.v1.ple_offload.progress import PleProgressMonitor
+import torch.nn as nn
+import zmq
+from cuda.bindings import driver as cuda_driver
+
+from vllm.config import VllmConfig
+from vllm.distributed.parallel_state import get_dp_group, get_pp_group, get_tp_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.ple_offload_layer import (
+ CpuGpuSemaphore,
+ PleOffloadLayer,
+)
+from vllm.v1.ple_offload.protocol import (
+ PleOffloadRegistration,
+ PleOffloadRequest,
+)
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm_ple_test_fault import _FAULT
+
+logger = init_logger(__name__)
+
+# CUDA registration alone does not own Python storage. Captured graphs retain
+# raw addresses, so every successfully registered host allocation remains
+# strongly referenced until process exit.
+_PROCESS_LIFETIME_REGISTERED_STORAGE: list[torch.Tensor] = []
+
+
+def _unlink_filename_shared_storage(
+ tensor: torch.Tensor,
+ reserved_filenames: set[bytes],
+) -> None:
+ """Reserve and unlink one already filename-shared CPU storage."""
+ try:
+ storage = tensor.untyped_storage()
+ if storage.device.type != "cpu" or not storage.is_shared():
+ return
+ try:
+ storage._get_shared_fd()
+ except RuntimeError:
+ # Filename-shared storages do not expose a shared fd. Only now is
+ # it safe to query the filename: doing so on ordinary or fd-shared
+ # storage would convert it to a new filename-backed allocation.
+ pass
+ else:
+ return
+ _, filename, _ = storage._share_filename_cpu_()
+ if not isinstance(filename, bytes):
+ filename = os.fsencode(filename)
+ path = os.path.join("/dev/shm", os.fsdecode(filename).lstrip("/"))
+ except Exception:
+ logger.warning(
+ "Failed to inspect a PLE shared-memory storage during teardown",
+ exc_info=True,
+ )
+ return
+
+ # RefcountedMapAllocator unlinks the name when its shared counter reaches
+ # zero. Reserve one count before manual unlink so later storage destruction
+ # cannot try to unlink the already-removed name. The filename identifies
+ # the shared counter and also deduplicates tensor aliases and repeated close.
+ if filename in reserved_filenames:
+ return
+ try:
+ storage._shared_incref()
+ except Exception:
+ logger.warning(
+ "Failed to reserve the PLE shared-memory counter for %s",
+ path,
+ exc_info=True,
+ )
+ return
+ reserved_filenames.add(filename)
+
+ try:
+ os.unlink(path)
+ except FileNotFoundError:
+ # Another teardown participant already removed the name. Retain the
+ # reservation because this allocator must still avoid a later unlink.
+ pass
+ except OSError as error:
+ # The name remains available, so restore the allocator's original
+ # counter and permit a later release attempt to retry.
+ try:
+ storage._shared_decref()
+ except Exception:
+ # The reservation state is uncertain after a failed private call.
+ # Keep it marked to avoid adding another reservation on retry.
+ logger.warning(
+ "Failed to roll back the PLE shared-memory counter for %s",
+ path,
+ exc_info=True,
+ )
+ else:
+ reserved_filenames.discard(filename)
+ logger.warning("Failed to unlink PLE shared memory %s: %s", path, error)
+
+
+@dataclass(frozen=True)
+class _PendingPleOffloadRequest:
+ """Bind request metadata to its MRV2 D2H completion event."""
+
+ request: PleOffloadRequest
+ d2h_done_event: torch.cuda.Event
+ eligible: bool = False
+ transaction_id: str | None = None
+ correlation_id: str | None = None
+
+
+def _cuda_check(result: Any, operation: str) -> Any:
+ """Check the ``(CUresult, ...)`` tuple returned by cuda-python calls."""
+ error = result[0] if isinstance(result, tuple) else result
+ if error.value != 0:
+ raise RuntimeError(f"{operation} failed: {error}")
+ return result
+
+
+class PleOffloadConnector:
+ """Connect a model runner V2 GPU worker to the PLE CPU worker."""
+
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ model: nn.Module,
+ device: torch.device,
+ ipc_addr: str,
+ *,
+ input_ids_source: torch.Tensor,
+ query_start_loc_source: torch.Tensor,
+ ngram_context_source: torch.Tensor | None,
+ ) -> None:
+ self.device = device
+ self.rank = int(vllm_config.parallel_config.rank)
+ self.dp_rank = get_dp_group().rank_in_group
+ self.tp_rank = get_tp_group().rank_in_group
+ self.pp_rank = get_pp_group().rank_in_group
+ # Host-mapped pull design: per layer a shared host output buffer and a
+ # shared host int32 flag (element 0), filled/stored by the offload
+ # process with plain CPU writes. Page-locked + device-mapped here only
+ # after registration, because serialising for registration may move
+ # the shared storage to a new address.
+ self._host_output_bufs: dict[str, torch.Tensor] = {}
+ self._host_flags: dict[str, torch.Tensor] = {}
+ self._reserved_host_storage_filenames: set[bytes] = set()
+ self._gpu_output_bufs: dict[str, torch.Tensor] = {}
+ self._layers = self._setup_layers(vllm_config, model)
+ self._progress = PleProgressMonitor.from_environment(
+ rank=self.rank,
+ layer_names=tuple(self._layers),
+ )
+
+ # TP0 registers shared buffers with CUDA for asynchronous D2H copies.
+ scheduler_config = vllm_config.scheduler_config
+ self._input_ids_buf = torch.empty(
+ scheduler_config.max_num_batched_tokens,
+ dtype=torch.int32,
+ device="cpu",
+ ).share_memory_()
+ self._query_start_loc_buf = torch.empty(
+ scheduler_config.max_num_seqs + 1,
+ dtype=torch.int32,
+ device="cpu",
+ ).share_memory_()
+ self._ngram_context_buf = None
+ config = vllm_config.model_config.hf_text_config
+ ngram_context_len = int(config.ngram_size) - 1
+ if ngram_context_len > 0:
+ self._ngram_context_buf = torch.empty(
+ scheduler_config.max_num_seqs,
+ ngram_context_len,
+ dtype=torch.int32,
+ device="cpu",
+ ).share_memory_()
+
+ # Runner input allocations are address-stable, so bind them once and
+ # pass only batch sizes through the per-forward request queue.
+ self._input_ids_source = input_ids_source
+ self._query_start_loc_source = query_start_loc_source
+ self._ngram_context_source = ngram_context_source
+ self._validate_input_sources()
+
+ self._pinned_input_buffers: list[torch.Tensor] = []
+ # Device-mapped host outputs/flags referenced by raw address from
+ # captured graphs: never unregistered while the process lives
+ # (unregistering under a pending graph wait unmaps it).
+ self._mapped_host_buffers: list[torch.Tensor] = []
+ # True while dummy/capture work (seeds, dummy forwards, capture
+ # cleanup resets) has run since the last real forward. Releases in
+ # that state complete synchronously, and the next real forward first
+ # passes a TP barrier, so no rank can publish a real request while
+ # another rank still has a dummy/cleanup reset pending.
+ # In this tree real-path kernel warm-ups run BEFORE
+ # capture and capture adds seeds/resets after them, so this must hold
+ # at every dummy->real transition, not only at the first real one.
+ self._after_dummy = True
+ self._request_queue: queue.Queue[_PendingPleOffloadRequest | None] = (
+ queue.Queue(maxsize=vllm_config.max_concurrent_batches)
+ )
+ self._dropped_events: list[torch.cuda.Event] = []
+ self._request_thread: threading.Thread | None = None
+ self._request_thread_ready = threading.Event()
+ self._zmq_ctx: zmq.Context | None = None
+ self._registration_socket: zmq.Socket | None = None
+ self._d2h_event_pool: queue.Queue[torch.cuda.Event] | None = None
+
+ try:
+ self._zmq_ctx = zmq.Context()
+ self._registration_socket = self._zmq_ctx.socket(zmq.PUSH)
+ self._registration_socket.connect(ipc_addr)
+ payload = self._serialize_registration(vllm_config)
+
+ # Serialising moved the shared storages to their final addresses.
+ # Page-lock/map exactly those, THEN publish, so a received
+ # registration always means "storage usable".
+ with torch.accelerator.device_index(self.device.index):
+ self._map_host_outputs(vllm_config)
+ if self.tp_rank == 0:
+ self._pin_input_buffers()
+ self._send_registration(payload, ipc_addr)
+
+ if self.tp_rank == 0:
+ with torch.accelerator.device_index(self.device.index):
+ self._d2h_event_pool = queue.Queue(
+ maxsize=vllm_config.max_concurrent_batches
+ )
+ for _ in range(vllm_config.max_concurrent_batches):
+ self._d2h_event_pool.put_nowait(torch.cuda.Event())
+ self._start_request_thread(ipc_addr)
+ except Exception:
+ self.close()
+ raise
+
+ def _setup_layers(
+ self,
+ vllm_config: VllmConfig,
+ model: nn.Module,
+ ) -> dict[str, PleOffloadLayer]:
+ """Attach output buffers and semaphores to GPU PLE placeholders."""
+ layers = {
+ name: module
+ for name, module in model.named_modules()
+ if isinstance(module, PleOffloadLayer)
+ }
+ if not layers:
+ raise RuntimeError(
+ "VLLM_PLE_CPU_OFFLOAD is enabled, but the model has no PleOffloadLayer"
+ )
+
+ config = vllm_config.model_config.hf_text_config
+ max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+ for name, layer in layers.items():
+ if vllm_config.load_config.load_format == "dummy":
+ layer.initialize_dummy_offload_metadata(self.device)
+ out_dim = layer.get_offload_output_dim(int(config.ple_embed_dim))
+ out_dtype = layer.get_offload_output_dtype(vllm_config.model_config.dtype)
+ # The captured forward DMAs the host result into this device
+ # buffer after the host flag is set; consumers read it as before.
+ self._gpu_output_bufs[name] = torch.empty(
+ max_num_tokens, out_dim, dtype=out_dtype, device=self.device
+ )
+ # Zero-initialised so dummy/profile forwards copy finite values.
+ self._host_output_bufs[name] = torch.zeros(
+ max_num_tokens, out_dim, dtype=out_dtype, device="cpu"
+ ).share_memory_()
+ # 16 x int32 = one 64-byte line; element 0 is the flag.
+ self._host_flags[name] = torch.zeros(
+ 16, dtype=torch.int32, device="cpu"
+ ).share_memory_()
+ return layers
+
+ def _map_host_outputs(self, vllm_config: VllmConfig) -> None:
+ """Page-lock + device-map host outputs/flags and arm the placeholders."""
+ for name, layer in self._layers.items():
+ layer._ple_offload_name = name
+ host_out = self._host_output_bufs[name]
+ host_flag = self._host_flags[name]
+ for buffer in (host_out, host_flag):
+ self._register_host_buffer(buffer, device_map=True)
+ flag_dev_ptr = _cuda_check(
+ cuda_driver.cuMemHostGetDevicePointer(host_flag.data_ptr(), 0),
+ "cuMemHostGetDevicePointer(PLE host flag)",
+ )[1]
+ out_dev_ptr = _cuda_check(
+ cuda_driver.cuMemHostGetDevicePointer(host_out.data_ptr(), 0),
+ "cuMemHostGetDevicePointer(PLE host output)",
+ )[1]
+ row_bytes = host_out.shape[1] * host_out.element_size()
+ gpu_buf = self._gpu_output_bufs[name]
+ if (
+ gpu_buf.shape[0] != host_out.shape[0]
+ or gpu_buf.shape[1] * gpu_buf.element_size() != row_bytes
+ ):
+ raise RuntimeError(f"PLE host/GPU buffer geometry mismatch for {name}")
+ # Stream memops use the flag's DEVICE pointer; the H2D copy uses
+ # this worker's HOST pointer (no raw addresses cross
+ # processes; the offload writes through its own mapping).
+ layer.setup_cross_process_offload(
+ gpu_buf,
+ CpuGpuSemaphore.from_host_mapped(host_flag, int(flag_dev_ptr)),
+ host_output_ptr=host_out.data_ptr(),
+ host_row_bytes=row_bytes,
+ )
+ logger.info(
+ "PleOffload: host-mapped pull for %s (rows %d x %d B, flag dev "
+ "0x%x, out dev 0x%x, host==dev %s)",
+ name,
+ host_out.shape[0],
+ row_bytes,
+ int(flag_dev_ptr),
+ int(out_dev_ptr),
+ int(out_dev_ptr) == host_out.data_ptr(),
+ )
+
+ def _pin_input_buffers(self) -> None:
+ """Page-lock shared input allocations without replacing their storage."""
+ buffers = [self._input_ids_buf, self._query_start_loc_buf]
+ if self._ngram_context_buf is not None:
+ buffers.append(self._ngram_context_buf)
+ for buffer in buffers:
+ self._register_host_buffer(buffer, device_map=False)
+
+ def _register_host_buffer(self, buffer: torch.Tensor, device_map: bool) -> None:
+ """Page-lock (and optionally device-map) a shared CPU allocation."""
+ if buffer.device.type != "cpu" or not buffer.is_shared():
+ raise RuntimeError("PLE host buffers must be shared CPU tensors")
+ if not buffer.is_contiguous():
+ raise RuntimeError("PLE host buffers must be contiguous")
+ flags = cuda_driver.CU_MEMHOSTREGISTER_PORTABLE
+ if device_map:
+ flags |= cuda_driver.CU_MEMHOSTREGISTER_DEVICEMAP
+ _cuda_check(
+ cuda_driver.cuMemHostRegister(
+ buffer.data_ptr(),
+ buffer.numel() * buffer.element_size(),
+ flags,
+ ),
+ "cuMemHostRegister(PLE host buffer)",
+ )
+ _PROCESS_LIFETIME_REGISTERED_STORAGE.append(buffer)
+ if device_map:
+ self._mapped_host_buffers.append(buffer)
+ else:
+ self._pinned_input_buffers.append(buffer)
+ if not buffer.is_pinned():
+ raise RuntimeError("CUDA did not page-lock a PLE host buffer")
+
+ def _unpin_input_buffers(self) -> None:
+ """Release CUDA registrations after the request thread has stopped."""
+ for buffer in reversed(self._pinned_input_buffers):
+ try:
+ _cuda_check(
+ cuda_driver.cuMemHostUnregister(buffer.data_ptr()),
+ "cuMemHostUnregister(PLE input buffer)",
+ )
+ except RuntimeError:
+ logger.exception("Failed to unregister a PLE input buffer")
+ self._pinned_input_buffers.clear()
+
+ def _serialize_registration(self, vllm_config: VllmConfig) -> Any:
+ """Build and serialise the registration (may move shared storages)."""
+ # Each GPU worker owns distinct host output buffers and flags, while
+ # TP0's shared inputs become the request source for its DP rank.
+ registration = PleOffloadRegistration(
+ worker_id=(
+ self.dp_rank * vllm_config.parallel_config.world_size
+ + vllm_config.parallel_config.rank
+ ),
+ tp_rank=self.tp_rank,
+ dp_rank=self.dp_rank,
+ host_output_buffers=dict(self._host_output_bufs),
+ host_flag_tensors=dict(self._host_flags),
+ input_ids_buf=self._input_ids_buf,
+ query_start_loc_buf=self._query_start_loc_buf,
+ ngram_context_buf=self._ngram_context_buf,
+ )
+
+ # ForkingPickler transmits tensors through shared-memory and CUDA IPC.
+ import torch.multiprocessing as torch_mp
+
+ original_strategy = torch_mp.get_sharing_strategy()
+ torch_mp.set_sharing_strategy("file_system")
+ try:
+ payload = ForkingPickler.dumps(registration)
+ finally:
+ torch_mp.set_sharing_strategy(original_strategy)
+ return payload
+
+ def _send_registration(self, payload: Any, ipc_addr: str) -> None:
+ """Publish the serialised registration (storages already mapped)."""
+ assert self._registration_socket is not None
+ self._registration_socket.send(payload)
+
+ logger.info(
+ "PleOffload: registered %d PleOffloadLayer(s) "
+ "(dp_rank=%d, tp_rank=%d, ipc_addr=%s): %s",
+ len(self._layers),
+ self.dp_rank,
+ self.tp_rank,
+ ipc_addr,
+ sorted(self._layers),
+ )
+
+ def _start_request_thread(self, ipc_addr: str) -> None:
+ """Start the thread that publishes batches after inputs are ready."""
+ self._request_thread = threading.Thread(
+ target=self._request_loop,
+ args=(ipc_addr,),
+ name=f"ple-offload-dp{self.dp_rank}",
+ daemon=True,
+ )
+ self._request_thread.start()
+ if not self._request_thread_ready.wait(timeout=10):
+ raise RuntimeError("Timed out starting the PLE request thread")
+
+ def _request_loop(self, ipc_addr: str) -> None:
+ """Wait for staged inputs, then notify the CPU worker."""
+ socket: zmq.Socket | None = None
+ try:
+ if self._zmq_ctx is None:
+ raise RuntimeError("PLE ZMQ context closed before thread startup")
+ socket = self._zmq_ctx.socket(zmq.PUSH)
+ socket.connect(ipc_addr)
+ self._request_thread_ready.set()
+ while True:
+ request = self._request_queue.get()
+ if request is None:
+ return
+ self._process_request(request, socket)
+ except Exception:
+ logger.exception("PLE request thread failed")
+ os._exit(1)
+ finally:
+ self._request_thread_ready.set()
+ if socket is not None:
+ socket.close(linger=0)
+
+ def _process_request(
+ self, pending: _PendingPleOffloadRequest, socket: zmq.Socket
+ ) -> None:
+ """Wait for one staged batch and publish its request."""
+ request = pending.request
+ event_pool = self._d2h_event_pool
+ event = pending.d2h_done_event
+ assert event_pool is not None, "PLE D2H event pool is not initialized"
+ if _FAULT is not None and _FAULT.fires(
+ "block_input_completion",
+ role="gpu",
+ unit="connector-request",
+ eligible=pending.eligible,
+ tp_rank=self.tp_rank,
+ pp_rank=self.pp_rank,
+ dp_rank=self.dp_rank,
+ transaction_id=pending.transaction_id,
+ correlation_id=pending.correlation_id,
+ ):
+ threading.Event().wait()
+
+ with torch.accelerator.device_index(self.device.index):
+ event.synchronize()
+
+ payload = msgspec.msgpack.encode(request)
+ if _FAULT is None:
+ socket.send(payload)
+ else:
+ # Armed runs only: correlation metadata rides as a second frame.
+ metadata = msgspec.msgpack.encode(
+ {
+ "serving": pending.eligible,
+ "transaction": pending.transaction_id,
+ "correlation": pending.correlation_id,
+ }
+ )
+ socket.send_multipart([payload, metadata])
+ event_pool.put_nowait(event)
+
+ def _validate_input_sources(self) -> None:
+ """Validate fixed runner sources against shared input buffers."""
+ sources = [
+ ("input_ids", self._input_ids_source, self._input_ids_buf),
+ (
+ "query_start_loc",
+ self._query_start_loc_source,
+ self._query_start_loc_buf,
+ ),
+ ]
+ if (self._ngram_context_source is None) != (self._ngram_context_buf is None):
+ raise ValueError("PLE ngram_context source and buffer must match")
+ if self._ngram_context_source is not None:
+ assert self._ngram_context_buf is not None
+ sources.append(
+ (
+ "ngram_context",
+ self._ngram_context_source,
+ self._ngram_context_buf,
+ )
+ )
+
+ for name, source, buffer in sources:
+ if (
+ source.device != self.device
+ or source.dtype != buffer.dtype
+ or source.ndim != buffer.ndim
+ or source.shape[0] < buffer.shape[0]
+ or source.shape[1:] != buffer.shape[1:]
+ ):
+ raise ValueError(f"PLE {name} source is incompatible")
+
+ def _enqueue_cuda_inputs(
+ self,
+ request: PleOffloadRequest,
+ d2h_done_event: torch.cuda.Event,
+ ) -> None:
+ """Stage inputs on the model stream and record D2H completion."""
+ with (
+ record_function_or_nullcontext("ple_offload: enqueue_cuda_inputs"),
+ torch.accelerator.device_index(self.device.index),
+ ):
+ stream = torch.cuda.current_stream(self.device)
+ self._input_ids_buf[: request.num_tokens].copy_(
+ self._input_ids_source[: request.num_tokens],
+ non_blocking=True,
+ )
+ self._query_start_loc_buf[: request.num_reqs + 1].copy_(
+ self._query_start_loc_source[: request.num_reqs + 1],
+ non_blocking=True,
+ )
+ if self._ngram_context_buf is not None:
+ assert self._ngram_context_source is not None
+ self._ngram_context_buf[: request.num_reqs].copy_(
+ self._ngram_context_source[: request.num_reqs],
+ non_blocking=True,
+ )
+ d2h_done_event.record(stream)
+
+ def _launch(
+ self,
+ num_reqs: int,
+ num_tokens: int,
+ *,
+ serving: bool = False,
+ transaction_id: str | None = None,
+ correlation_id: str | None = None,
+ ) -> None:
+ """Stage or queue one batch for request publication."""
+ # Inputs are replicated across TP ranks. One request per DP rank drives
+ # the CPU result fan-out to every registered TP output buffer.
+ if self.tp_rank != 0:
+ return
+
+ request = PleOffloadRequest(
+ dp_rank=self.dp_rank,
+ num_tokens=num_tokens,
+ num_reqs=num_reqs,
+ )
+ assert self._d2h_event_pool is not None, "PLE D2H event pool is not initialized"
+ try:
+ d2h_done_event = self._d2h_event_pool.get_nowait()
+ except queue.Empty as exc:
+ raise RuntimeError(
+ "PLE has more requests than configured concurrent batches"
+ ) from exc
+ self._enqueue_cuda_inputs(request, d2h_done_event)
+ if _FAULT is not None and _FAULT.fires(
+ "drop_request",
+ role="gpu",
+ unit="connector-request",
+ eligible=serving,
+ tp_rank=self.tp_rank,
+ pp_rank=self.pp_rank,
+ dp_rank=self.dp_rank,
+ transaction_id=transaction_id,
+ correlation_id=correlation_id,
+ ):
+ # Deliberately retain the event: the fault is loss after D2H, not
+ # premature event-pool reuse while that D2H may still be in flight.
+ self._dropped_events.append(d2h_done_event)
+ return
+ self._request_queue.put_nowait(
+ _PendingPleOffloadRequest(
+ request,
+ d2h_done_event,
+ eligible=serving,
+ transaction_id=transaction_id,
+ correlation_id=correlation_id,
+ )
+ )
+
+ def _fail_closed(self, where: str, error: BaseException) -> None:
+ ple_fail_closed(
+ where=where,
+ rank=self.rank,
+ layer=",".join(self._layers),
+ sequence=self._progress.sequence,
+ error=error,
+ )
+
+ @property
+ def has_pending_completions(self) -> bool:
+ return self._progress.has_pending_completions
+
+ def begin_forward(self, *, dummy_run: bool, num_tokens: int) -> int:
+ """Begin the bounded transaction before any runner fault boundary."""
+ try:
+ return self._progress.begin(
+ startup=dummy_run,
+ num_tokens=num_tokens,
+ )
+ except Exception as error:
+ self._fail_closed("prepare_forward", error)
+
+ def fault_transaction(self, sequence: int) -> str:
+ """Read diagnostic state only when the caller has established arming."""
+ return self._progress.describe_submission(sequence)
+
+ def completion_is_pending(self, sequence: int) -> bool:
+ """Return exact-sequence incomplete completion state for cut 6."""
+ return self._progress.completion_is_pending(sequence)
+
+ def prepare_forward(
+ self,
+ num_reqs: int,
+ num_tokens: int,
+ dummy_run: bool,
+ *,
+ consumer_num_tokens: int | None = None,
+ serving: bool = False,
+ transaction_id: str | None = None,
+ correlation_id: str | None = None,
+ ) -> None:
+ """Prepare one PLE transaction; any ordinary exception is fatal."""
+ try:
+ self._prepare_forward_impl(
+ num_reqs,
+ num_tokens,
+ dummy_run,
+ consumer_num_tokens=consumer_num_tokens,
+ serving=serving,
+ transaction_id=transaction_id,
+ correlation_id=correlation_id,
+ )
+ except Exception as error:
+ self._fail_closed("prepare_forward", error)
+
+ def _prepare_forward_impl(
+ self,
+ num_reqs: int,
+ num_tokens: int,
+ dummy_run: bool,
+ *,
+ consumer_num_tokens: int | None = None,
+ serving: bool = False,
+ transaction_id: str | None = None,
+ correlation_id: str | None = None,
+ ) -> None:
+ """Submit real inputs or satisfy the PLE wait for a dummy forward."""
+ if consumer_num_tokens is None:
+ consumer_num_tokens = num_tokens
+ if int(num_tokens) != int(consumer_num_tokens):
+ raise RuntimeError(
+ "PLE request/consumer token-count mismatch: "
+ f"request={num_tokens}, consumer={consumer_num_tokens}"
+ )
+ if not 0 <= int(num_tokens) <= self._input_ids_buf.shape[0]:
+ raise RuntimeError(
+ f"PLE token count {num_tokens} exceeds connector capacity "
+ f"{self._input_ids_buf.shape[0]}"
+ )
+ if not 0 <= int(num_reqs) < self._query_start_loc_buf.shape[0]:
+ raise RuntimeError(
+ f"PLE request count {num_reqs} exceeds connector capacity "
+ f"{self._query_start_loc_buf.shape[0] - 1}"
+ )
+ self._rowdump_n = None if dummy_run else num_tokens
+ if dummy_run:
+ self._signal_dummy_outputs_impl(num_tokens)
+ return
+ if self._after_dummy:
+ # Every dummy/cleanup reset on THIS rank already completed
+ # (synchronous releases); the barrier extends that to every
+ # stage-0 TP rank before TP0 can publish (CPU group, no GPU work;
+ # all TP ranks run the same forwards in the same order).
+ get_tp_group().barrier()
+ self._after_dummy = False
+ logger.info(
+ "PleOffload: dummy->real transition (dp_rank=%d, tp_rank=%d): "
+ "releases were synchronous, TP barrier passed",
+ self.dp_rank,
+ self.tp_rank,
+ )
+ self._launch(
+ num_reqs,
+ num_tokens,
+ serving=serving,
+ transaction_id=transaction_id,
+ correlation_id=correlation_id,
+ )
+
+ def signal_dummy_outputs(self, num_tokens: int) -> None:
+ """Locally satisfy PLE waits for dummy and capture forwards."""
+ self._progress.begin(startup=True, num_tokens=num_tokens)
+ self._signal_dummy_outputs_impl(num_tokens)
+
+ def _signal_dummy_outputs_impl(self, num_tokens: int) -> None:
+ """Signal dummy outputs after the progress transaction is armed."""
+ # Dummy and capture forwards do not send CPU requests, but every PLE
+ # placeholder still waits for a completed output semaphore.
+ self._after_dummy = True
+ stream = torch.cuda.current_stream(self.device)
+ for layer in self._layers.values():
+ layer._gpu_output_buffer[:num_tokens].zero_()
+ layer._sem.signal(stream)
+
+ def release_outputs(
+ self,
+ *,
+ serving: bool = False,
+ transaction_id: str | None = None,
+ correlation_id: str | None = None,
+ ) -> None:
+ """Release one PLE transaction; any ordinary exception is fatal."""
+ try:
+ self._release_outputs_impl(
+ serving=serving,
+ transaction_id=transaction_id,
+ correlation_id=correlation_id,
+ )
+ self._progress.complete()
+ except Exception as error:
+ self._fail_closed("release_outputs", error)
+
+ def _release_outputs_impl(
+ self,
+ *,
+ serving: bool = False,
+ transaction_id: str | None = None,
+ correlation_id: str | None = None,
+ ) -> None:
+ """Mark GPU output buffers reusable after the model consumes them."""
+ if _FAULT is not None and _FAULT.fires(
+ "exc_in_release",
+ role="gpu",
+ unit="connector-release",
+ eligible=serving,
+ tp_rank=self.tp_rank,
+ pp_rank=self.pp_rank,
+ dp_rank=self.dp_rank,
+ transaction_id=transaction_id,
+ correlation_id=correlation_id,
+ ):
+ raise RuntimeError("PLE test fault: exc_in_release")
+ n = getattr(self, "_rowdump_n", None)
+ self._rowdump_n = None
+ if _rowdump.enabled() and n is not None and not torch.cuda.is_current_stream_capturing():
+ torch.cuda.current_stream(self.device).synchronize()
+ seq = getattr(self, "_rowdump_seq", 0)
+ self._rowdump_seq = seq + 1
+ for layer in self._layers.values():
+ _rowdump.write(f"gpu-tp{self.tp_rank}", 0, self.tp_rank, seq,
+ self._input_ids_source[:n], layer._gpu_output_buffer[:n])
+ # Reset only after the consumer forward so the CPU worker cannot
+ # overwrite an output that a GPU PLE placeholder may still read.
+ stream = torch.cuda.current_stream(self.device)
+ for layer in self._layers.values():
+ layer.release_offloaded_output(stream)
+ # Retain the serving deadline until the model stream has actually
+ # executed every reset above. The monitor thread is the sole query site.
+ completion_event = torch.cuda.Event()
+ completion_event.record(stream)
+ self._progress.track_completion(completion_event)
+ # After dummy/capture work (profile, warm-up sizes, capture seed,
+ # adaptive-verification dummy replays, capture cleanup): complete the
+ # reset before returning, so it cannot still be pending when a real
+ # request lets the offload publish. Real-forward releases stay async.
+ # Never inside a stream capture.
+ if self._after_dummy and not torch.cuda.is_current_stream_capturing():
+ stream.synchronize()
+
+ def release_host_storage(self) -> None:
+ """Unlink host-storage names while retaining mappings and tensors."""
+ buffers = [
+ *self._host_output_bufs.values(),
+ *self._host_flags.values(),
+ self._input_ids_buf,
+ self._query_start_loc_buf,
+ ]
+ if self._ngram_context_buf is not None:
+ buffers.append(self._ngram_context_buf)
+ for buffer in buffers:
+ _unlink_filename_shared_storage(
+ buffer,
+ self._reserved_host_storage_filenames,
+ )
+
+ def close(self) -> None:
+ """Stop transport and unlink names, retaining host mappings."""
+ request_thread = self._request_thread
+ if request_thread is not None and request_thread.is_alive():
+ try:
+ self._request_queue.put(None, timeout=5)
+ except queue.Full:
+ logger.error("Timed out stopping the PLE request thread")
+ request_thread.join(timeout=5)
+ if request_thread is not None and request_thread.is_alive():
+ # The thread may still access the registered buffers or ZMQ context.
+ # It may also still race another process opening a shared name, so
+ # this deferred-cleanup branch must not unlink host storage.
+ logger.error("PLE request thread did not stop; deferring resource cleanup")
+ self._progress.close()
+ return
+ self._request_thread = None
+
+ # Inputs, outputs, and flags remain registered and strongly referenced.
+ # A captured graph or delayed request can still contain their addresses;
+ # the CUDA driver releases registrations when the process exits.
+ registered = len(self._pinned_input_buffers) + len(self._mapped_host_buffers)
+ if registered:
+ logger.debug(
+ "PleOffload: %d registered host buffers retained until exit",
+ registered,
+ )
+ self._d2h_event_pool = None
+ if self._registration_socket is not None:
+ self._registration_socket.close(linger=0)
+ self._registration_socket = None
+ if self._zmq_ctx is not None:
+ self._zmq_ctx.term()
+ self._zmq_ctx = None
+ self._progress.close()
+
+ # __init__'s failure handler can reach here before the child has rebuilt
+ # a registration. Startup is already failing in that case; if a payload
+ # was sent, a late rebuild may also observe the removed name. Successful
+ # startup accepts every registration before an external close begins.
+ #
+ # NOTE: Graceful shutdown unlinks the PLE shared-memory names and the
+ # PLE ipc socket. The remaining 17 /tmp UUID sockets per boot are
+ # upstream vLLM VLLM_RPC_BASE_PATH behaviour; a tmpfiles/launch-profile
+ # sweep is recommended.
+ # Keep this last: unlinking preserves every existing CPU/CUDA mapping,
+ # but no transport user may still need to open a name at this point.
+ self.release_host_storage()
diff --git a/vllm/v1/ple_offload/fatal.py b/vllm/v1/ple_offload/fatal.py
new file mode 100644
index 0000000..a80cad4
--- /dev/null
+++ b/vllm/v1/ple_offload/fatal.py
@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fail-closed helpers for the PLE single-slot transport."""
+
+import functools
+import os
+import time
+import traceback
+from collections.abc import Callable
+from typing import Any, NoReturn, TypeVar
+
+_F = TypeVar("_F", bound=Callable[..., Any])
+
+_TRUSTED_ERROR_LABELS = (
+ (RuntimeError, "RuntimeError"),
+ (TimeoutError, "TimeoutError"),
+ (ValueError, "ValueError"),
+ (OSError, "OSError"),
+ (EOFError, "EOFError"),
+ (BrokenPipeError, "BrokenPipeError"),
+ (AssertionError, "AssertionError"),
+ (KeyError, "KeyError"),
+ (IndexError, "IndexError"),
+ (TypeError, "TypeError"),
+ (MemoryError, "MemoryError"),
+ (KeyboardInterrupt, "KeyboardInterrupt"),
+)
+
+
+def ple_fail_closed(
+ *,
+ where: str,
+ rank: int,
+ layer: str,
+ sequence: int,
+ error: BaseException,
+) -> NoReturn:
+ """Emit one best-effort nonblocking diagnostic and terminate immediately."""
+ try:
+ os.set_blocking(2, False)
+ error_type = type(error)
+ error_name = "Exception"
+ trusted_error = False
+ for trusted_type, trusted_name in _TRUSTED_ERROR_LABELS:
+ if error_type is trusted_type:
+ error_name = trusted_name
+ trusted_error = True
+ break
+
+ safe_where = where[:128] if type(where) is str else "unknown"
+ safe_layer = layer[:256] if type(layer) is str else "unknown"
+ safe_rank = rank if type(rank) is int else -1
+ safe_sequence = sequence if type(sequence) is int else -1
+ detail_suffix = ""
+ if trusted_error:
+ args = BaseException.args.__get__(error, error_type)
+ if (
+ type(args) is tuple
+ and len(args) == 1
+ and type(args[0]) is str
+ ):
+ detail = args[0][:512].replace("\n", "\\n")
+ detail_suffix = f":{detail}"
+ message = (
+ f"PLE fatal: where={safe_where} rank={safe_rank} layer={safe_layer} "
+ f"sequence={safe_sequence} error={error_name}{detail_suffix} "
+ f"pid={os.getpid()} "
+ f"monotonic={time.monotonic():.9f} wall={time.time():.9f}\n"
+ )
+ os.write(2, message.encode("utf-8", "backslashreplace")[:2048])
+ except BaseException:
+ pass
+ # THE TRACEBACK, unconditionally, in its own try so it cannot cost us the
+ # line above.
+ #
+ # This path os._exit(1)s, which runs no handlers and flushes nothing, so
+ # whatever is not written here is gone. The line above names the TYPE only
+ # when the type is one of _TRUSTED_ERROR_LABELS and flattens everything
+ # else to the literal string "Exception" -- and the flattened set includes
+ # the two that matter most here: torch.AcceleratorError (a RuntimeError
+ # SUBCLASS, so the `is` comparison misses it) and Triton CompilationError.
+ # An earlier boot died exactly that way: both PP0 ranks
+ # reported "error=Exception" and there was no traceback anywhere in the
+ # log, in any rank, so the cause could not be named at all and the next
+ # step was a re-run rather than a fix.
+ #
+ # A fail-closed path is where evidence matters most, not least: it fires
+ # once, it kills the process, and it is the only witness. Printing the type
+ # of a wrapped exception without the frames that raised it is not a
+ # diagnostic.
+ try:
+ # Three-arg form: it does not depend on being inside an active except
+ # block, and this function is called from one frame further out than
+ # the raise.
+ tb = "".join(
+ traceback.format_exception(type(error), error, error.__traceback__)
+ )
+ raw = tb.encode("utf-8", "backslashreplace")
+ # Truncate from the END, not the start. format_exception puts the
+ # frames first and the exception TYPE AND MESSAGE last, so a head-side
+ # cut throws away the one line that names what went wrong and keeps the
+ # import-time frames nobody needs. A long chained traceback is exactly
+ # the case where that matters.
+ clipped = len(raw) > 8192
+ raw = raw[-8192:]
+ os.write(
+ 2,
+ b"PLE fatal traceback (tail) follows:\n" if clipped
+ else b"PLE fatal traceback follows:\n",
+ )
+ os.write(2, raw)
+ if not raw.endswith(b"\n"):
+ os.write(2, b"\n")
+ except BaseException:
+ # Never silent HERE of all places: without this the reader cannot tell
+ # "the wrapper predates the traceback fix" from "the traceback itself
+ # could not be produced", and those call for different next steps.
+ # Still guarded, because os._exit below must run whatever happens: an
+ # escape from this handler would leave the worker alive in the state
+ # the fail-closed path exists to end.
+ try:
+ os.write(2, b"PLE fatal: traceback unavailable\n")
+ except BaseException:
+ pass
+ os._exit(1)
+
+
+def fail_ple_execute_model(func: _F) -> _F:
+ """Fail a non-shutdown worker-step escape only on PLE-owning ranks.
+
+ This includes ``sample_tokens``, the second half of the split execute-model
+ step. A rank without a PLE connector rethrows the original exception.
+ """
+
+ @functools.wraps(func)
+ def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
+ try:
+ return func(self, *args, **kwargs)
+ except SystemExit:
+ # WorkerProc's signal handler uses SystemExit for clean shutdown.
+ raise
+ except BaseException as error:
+ connector = getattr(self, "_ple_offload_connector", None)
+ if connector is None:
+ model_runner = getattr(self, "model_runner", None)
+ connector = getattr(
+ model_runner, "_ple_offload_connector", None
+ )
+ if connector is None:
+ raise
+ connector._fail_closed(func.__name__, error)
+
+ return wrapped # type: ignore[return-value]
diff --git a/vllm/v1/ple_offload/progress.py b/vllm/v1/ple_offload/progress.py
new file mode 100644
index 0000000..bf88291
--- /dev/null
+++ b/vllm/v1/ple_offload/progress.py
@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Independent progress deadline for PLE transactions."""
+
+import math
+import os
+import threading
+import time
+from collections import deque
+from dataclasses import dataclass
+from typing import Any
+
+from vllm.v1.ple_offload.fatal import ple_fail_closed
+
+_DEFAULT_SERVING_DEADLINE_S = 300.0
+_DEFAULT_STARTUP_DEADLINE_S = 1800.0
+_MAX_PROGRESS_DEADLINE_S = 86400.0
+_POLL_INTERVAL_S = 0.1
+
+
+def _validate_deadline(value: float, name: str) -> float:
+ value = float(value)
+ if (
+ not math.isfinite(value)
+ or value <= 0
+ or value > _MAX_PROGRESS_DEADLINE_S
+ ):
+ raise ValueError(
+ f"{name} must be finite, greater than zero, and no greater than "
+ f"{_MAX_PROGRESS_DEADLINE_S:g} seconds"
+ )
+ return value
+
+
+def _environment_deadline(name: str, default: float) -> float:
+ value = float(os.environ.get(name, str(default)))
+ return _validate_deadline(value, name)
+
+
+@dataclass(frozen=True)
+class _Submission:
+ sequence: int
+ started: float
+ startup: bool
+ num_tokens: int
+ deadline: float
+ depth: int
+
+
+@dataclass(frozen=True)
+class _PendingCompletion:
+ sequence: int
+ started: float
+ startup: bool
+ num_tokens: int
+ deadline: float
+ event: Any
+
+
+class PleProgressMonitor:
+ """Bound host submission and the resulting ordered GPU completions."""
+
+ def __init__(
+ self,
+ *,
+ rank: int,
+ layer_names: tuple[str, ...],
+ serving_deadline_s: float,
+ startup_deadline_s: float,
+ ) -> None:
+ serving_deadline_s = _validate_deadline(
+ serving_deadline_s, "serving_deadline_s"
+ )
+ startup_deadline_s = _validate_deadline(
+ startup_deadline_s, "startup_deadline_s"
+ )
+ self._rank = rank
+ self._layer = ",".join(layer_names)
+ self._serving_deadline_s = serving_deadline_s
+ self._startup_deadline_s = startup_deadline_s
+ self._condition = threading.Condition()
+ self._submission: _Submission | None = None
+ self._pending: deque[_PendingCompletion] = deque()
+ self._sequence = 0
+ self._stopping = False
+ self._thread = threading.Thread(
+ target=self._run,
+ name=f"ple-progress-rank{rank}",
+ daemon=True,
+ )
+ self._thread.start()
+
+ @classmethod
+ def from_environment(
+ cls, *, rank: int, layer_names: tuple[str, ...]
+ ) -> "PleProgressMonitor":
+ return cls(
+ rank=rank,
+ layer_names=layer_names,
+ serving_deadline_s=_environment_deadline(
+ "VLLM_PLE_PROGRESS_DEADLINE_S",
+ _DEFAULT_SERVING_DEADLINE_S,
+ ),
+ startup_deadline_s=_environment_deadline(
+ "VLLM_PLE_STARTUP_DEADLINE_S",
+ _DEFAULT_STARTUP_DEADLINE_S,
+ ),
+ )
+
+ @property
+ def sequence(self) -> int:
+ with self._condition:
+ return self._sequence
+
+ @property
+ def has_pending_completions(self) -> bool:
+ with self._condition:
+ return bool(self._pending)
+
+ def describe_submission(self, sequence: int) -> str:
+ """Return the armed transaction descriptor for the active sequence."""
+ with self._condition:
+ active = self._submission
+ if active is None or active.sequence != sequence:
+ raise RuntimeError("PLE fault correlation has no active submission")
+ return (
+ f"sequence:{active.sequence}"
+ f"@begin:{active.started:.9f}"
+ )
+
+ def completion_is_pending(self, sequence: int) -> bool:
+ """Whether this serving sequence has a retained, incomplete event.
+
+ This method is called only behind the armed fault guard. An event that
+ has completed on the GPU but has not yet been dequeued is not counted.
+ """
+ with self._condition:
+ pending = next(
+ (
+ item
+ for item in self._pending
+ if item.sequence == sequence and not item.startup
+ ),
+ None,
+ )
+ if pending is None:
+ return False
+ return not bool(pending.event.query())
+
+ def begin(self, *, startup: bool, num_tokens: int) -> int:
+ now = time.monotonic()
+ budget = (
+ self._startup_deadline_s if startup else self._serving_deadline_s
+ )
+ with self._condition:
+ active = self._submission
+ if active is not None:
+ if not startup or not active.startup:
+ raise RuntimeError(
+ "Overlapping PLE host submissions violate the serialized "
+ "single-slot protocol"
+ )
+ sequence = active.sequence
+ self._submission = _Submission(
+ sequence=sequence,
+ started=active.started,
+ startup=True,
+ num_tokens=num_tokens,
+ deadline=now + budget,
+ depth=active.depth + 1,
+ )
+ else:
+ self._sequence += 1
+ sequence = self._sequence
+ self._submission = _Submission(
+ sequence=sequence,
+ started=now,
+ startup=startup,
+ num_tokens=num_tokens,
+ deadline=now + budget,
+ depth=1,
+ )
+ self._condition.notify_all()
+ return sequence
+
+ def track_completion(self, event: Any) -> None:
+ """Retain a release event until the monitor observes GPU completion."""
+ if not callable(getattr(event, "query", None)):
+ raise TypeError("PLE completion event must provide query()")
+ with self._condition:
+ active = self._submission
+ if active is None:
+ raise RuntimeError("PLE completion has no matching submission")
+ self._pending.append(
+ _PendingCompletion(
+ sequence=active.sequence,
+ started=active.started,
+ startup=active.startup,
+ num_tokens=active.num_tokens,
+ deadline=active.deadline,
+ event=event,
+ )
+ )
+ self._condition.notify_all()
+
+ def complete(self) -> None:
+ with self._condition:
+ active = self._submission
+ if active is None or active.depth <= 0:
+ raise RuntimeError("PLE release has no matching transaction")
+ if active.depth == 1:
+ self._submission = None
+ else:
+ self._submission = _Submission(
+ sequence=active.sequence,
+ started=active.started,
+ startup=active.startup,
+ num_tokens=active.num_tokens,
+ deadline=active.deadline,
+ depth=active.depth - 1,
+ )
+ self._condition.notify_all()
+
+ def close(self) -> None:
+ with self._condition:
+ self._stopping = True
+ self._condition.notify_all()
+ self._thread.join(timeout=1)
+
+ def _run(self) -> None:
+ try:
+ self._monitor()
+ except BaseException as error:
+ ple_fail_closed(
+ where="progress-monitor-thread",
+ rank=self._rank,
+ layer=self._layer,
+ sequence=self._sequence,
+ error=error,
+ )
+
+ def _monitor(self) -> None:
+ while True:
+ with self._condition:
+ if self._stopping:
+ return
+ submission = self._submission
+ pending = self._pending[0] if self._pending else None
+ if submission is None and pending is None:
+ self._condition.wait(timeout=_POLL_INTERVAL_S)
+ continue
+
+ # CUDA completion is observed only here. Never synchronize and never
+ # use the host flag as evidence that the queued reset completed.
+ if pending is not None and pending.event.query():
+ with self._condition:
+ if self._pending and self._pending[0] is pending:
+ self._pending.popleft()
+ self._condition.notify_all()
+ continue
+
+ now = time.monotonic()
+ if pending is not None and now >= pending.deadline:
+ with self._condition:
+ if not self._pending or self._pending[0] is not pending:
+ continue
+ mode = "startup" if pending.startup else "serving"
+ ple_fail_closed(
+ where=(
+ f"progress-deadline:{mode}:gpu-completion:"
+ f"tokens={pending.num_tokens}"
+ ),
+ rank=self._rank,
+ layer=self._layer,
+ sequence=pending.sequence,
+ error=TimeoutError("PLE GPU completion made no bounded progress"),
+ )
+
+ if submission is not None and now >= submission.deadline:
+ with self._condition:
+ if self._submission is not submission:
+ continue
+ mode = "startup" if submission.startup else "serving"
+ ple_fail_closed(
+ where=(
+ f"progress-deadline:{mode}:host-submission:"
+ f"tokens={submission.num_tokens}"
+ ),
+ rank=self._rank,
+ layer=self._layer,
+ sequence=submission.sequence,
+ error=TimeoutError("PLE host submission made no bounded progress"),
+ )
+
+ deadlines = []
+ if pending is not None:
+ deadlines.append(pending.deadline)
+ if submission is not None:
+ deadlines.append(submission.deadline)
+ timeout = _POLL_INTERVAL_S
+ if deadlines:
+ timeout = min(
+ timeout,
+ max(0.0, min(deadlines) - time.monotonic()),
+ )
+ with self._condition:
+ if self._stopping:
+ return
+ if (
+ self._submission is not submission
+ or (self._pending[0] if self._pending else None) is not pending
+ ):
+ continue
+ self._condition.wait(timeout=timeout)
diff --git a/vllm/v1/ple_offload/protocol.py b/vllm/v1/ple_offload/protocol.py
new file mode 100644
index 0000000..87eba13
--- /dev/null
+++ b/vllm/v1/ple_offload/protocol.py
@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""IPC message definitions for PLE CPU offload (host-mapped pull design).
+
+The offload process owns no CUDA context. Every GPU worker owns, per PLE
+layer, a shared host output buffer and a shared host int32 flag, both
+page-locked and device-mapped in the worker's context. The offload process
+waits for ``flag == 0``, writes the layer result into the host buffer and
+stores ``flag = 1`` with plain CPU stores; the GPU's captured forward waits on
+the host flag, DMAs host -> GPU on the model stream, and ``release_outputs``
+resets the flag after the forward.
+"""
+
+from dataclasses import dataclass
+
+import msgspec
+import torch
+
+# ---------------------------------------------------------------------------
+# IPC message dataclasses
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class PleOffloadRegistration:
+ """Sent once from each GPU worker during offload setup."""
+
+ worker_id: int
+ tp_rank: int
+ dp_rank: int
+ # Shared CPU tensors per PLE layer, filled by the offload process.
+ host_output_buffers: dict[str, torch.Tensor]
+ # Shared CPU int32 tensors per PLE layer; element 0 is the ready flag.
+ host_flag_tensors: dict[str, torch.Tensor]
+ # CPU tensors are allocated in shared memory and registered once.
+ input_ids_buf: torch.Tensor
+ query_start_loc_buf: torch.Tensor
+ ngram_context_buf: torch.Tensor | None
+
+
+@dataclass
+class PleOffloadRequest:
+ """Sent by each DP rank's TP rank zero at every inference step."""
+
+ dp_rank: int
+ num_tokens: int
+ num_reqs: int
+
+
+_PLE_OFFLOAD_REQUEST_DECODER = msgspec.msgpack.Decoder(PleOffloadRequest)
diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py
new file mode 100644
index 0000000..43caf80
--- /dev/null
+++ b/vllm/v1/ple_offload/worker.py
@@ -0,0 +1,1298 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dedicated CPU-offload process for PLE embedding layers.
+
+This module implements a standalone process that:
+1. Loads only the :class:`PleOffloadLayer` weights into CPU memory.
+2. Accepts per-step computation requests from GPU worker processes.
+3. Runs ``forward_impl()`` on CPU, writes the rows into every TP worker's
+ shared host output buffer for the requesting DP rank and sets that
+ worker's shared host flag (host-mapped pull design: this process holds no
+ CUDA context; the GPU waits on the host flag and DMAs host -> GPU inside
+ its captured forward).
+
+The TP workers within one DP rank receive identical inputs, so the CPU result
+is computed once per DP rank and fanned out to all of its TP ranks.
+
+Class structure mirrors the GPU worker pattern in multiproc_executor.py:
+
+ PleOffloadWorkerHandle -- handle held by the spawning GPU worker
+ PleOffloadWorker -- process lifecycle and READY handshake
+ PleOffloadRunner -- owns weights and serves inference requests
+"""
+
+import contextlib
+import os
+import pickle
+import select
+import signal
+import socket
+import struct
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from collections.abc import Iterable
+from dataclasses import dataclass
+from multiprocessing.connection import Connection
+from typing import Any, cast
+
+import msgspec
+import torch
+import torch.distributed as dist
+import zmq
+
+import vllm.envs as envs
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed.parallel_state import (
+ ensure_model_parallel_initialized,
+ init_distributed_environment,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.ple_offload_layer import (
+ PleOffloadLayer,
+ mark_as_offload_worker,
+)
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
+from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
+from vllm.model_executor.model_loader.utils import (
+ initialize_model,
+ process_weights_after_loading,
+)
+from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights
+from vllm.utils.system_utils import decorate_logs
+from vllm.utils.torch_utils import set_default_torch_dtype
+from vllm.v1.ple_offload.protocol import (
+ _PLE_OFFLOAD_REQUEST_DECODER,
+ PleOffloadRegistration,
+ PleOffloadRequest,
+)
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm_ple_test_fault import (
+ _FAULT,
+ _FaultReconstructionPayload,
+)
+
+logger = init_logger(__name__)
+
+# Publication fence for the host-mapped pull: a locked read-modify-write (the
+# uncontended lock acquire) is a full barrier on x86 and also drains
+# write-combining buffers, so every row store of a result is globally visible
+# before the flag store that tells the GPU to DMA it (do not rely
+# on "memcpy uses no non-temporal stores at this size").
+_PUBLISH_FENCE = threading.Lock()
+
+
+def _unlink_ipc_socket(ipc_addr: str) -> None:
+ """Best-effort unlink of the exact filesystem IPC endpoint PLE owns."""
+ if not ipc_addr.startswith("ipc://"):
+ return
+ path = ipc_addr.removeprefix("ipc://")
+ # PLE receives an absolute UUID path from get_open_zmq_ipc_path(). Avoid
+ # treating abstract or relative ZMQ addresses as filesystem ownership.
+ if not os.path.isabs(path):
+ return
+ try:
+ os.unlink(path)
+ except FileNotFoundError:
+ pass
+ except OSError as error:
+ logger.warning("Failed to unlink PLE IPC socket %s: %s", path, error)
+
+
+def _publish_fence() -> None:
+ with _PUBLISH_FENCE:
+ pass
+
+
+def _offload_child_environment() -> dict[str, str]:
+ """Build the offload environment without mutating the parent."""
+ child_env = dict(os.environ)
+ child_env["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1"
+ child_env.setdefault("CUTLASS_CUDA_DEVICE_ID", "0")
+ child_env.pop("VLLM_PP_LAYER_PARTITION", None)
+
+ home_device = child_env.get("VLLM_PLE_OFFLOAD_HOME_DEVICE", "").strip()
+ if home_device:
+ parent_visible = child_env.get("CUDA_VISIBLE_DEVICES", "").strip()
+ try:
+ physical = (
+ parent_visible.split(",")[int(home_device)].strip()
+ if parent_visible
+ else home_device
+ )
+ except (IndexError, ValueError) as error:
+ raise ValueError("Invalid VLLM_PLE_OFFLOAD_HOME_DEVICE") from error
+ child_env["CUDA_VISIBLE_DEVICES"] = physical
+ return child_env
+
+
+def _prepare_cuda_free_process() -> None:
+ """Keep the PLE offload process free of any CUDA context.
+
+ The host-mapped pull design issues no CUDA work here, yet model
+ construction reached CUDA incidentally (found by the offload-startup
+ replays):
+ 1. ``VLLM_DISABLE_SHARED_EXPERTS_STREAM=1``: each MoE layer's shared-experts
+ runner calls ``aux_stream()`` -> ``torch.cuda.Stream()``, whose C++
+ constructor creates the primary context without calling
+ ``torch.cuda._lazy_init`` (``torch.cuda.is_initialized()`` stays False).
+ This is the context the replay caught (stack dumped inside
+ ``Stream.__new__``). This process never runs an MoE layer, so it never
+ needs the stream. vLLM reads the variable when the layer is built; this
+ process does not enable vLLM's env cache.
+ 2. ``CUTLASS_CUDA_DEVICE_ID``: CUTLASS python (``cutlass_cppgen``) runs
+ ``initialize_cuda_context()`` -> ``cudart.cudaFree(0)`` at import unless
+ it is set. Not observed on this import path; kept as a cheap defence.
+ 3. The torch guard (``_install_cuda_guard``).
+ """
+ os.environ["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1"
+ os.environ.setdefault("CUTLASS_CUDA_DEVICE_ID", "0")
+ _install_cuda_guard()
+
+
+def _install_cuda_guard() -> None:
+ """Make torch-level CUDA initialization raise in the PLE offload process.
+
+ An offload-startup check found the host-mapped pull still got a 322 MiB context
+ on its home card: importing the Qwen4Exp model imports the vendored FLA
+ ops, whose utils call triton's ``get_current_target()`` at import time,
+ i.e. ``torch.cuda.current_device()`` -> lazy init. FLA catches
+ ``RuntimeError`` and falls back to "cpu" (and then skips its Hopper
+ probe), so raising here keeps this process context-free. Any other call
+ site fails loudly at offload startup with its stack logged, instead of
+ silently taking VRAM on a stage-0 card. The GPU workers are untouched.
+
+ torch also has C++ entry points that create the context without calling
+ ``torch.cuda._lazy_init``: stream construction (``torch.cuda.Stream``) and
+ the pinned host allocator (``Tensor.pin_memory``). Both raise here too;
+ this process needs neither (all of its buffers are plain CPU memory).
+ ``torch.cuda.is_initialized()`` only reflects the Python-level path, so
+ the per-PID GPU sampler, not that flag, is the authority on contexts.
+ """
+ import traceback
+
+ if torch.cuda.is_initialized():
+ # Fail closed: a context that already exists would stay on a
+ # stage-0 card for the whole run.
+ raise RuntimeError(
+ "CUDA is already initialized in the PLE offload process before its "
+ "guard; the host-mapped pull offload must stay CUDA-free"
+ )
+
+ def _refuse(what: str) -> None:
+ stack = "".join(traceback.format_stack(limit=12)[:-2])
+ logger.info(
+ "PLE offload CUDA guard: blocked a CUDA initialization attempt "
+ "(%s; the caller must fall back to CPU or fail); call site:\n%s",
+ what,
+ stack,
+ )
+ raise RuntimeError(
+ f"CUDA is disabled in the PLE offload process (host-mapped pull): {what}"
+ )
+
+ def _blocked_lazy_init() -> None:
+ _refuse("torch.cuda lazy init")
+
+ def _blocked_stream(cls, *args, **kwargs):
+ _refuse("torch.cuda.Stream construction")
+
+ def _blocked_pin_memory(self, *args, **kwargs):
+ _refuse("Tensor.pin_memory")
+
+ torch.cuda._lazy_init = _blocked_lazy_init # type: ignore[assignment]
+ torch.cuda.Stream.__new__ = staticmethod(_blocked_stream) # type: ignore[assignment]
+ torch.Tensor.pin_memory = _blocked_pin_memory # type: ignore[assignment]
+ logger.info(
+ "PLE offload CUDA guard installed (torch CUDA lazy init, stream "
+ "construction and pin_memory raise)"
+ )
+
+
+@contextlib.contextmanager
+def _platform_device_is_meta():
+ """During the offload's meta-model construction, report the meta device.
+
+ Some layer constructors pass an explicit ``current_platform.current_device()``
+ (e.g. the GDN output norm in qwen_gdn_linear_attn.py), which overrides the
+ ``torch.device("meta")`` context and would allocate on CUDA. In the offload
+ process nothing but the CPU PLE layers is ever materialized, so those
+ tensors belong on meta too. The patch is an instance attribute removed on
+ exit, restoring the platform's own method.
+ """
+ from vllm.platforms import current_platform
+
+ current_platform.current_device = lambda *args, **kwargs: torch.device("meta")
+ try:
+ yield
+ finally:
+ with contextlib.suppress(AttributeError):
+ del current_platform.current_device
+
+
+# Bounded wait for the GPU to release a result (seconds). A timed-out ownership
+# transfer is fatal: the offload process raises and exits, never overwrites.
+_PLE_WAIT_TIMEOUT_S = float(os.environ.get("VLLM_PLE_OFFLOAD_WAIT_TIMEOUT", "120"))
+
+
+@dataclass
+class PleOffloadOutputTarget:
+ """Host output destination and ready flag for one TP worker.
+
+ Both are shared CPU tensors owned by the GPU worker, page-locked and
+ device-mapped in that worker's context. This process only reads/writes
+ them with plain CPU loads/stores (no CUDA): wait for ``flag == 0`` (the GPU
+ released the previous result), write the rows, then store ``flag = 1``.
+ """
+
+ tp_rank: int
+ host_output_buffer: torch.Tensor
+ host_flag: torch.Tensor # int32, element 0 is the flag
+
+
+@dataclass
+class PleOffloadInputBuffers:
+ """Shared-memory input buffers registered for one DP rank."""
+
+ input_ids_buf: torch.Tensor # int32 (max_num_tokens,)
+ query_start_loc_buf: torch.Tensor # int32 (max_num_reqs + 1,)
+ ngram_context_buf: torch.Tensor | None # int32 (max_num_reqs, ngram_context_len)
+
+
+@dataclass(frozen=True)
+class _ReceivedPleOffloadRequest:
+ request: PleOffloadRequest
+ serving: bool
+ transaction_id: str | None
+ correlation_id: str | None
+
+
+def _decode_offload_request(
+ frames: list[bytes],
+) -> _ReceivedPleOffloadRequest:
+ if len(frames) != 2:
+ raise RuntimeError("PLE offload request must contain two frames")
+ request = _PLE_OFFLOAD_REQUEST_DECODER.decode(frames[0])
+ metadata = msgspec.msgpack.decode(frames[1])
+ if type(metadata) is not dict or set(metadata) != {
+ "serving",
+ "transaction",
+ "correlation",
+ }:
+ raise RuntimeError("PLE offload request metadata is invalid")
+ serving = metadata["serving"]
+ transaction = metadata["transaction"]
+ correlation = metadata["correlation"]
+ if type(serving) is not bool:
+ raise RuntimeError("PLE offload serving metadata is invalid")
+ if transaction is not None and type(transaction) is not str:
+ raise RuntimeError("PLE offload transaction metadata is invalid")
+ if correlation is not None and type(correlation) is not str:
+ raise RuntimeError("PLE offload correlation metadata is invalid")
+ return _ReceivedPleOffloadRequest(
+ request=request,
+ serving=serving,
+ transaction_id=transaction,
+ correlation_id=correlation,
+ )
+
+
+def _receive_offload_request(
+ socket: zmq.Socket, flags: int = 0
+) -> _ReceivedPleOffloadRequest | PleOffloadRequest:
+ """Disarmed: the unchanged single-frame request. Armed: request + metadata."""
+ if _FAULT is None:
+ return _PLE_OFFLOAD_REQUEST_DECODER.decode(socket.recv(flags))
+ return _decode_offload_request(socket.recv_multipart(flags))
+
+
+def _send_framed_payload(fd: int, payload: bytes, deadline: float) -> None:
+ """Send one launch payload without blocking beyond an absolute deadline."""
+ frame = memoryview(struct.pack("!Q", len(payload)) + payload)
+ offset = 0
+ os.set_blocking(fd, False)
+ try:
+ while offset < len(frame):
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ raise TimeoutError("PLE launch payload transmission timed out")
+ _, writable, _ = select.select([], [fd], [], remaining)
+ if not writable:
+ raise TimeoutError("PLE launch payload transmission timed out")
+ try:
+ sent = os.write(fd, frame[offset:])
+ except BlockingIOError:
+ continue
+ if sent <= 0:
+ raise BrokenPipeError("PLE launch control socket closed")
+ offset += sent
+ if (
+ _FAULT is not None
+ and _FAULT.cut == "launch_park_before_payload_read"
+ ):
+ _FAULT.acknowledge_launch_send(
+ bytes_sent=offset,
+ frame_size=len(frame),
+ )
+ finally:
+ with contextlib.suppress(OSError):
+ os.set_blocking(fd, True)
+
+
+_PLE_READY_MAX_PAYLOAD_BYTES = 1 << 20
+
+class _BoundedStartupResponseWriter:
+ """Preflight startup responses against the parent receiver's frame cap."""
+
+ def __init__(self, writer: Connection) -> None:
+ self._writer = writer
+
+ def send(self, message: Any) -> None:
+ payload = pickle.dumps(message, protocol=pickle.HIGHEST_PROTOCOL)
+ if len(payload) > _PLE_READY_MAX_PAYLOAD_BYTES:
+ payload = pickle.dumps(
+ {
+ "status": "FAILURE",
+ "error": "PLE startup response exceeded size limit",
+ },
+ protocol=pickle.HIGHEST_PROTOCOL,
+ )
+ if len(payload) > _PLE_READY_MAX_PAYLOAD_BYTES:
+ raise RuntimeError(
+ "Bounded PLE startup failure response exceeds its frame cap"
+ )
+ self._writer.send_bytes(payload)
+
+ def close(self) -> None:
+ self._writer.close()
+
+
+def _read_exact_before_deadline(
+ fd: int,
+ size: int,
+ deadline: float,
+) -> bytes:
+ """Read exactly ``size`` bytes without crossing an absolute deadline."""
+ chunks = bytearray()
+ while len(chunks) < size:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ raise TimeoutError(
+ "PLE offload worker did not become ready before its "
+ "absolute launch deadline"
+ )
+ readable, _, _ = select.select([fd], [], [], remaining)
+ if not readable:
+ raise TimeoutError(
+ "PLE offload worker did not become ready before its "
+ "absolute launch deadline"
+ )
+ try:
+ chunk = os.read(fd, size - len(chunks))
+ except BlockingIOError:
+ continue
+ if not chunk:
+ raise EOFError("PLE startup response was truncated")
+ chunks.extend(chunk)
+ return bytes(chunks)
+
+
+def _recv_bounded_connection_message(
+ reader: Connection,
+ deadline: float,
+) -> Any:
+ """Receive one bounded multiprocessing Connection frame by ``deadline``."""
+ fd = reader.fileno()
+ os.set_blocking(fd, False)
+ try:
+ header = _read_exact_before_deadline(fd, 4, deadline)
+ size = struct.unpack("!i", header)[0]
+ if size == -1:
+ size = struct.unpack(
+ "!Q", _read_exact_before_deadline(fd, 8, deadline)
+ )[0]
+ elif size < 0:
+ raise RuntimeError("PLE startup response has an invalid frame size")
+ if size > _PLE_READY_MAX_PAYLOAD_BYTES:
+ raise RuntimeError(
+ f"PLE startup response exceeds "
+ f"{_PLE_READY_MAX_PAYLOAD_BYTES} bytes"
+ )
+ frame = _read_exact_before_deadline(fd, size, deadline)
+ finally:
+ with contextlib.suppress(OSError):
+ os.set_blocking(fd, True)
+ try:
+ return pickle.loads(frame)
+ except BaseException as error:
+ raise RuntimeError(
+ "PLE startup response is not a valid serialized message"
+ ) from error
+
+
+@dataclass
+class PleOffloadWorkerHandle:
+ """Resources owned by the GPU worker that spawned the offload process."""
+
+ proc: subprocess.Popen
+ death_writer: Connection | None
+ ready_pipe_reader: Connection | None
+ launch_deadline: float
+ ipc_addr: str | None = None
+
+ def close(self, *, graceful: bool = True) -> None:
+ """Release all process resources. Safe to call more than once."""
+ if self.ready_pipe_reader is not None:
+ self.ready_pipe_reader.close()
+ self.ready_pipe_reader = None
+ if graceful and self.proc.poll() is None:
+ # Keep parent-death supervision connected while SIGTERM lets
+ # proc_main run its shutdown event and finally cleanup.
+ self.proc.terminate()
+ try:
+ self.proc.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ pass
+ if self.death_writer is not None:
+ self.death_writer.close()
+ self.death_writer = None
+ if not graceful and self.proc.poll() is None:
+ # Expired launch: retain the immediate EOF -> TERM path.
+ self.proc.terminate()
+ try:
+ self.proc.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ pass
+ if self.proc.poll() is None:
+ self.proc.kill()
+ try:
+ self.proc.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ pass
+ # The child normally removes the endpoint in proc_main's finally.
+ # Once it has exited, this exact PLE-owned address is safe to remove as
+ # a fallback for SIGKILL or an interrupted child cleanup.
+ if self.proc.poll() is not None and self.ipc_addr is not None:
+ _unlink_ipc_socket(self.ipc_addr)
+
+
+def _init_offload_distributed() -> None:
+ """Initialize the single-rank Gloo world required by TP-aware layers."""
+ if dist.is_initialized():
+ return
+
+ # VocabParallelEmbedding reads the TP process group during construction.
+ # The offload process owns the full embedding table, so it uses an isolated
+ # TP1/PP1 Gloo world and never joins the GPU workers' NCCL groups.
+ store_dir = tempfile.mkdtemp(prefix="vllm_ple_offload_")
+ init_distributed_environment(
+ world_size=1,
+ rank=0,
+ distributed_init_method=f"file://{store_dir}/store",
+ local_rank=0,
+ backend="gloo",
+ )
+ # initialize_model_parallel reads the active VllmConfig in the current
+ # vLLM version. Explicitly configure DP1/TP1/PP1 to match the isolated
+ # world, regardless of any DP environment variables inherited from the GPU
+ # worker. The real DP/TP configuration is used later for model construction,
+ # registration, and request routing.
+ offload_config = VllmConfig()
+ offload_parallel_config = offload_config.parallel_config
+ offload_parallel_config.data_parallel_size = 1
+ offload_parallel_config.data_parallel_size_local = 1
+ offload_parallel_config.data_parallel_rank = 0
+ offload_parallel_config.data_parallel_rank_local = 0
+ offload_parallel_config.data_parallel_index = 0
+ offload_parallel_config.tensor_parallel_size = 1
+ offload_parallel_config.pipeline_parallel_size = 1
+ offload_parallel_config.prefill_context_parallel_size = 1
+ offload_parallel_config.decode_context_parallel_size = 1
+ offload_parallel_config.world_size = 1
+ offload_parallel_config.nnodes = 1
+ offload_parallel_config.node_rank = 0
+ with set_current_vllm_config(offload_config):
+ ensure_model_parallel_initialized(
+ tensor_model_parallel_size=1,
+ pipeline_model_parallel_size=1,
+ backend="gloo",
+ )
+ logger.info(
+ "Distributed environment initialized (backend=gloo, rank=0, world_size=1)."
+ )
+
+
+class PleOffloadWorker:
+ """Manage process creation, READY handshake, and the child entry point."""
+
+ READY_STR = "READY"
+
+ @staticmethod
+ def make_process(
+ vllm_config: VllmConfig,
+ num_workers: int,
+ ipc_addr: str,
+ ) -> PleOffloadWorkerHandle:
+ """Start one CPU offload process in a fresh interpreter."""
+ launch_deadline = (
+ time.monotonic() + float(envs.VLLM_PLE_OFFLOAD_READY_TIMEOUT)
+ )
+ config_payload = pickle.dumps(
+ vllm_config, protocol=pickle.HIGHEST_PROTOCOL
+ )
+ child_env = _offload_child_environment()
+
+ # READY and parent-death supervision must never compete for reads.
+ parent_control, child_control = socket.socketpair()
+ parent_death, child_death = socket.socketpair()
+ ready_reader = Connection(parent_control.detach())
+ death_writer = Connection(parent_death.detach())
+ try:
+ proc = subprocess.Popen(
+ [
+ sys.executable,
+ "-m",
+ "vllm_ple_offload_bootstrap",
+ str(child_control.fileno()),
+ str(child_death.fileno()),
+ ],
+ env=child_env,
+ close_fds=True,
+ pass_fds=(child_control.fileno(), child_death.fileno()),
+ )
+ except BaseException:
+ ready_reader.close()
+ death_writer.close()
+ raise
+ finally:
+ child_control.close()
+ child_death.close()
+
+ handle = PleOffloadWorkerHandle(
+ proc=proc,
+ death_writer=death_writer,
+ ready_pipe_reader=ready_reader,
+ launch_deadline=launch_deadline,
+ ipc_addr=ipc_addr,
+ )
+ try:
+ config_for_launch: bytes | _FaultReconstructionPayload = config_payload
+ if (
+ _FAULT is not None
+ and _FAULT.cut == "launch_block_reconstruction"
+ ):
+ config_for_launch = _FaultReconstructionPayload(config_payload)
+ launch_state = {
+ "vllm_config_payload": config_for_launch,
+ "num_workers": num_workers,
+ "ipc_addr": ipc_addr,
+ }
+ if (
+ _FAULT is not None
+ and _FAULT.cut == "launch_park_before_payload_read"
+ ):
+ # Larger than ordinary AF_UNIX socket buffers: an acknowledged
+ # non-reading child leaves the parent inside framed transmission.
+ launch_state["_test_padding"] = b"\0" * (8 << 20)
+ launch_payload = pickle.dumps(
+ launch_state,
+ protocol=pickle.HIGHEST_PROTOCOL,
+ )
+ _send_framed_payload(
+ ready_reader.fileno(), launch_payload, launch_deadline
+ )
+ except BaseException:
+ handle.close(graceful=False)
+ raise
+
+ return handle
+
+ @staticmethod
+ def wait_for_ready(handle: PleOffloadWorkerHandle) -> None:
+ """Wait until weights and all GPU registrations are ready to serve."""
+ reader = handle.ready_pipe_reader
+ if reader is None:
+ return
+ try:
+ message = _recv_bounded_connection_message(
+ reader, handle.launch_deadline
+ )
+ if type(message) is not dict:
+ raise RuntimeError("PLE offload worker sent an invalid response")
+ status = dict.get(message, "status")
+ if type(status) is not str:
+ raise RuntimeError(
+ "PLE offload worker sent an invalid response status"
+ )
+ if status == PleOffloadWorker.READY_STR:
+ pass
+ elif status == "FAILURE":
+ error_detail = dict.get(message, "error")
+ if type(error_detail) is not str:
+ error_detail = "unknown error"
+ raise RuntimeError(
+ "PLE offload worker failed during startup: "
+ f"{error_detail}"
+ )
+ else:
+ raise RuntimeError("PLE offload worker sent unsupported status")
+ layer_names = dict.get(message, "layer_names")
+ if (
+ type(layer_names) is not list
+ or any(
+ type(name) is not str for name in layer_names
+ )
+ ):
+ raise RuntimeError(
+ "PLE offload worker sent invalid READY layer_names"
+ )
+ logger.info(
+ "Worker ready - %d PleOffloadLayer(s): %s",
+ len(layer_names),
+ layer_names,
+ )
+ except EOFError as error:
+ handle.close(graceful=False)
+ raise RuntimeError("PLE offload worker exited during startup") from error
+ except BaseException:
+ handle.close(graceful=False)
+ raise
+ finally:
+ if handle.ready_pipe_reader is not None:
+ handle.ready_pipe_reader.close()
+ handle.ready_pipe_reader = None
+
+ @staticmethod
+ def proc_main(
+ vllm_config: VllmConfig,
+ num_workers: int,
+ ipc_addr: str,
+ ready_pipe: Connection,
+ ) -> None:
+ """Load PLE weights, accept registrations, and run the request loop."""
+ decorate_logs("PleOffloadWorker")
+ ready_writer = _BoundedStartupResponseWriter(ready_pipe)
+ shutdown_event = threading.Event()
+
+ def handle_signal(_signum: int, _frame: object) -> None:
+ shutdown_event.set()
+
+ signal.signal(signal.SIGTERM, handle_signal)
+ signal.signal(signal.SIGINT, handle_signal)
+
+ zmq_context: zmq.Context | None = None
+ pull_socket: zmq.Socket | None = None
+ owns_ipc_socket = False
+ try:
+ # The flag lets PleOffloadLayer subclasses execute their complete
+ # constructors instead of becoming empty GPU-worker placeholders.
+ mark_as_offload_worker()
+
+ # The bootstrap established child-only visibility before importing
+ # torch or vLLM. Retain the Python-level guards for known CUDA
+ # initialization paths and fail if initialization already occurred.
+ _prepare_cuda_free_process()
+ home_device = os.environ.get("VLLM_PLE_OFFLOAD_HOME_DEVICE", "").strip()
+ if home_device:
+ logger.info(
+ "PLE offload bootstrap visibility is CUDA_VISIBLE_DEVICES=%s "
+ "(host-mapped pull: no CUDA context expected)",
+ os.environ.get("CUDA_VISIBLE_DEVICES", ""),
+ )
+
+ # Initialize Gloo before installing the real VllmConfig. This keeps
+ # the CPU process in an isolated rank-zero, world-size-one group.
+ # The offload worker is a PP=1 world; a VLLM_PP_LAYER_PARTITION set
+ # for the GPU workers (e.g. "25,23") would make make_layers raise
+ # "len(partitions) does not match pp_size=1" while the meta model
+ # is built here. Drop it in this process only.
+ os.environ.pop("VLLM_PP_LAYER_PARTITION", None)
+ _init_offload_distributed()
+
+ # Model components read the active VllmConfig while the meta model
+ # is constructed, so keep the context around runner initialization.
+ with set_current_vllm_config(vllm_config):
+ runner = PleOffloadRunner(vllm_config)
+
+ zmq_context = zmq.Context()
+ pull_socket = zmq_context.socket(zmq.PULL)
+ pull_socket.bind(ipc_addr)
+ owns_ipc_socket = True
+ logger.info(
+ "Bound IPC address %s; waiting for %d GPU worker registration(s).",
+ ipc_addr,
+ num_workers,
+ )
+
+ # READY means that the process can immediately serve requests. Wait
+ # for every DP/TP worker to register before notifying the parent.
+ runner.accept_registrations(pull_socket, num_workers)
+ # Checked: the host-mapped pull design must not create a context.
+ logger.info(
+ "PLE offload host-mapped pull: CUDA initialized in this process = %s",
+ torch.cuda.is_initialized(),
+ )
+ ready_writer.send(
+ {
+ "status": PleOffloadWorker.READY_STR,
+ "layer_names": sorted(runner.layer_names),
+ }
+ )
+ ready_writer.close()
+ ready_writer = None # type: ignore[assignment]
+
+ runner.busy_loop(pull_socket, shutdown_event)
+ except Exception as error:
+ logger.exception("Unexpected failure in PLE offload worker.")
+ if ready_writer is not None:
+ with contextlib.suppress(Exception):
+ ready_writer.send(
+ {"status": "FAILURE", "error": repr(error)[:4096]}
+ )
+ raise
+ finally:
+ if pull_socket is not None:
+ pull_socket.close(linger=0)
+ if zmq_context is not None:
+ zmq_context.term()
+ if owns_ipc_socket:
+ _unlink_ipc_socket(ipc_addr)
+ if ready_writer is not None:
+ ready_writer.close()
+
+
+class PleOffloadRunner:
+ """Own all discovered PLE tables and serve every local DP rank."""
+
+ def __init__(self, vllm_config: VllmConfig) -> None:
+ self.vllm_config = vllm_config
+ self._clamp_input_ids = (
+ getattr(vllm_config, "speculative_config", None) is not None
+ )
+ # name -> PleOffloadLayer (CPU)
+ self._layers: dict[str, PleOffloadLayer] = {}
+ # dp_rank -> layer_name -> one destination per TP rank
+ self._worker_targets: dict[int, dict[str, list[PleOffloadOutputTarget]]] = {}
+ # Each (dp_rank, layer_name) pair owns a plain CPU scratch buffer that
+ # forward_impl writes; the rows are then copied into every TP worker's
+ # shared host buffer. Nothing in this process is pinned or on a GPU.
+ self._scratch_bufs: dict[int, dict[str, torch.Tensor]] = {}
+ self._shutdown_event: threading.Event | None = None
+ self._step = 0 # number of handled request batches (diagnostics)
+ self._request_index = 0
+ # Shared-memory inputs are registered once per DP rank by TP rank zero.
+ self._input_bufs: dict[int, PleOffloadInputBuffers] = {}
+ self._load_weights()
+
+ @property
+ def layer_names(self) -> list[str]:
+ """Return PleOffloadLayer names in model traversal order."""
+ return list(self._layers)
+
+ def _load_weights(self) -> None:
+ """Load only :class:`PleOffloadLayer` subtrees into CPU memory.
+
+ Strategy:
+ 1. Build the entire model on ``meta`` so non-offloaded parameters use no
+ physical memory. PleOffloadLayer constructors explicitly target CPU.
+ 2. Discover all PleOffloadLayer modules from the complete model.
+ 3. Stream the checkpoint through a prefix filter so only matching PLE
+ tensors are materialized and passed to ``model.load_weights``.
+ 4. Run post-load processing only on the CPU-owned PLE subtrees.
+ """
+ model_config = self.vllm_config.model_config
+ load_config = self.vllm_config.load_config
+
+ # Step 1: build complete structure, while only PLE subtrees allocate CPU
+ # memory. All transformer, MoE, and vision parameters remain on meta.
+ logger.info("Initializing model structure for PLE weight discovery ...")
+ model_dtype = cast(torch.dtype, model_config.dtype)
+ with (
+ set_default_torch_dtype(model_dtype),
+ torch.device("meta"),
+ _platform_device_is_meta(),
+ ):
+ model = initialize_model(
+ vllm_config=self.vllm_config,
+ model_config=model_config,
+ )
+
+ # Step 2: preserve named_modules DFS order so CPU execution follows the
+ # same layer order as the GPU model forward.
+ offload_layers = {
+ name: module
+ for name, module in model.named_modules()
+ if isinstance(module, PleOffloadLayer)
+ }
+ if not offload_layers:
+ raise RuntimeError(
+ "VLLM_PLE_CPU_OFFLOAD is enabled, but no PleOffloadLayer "
+ "was found in the initialized model"
+ )
+ logger.info(
+ "Found %d PleOffloadLayer(s): %s",
+ len(offload_layers),
+ sorted(offload_layers),
+ )
+ offload_prefixes = tuple(f"{name}." for name in offload_layers)
+
+ # Step 3: filter checkpoint tensors before model.load_weights(). The
+ # conditional-generation checkpoint uses HF names such as
+ # ``model.language_model.*`` while named_modules exposes mapped vLLM
+ # names such as ``language_model.model.*``. Apply the model mapper only
+ # for matching, then yield the original pair so load_weights performs
+ # its normal single mapping pass.
+ mapper = getattr(model, "hf_to_vllm_mapper", None)
+ matched_checkpoint_tensors = 0
+
+ def offload_only_iter(
+ weights: Iterable[tuple[str, torch.Tensor]],
+ ) -> Iterable[tuple[str, torch.Tensor]]:
+ nonlocal matched_checkpoint_tensors
+ for weight_name, tensor in weights:
+ mapped_name: str | None = weight_name
+ if mapper is not None:
+ mapped_names = mapper.apply_list([weight_name])
+ mapped_name = mapped_names[0] if mapped_names else None
+ if mapped_name is not None and mapped_name.startswith(offload_prefixes):
+ matched_checkpoint_tensors += 1
+ yield weight_name, tensor
+
+ loader = get_model_loader(load_config)
+ if isinstance(loader, DummyModelLoader):
+ logger.info(
+ "Initializing dummy weights for %d PleOffloadLayer(s) ...",
+ len(offload_layers),
+ )
+ for layer in offload_layers.values():
+ initialize_dummy_weights(layer, model_config)
+ elif isinstance(loader, DefaultModelLoader):
+ all_weights = loader.get_all_weights(model_config, model)
+ loaded_params = model.load_weights(offload_only_iter(all_weights))
+ if matched_checkpoint_tensors == 0:
+ raise RuntimeError(
+ "PLE offload checkpoint filter matched no weights for "
+ f"layers: {sorted(offload_layers)}"
+ )
+
+ expected_offload_params = {
+ f"{layer_name}.{param_name}"
+ for layer_name, layer in offload_layers.items()
+ for param_name, _ in layer.named_parameters()
+ }
+ loaded_offload_entries = {
+ name for name in loaded_params if name.startswith(offload_prefixes)
+ }
+ loaded_expected_params = expected_offload_params.intersection(loaded_params)
+ missing_offload_params = sorted(
+ expected_offload_params.difference(loaded_expected_params)
+ )
+ if missing_offload_params:
+ raise RuntimeError(
+ "PLE offload checkpoint did not load all materialized "
+ f"parameters: {missing_offload_params}"
+ )
+ logger.info(
+ "PLE offload matched %d checkpoint tensor(s), loaded %d "
+ "offload entries, and verified %d/%d materialized "
+ "parameter(s) for layers: %s",
+ matched_checkpoint_tensors,
+ len(loaded_offload_entries),
+ len(loaded_expected_params),
+ len(expected_offload_params),
+ sorted(offload_layers),
+ )
+ else:
+ raise NotImplementedError(
+ "PLE offload requires the default or dummy model loader, got "
+ f"{type(loader).__name__}"
+ )
+
+ # Step 4: post-load processing is restricted to CPU-owned PLE modules;
+ # the remainder of the model is still on meta and must not be visited.
+ for layer in offload_layers.values():
+ process_weights_after_loading(layer, model_config, torch.device("cpu"))
+
+ self._layers.update(offload_layers)
+ del model
+ logger.info("PLE weight loading complete.")
+
+ def accept_registrations(
+ self,
+ pull_socket: zmq.Socket,
+ num_workers: int,
+ ) -> None:
+ """Receive every local DP/TP worker's IPC and shared-memory buffers."""
+ logger.info("Waiting for %d GPU worker registration(s) ...", num_workers)
+ registrations: list[PleOffloadRegistration] = []
+ for index in range(num_workers):
+ item = pickle.loads(pull_socket.recv())
+ if not isinstance(item, PleOffloadRegistration):
+ raise RuntimeError(
+ "Expected PleOffloadRegistration during setup, got "
+ f"{type(item).__name__} ({index + 1}/{num_workers})"
+ )
+ registrations.append(item)
+ logger.info(
+ "GPU worker %d registered (dp_rank=%d, tp_rank=%d, layers=%s).",
+ item.worker_id,
+ item.dp_rank,
+ item.tp_rank,
+ sorted(item.host_output_buffers),
+ )
+
+ dp_size = self.vllm_config.parallel_config.data_parallel_size
+ tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+ if num_workers != dp_size * tp_size:
+ raise RuntimeError(
+ f"Expected {dp_size * tp_size} registrations for DP={dp_size}, "
+ f"TP={tp_size}, got {num_workers}"
+ )
+
+ registrations_by_dp: dict[int, list[PleOffloadRegistration]] = {}
+ for registration in registrations:
+ registrations_by_dp.setdefault(registration.dp_rank, []).append(
+ registration
+ )
+ if set(registrations_by_dp) != set(range(dp_size)):
+ raise RuntimeError(
+ f"Expected DP ranks {set(range(dp_size))}, "
+ f"got {set(registrations_by_dp)}"
+ )
+ for dp_rank, dp_registrations in registrations_by_dp.items():
+ tp_ranks = {registration.tp_rank for registration in dp_registrations}
+ if tp_ranks != set(range(tp_size)):
+ raise RuntimeError(
+ f"DP rank {dp_rank} expected TP ranks {set(range(tp_size))}, "
+ f"got {tp_ranks}"
+ )
+
+ for registration in registrations:
+ if set(registration.host_output_buffers) != set(self.layer_names) or set(
+ registration.host_flag_tensors
+ ) != set(self.layer_names):
+ raise RuntimeError(
+ "Registered PLE layers do not match CPU layers: "
+ f"registered={sorted(registration.host_output_buffers)} / "
+ f"{sorted(registration.host_flag_tensors)}, "
+ f"cpu={sorted(self.layer_names)}"
+ )
+ targets_for_dp = self._worker_targets.setdefault(registration.dp_rank, {})
+ for layer_name, host_buffer in registration.host_output_buffers.items():
+ host_flag = registration.host_flag_tensors[layer_name]
+ for tensor in (host_buffer, host_flag):
+ if tensor.device.type != "cpu" or not tensor.is_shared():
+ raise RuntimeError(
+ f"PLE host buffer for {layer_name} (dp="
+ f"{registration.dp_rank}, tp={registration.tp_rank}) "
+ "is not a shared CPU tensor"
+ )
+ if host_flag.dtype != torch.int32 or host_flag.numel() < 1:
+ raise RuntimeError(f"PLE host flag for {layer_name} must be int32")
+ target = PleOffloadOutputTarget(
+ tp_rank=registration.tp_rank,
+ host_output_buffer=host_buffer,
+ host_flag=host_flag,
+ )
+ targets_for_dp.setdefault(layer_name, []).append(target)
+ # All TP ranks in one DP group receive the same input, so buffers
+ # registered by TP rank zero are sufficient for that DP rank.
+ if registration.tp_rank == 0:
+ self._input_bufs[registration.dp_rank] = PleOffloadInputBuffers(
+ input_ids_buf=registration.input_ids_buf,
+ query_start_loc_buf=registration.query_start_loc_buf,
+ ngram_context_buf=registration.ngram_context_buf,
+ )
+
+ if set(self._input_bufs) != set(range(dp_size)):
+ raise RuntimeError(
+ "TP rank zero did not register PLE input buffers for every DP "
+ f"rank: expected={set(range(dp_size))}, got={set(self._input_bufs)}"
+ )
+
+ config = self.vllm_config.model_config.hf_text_config
+ max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
+ for dp_rank, layer_targets in self._worker_targets.items():
+ self._scratch_bufs[dp_rank] = {}
+ for layer_name, targets in layer_targets.items():
+ if len(targets) != tp_size:
+ raise RuntimeError(
+ f"PLE layer {layer_name} for DP rank {dp_rank} received "
+ f"{len(targets)} targets, expected {tp_size}"
+ )
+ targets.sort(key=lambda target: target.tp_rank)
+ out_dim = self._layers[layer_name].get_offload_output_dim(
+ int(config.ple_embed_dim)
+ )
+ out_dtype = self._layers[layer_name].get_offload_output_dtype(
+ self.vllm_config.model_config.dtype
+ )
+ for target in targets:
+ shape = tuple(target.host_output_buffer.shape)
+ if shape != (max_tokens, out_dim) or (
+ target.host_output_buffer.dtype != out_dtype
+ ):
+ raise RuntimeError(
+ f"PLE host buffer for {layer_name} (dp={dp_rank}, "
+ f"tp={target.tp_rank}) is {shape} "
+ f"{target.host_output_buffer.dtype}, expected "
+ f"{(max_tokens, out_dim)} {out_dtype}"
+ )
+ self._scratch_bufs[dp_rank][layer_name] = torch.empty(
+ max_tokens, out_dim, dtype=out_dtype, device="cpu"
+ )
+ logger.info(
+ "Registrations complete (dp_size=%d, tp_size=%d, layers=%s).",
+ dp_size,
+ tp_size,
+ sorted(self.layer_names),
+ )
+
+ @torch.inference_mode()
+ def busy_loop(
+ self,
+ pull_socket: zmq.Socket,
+ shutdown_event: threading.Event,
+ ) -> None:
+ """Decode and batch available requests by DP rank until shutdown."""
+ logger.info("Busy-loop started.")
+ self._shutdown_event = shutdown_event
+ poller = zmq.Poller()
+ poller.register(pull_socket, zmq.POLLIN)
+ while not shutdown_event.is_set():
+ if pull_socket not in dict(poller.poll(timeout=100)):
+ continue
+
+ requests = []
+ try:
+ requests.append(_receive_offload_request(pull_socket))
+ while True:
+ requests.append(
+ _receive_offload_request(pull_socket, zmq.NOBLOCK)
+ )
+ except zmq.Again:
+ pass
+ except (msgspec.DecodeError, RuntimeError) as error:
+ raise RuntimeError("Unexpected PLE offload request") from error
+
+ self._handle_requests(requests)
+
+ def _handle_requests(
+ self, requests: list[_ReceivedPleOffloadRequest]
+ ) -> None:
+ """Run requests layer-first so each DP rank can resume promptly."""
+ with record_function_or_nullcontext("ple_offload: handle_requests"):
+ self._handle_requests_impl(requests)
+
+ def _validate_requests_before_publication(
+ self, requests: list[PleOffloadRequest]
+ ) -> None:
+ """Validate a drained request batch before computation or publication."""
+ for request in requests:
+ dp_rank = request.dp_rank
+ if dp_rank not in self._worker_targets or dp_rank not in self._input_bufs:
+ raise RuntimeError(
+ f"No PLE buffers or output targets for dp_rank={dp_rank}; "
+ "refusing request"
+ )
+
+ input_bufs = self._input_bufs[dp_rank]
+ if not 0 < request.num_tokens <= input_bufs.input_ids_buf.numel():
+ raise RuntimeError(
+ f"PLE num_tokens={request.num_tokens} is outside capacity "
+ f"{input_bufs.input_ids_buf.numel()} for dp_rank={dp_rank}"
+ )
+ max_reqs = input_bufs.query_start_loc_buf.numel() - 1
+ if not 0 < request.num_reqs <= max_reqs:
+ raise RuntimeError(
+ f"PLE num_reqs={request.num_reqs} is outside capacity "
+ f"{max_reqs} for dp_rank={dp_rank}"
+ )
+ if (
+ input_bufs.query_start_loc_buf.dtype != torch.int32
+ or input_bufs.query_start_loc_buf.dim() != 1
+ ):
+ raise RuntimeError(
+ f"PLE query_start_loc for dp_rank={dp_rank} must be 1-D int32"
+ )
+ if (
+ input_bufs.ngram_context_buf is not None
+ and request.num_reqs > input_bufs.ngram_context_buf.shape[0]
+ ):
+ raise RuntimeError(
+ f"PLE num_reqs={request.num_reqs} exceeds ngram context "
+ f"capacity for dp_rank={dp_rank}"
+ )
+
+ query_start_loc = input_bufs.query_start_loc_buf[
+ : request.num_reqs + 1
+ ].tolist()
+ if (
+ query_start_loc[0] != 0
+ or query_start_loc[-1] > request.num_tokens
+ or any(a > b for a, b in zip(query_start_loc, query_start_loc[1:]))
+ ):
+ raise RuntimeError(
+ f"Invalid PLE query_start_loc for dp_rank={dp_rank}: "
+ f"{query_start_loc}, num_tokens={request.num_tokens}"
+ )
+
+ for layer_name, targets in self._worker_targets[dp_rank].items():
+ scratch = self._scratch_bufs[dp_rank][layer_name]
+ if request.num_tokens > scratch.shape[0] or any(
+ request.num_tokens > target.host_output_buffer.shape[0]
+ for target in targets
+ ):
+ raise RuntimeError(
+ f"PLE num_tokens={request.num_tokens} exceeds output "
+ f"capacity for {layer_name}, dp_rank={dp_rank}"
+ )
+
+ def _handle_requests_impl(
+ self,
+ requests: list[_ReceivedPleOffloadRequest | PleOffloadRequest],
+ ) -> None:
+ """Process one batch of PLE requests."""
+ received = [
+ item
+ if isinstance(item, _ReceivedPleOffloadRequest)
+ else _ReceivedPleOffloadRequest(item, False, None, None)
+ for item in requests
+ ]
+ plain_requests = [item.request for item in received]
+ self._validate_requests_before_publication(plain_requests)
+ received_by_dp: dict[int, _ReceivedPleOffloadRequest] = {}
+ for item in received:
+ request = item.request
+ if request.dp_rank not in self._worker_targets:
+ raise RuntimeError(
+ f"No PLE output targets for dp_rank={request.dp_rank}; "
+ "refusing request"
+ )
+ if request.dp_rank in received_by_dp:
+ # One input slot and one output slot per DP rank: a second
+ # request in the same batch would alias them. Fatal, never
+ # silently dropped.
+ raise RuntimeError(
+ f"Two PLE requests for dp_rank={request.dp_rank} in one batch "
+ f"(step {self._step}); the single-slot protocol is violated"
+ )
+ received_by_dp[request.dp_rank] = item
+ self._step += 1
+ if _FAULT is not None:
+ for item in received_by_dp.values():
+ if not item.serving:
+ continue
+ self._request_index += 1
+ transaction = f"request:{self._request_index}"
+ if item.transaction_id is not None:
+ transaction = f"{transaction}@{item.transaction_id}"
+ for cut in ("park_before_publish", "park_before_publish_ack"):
+ if _FAULT is not None and _FAULT.fires(
+ cut,
+ role="offload",
+ unit="offload-request",
+ eligible=True,
+ dp_rank=item.request.dp_rank,
+ transaction_id=transaction,
+ correlation_id=item.correlation_id,
+ ):
+ # No multiprocessing or fault-state lock is held here.
+ threading.Event().wait()
+
+ requests_by_dp = {
+ dp_rank: item.request
+ for dp_rank, item in received_by_dp.items()
+ }
+
+ # Speculative placeholders are not vocabulary IDs. Normalize each DP
+ # input once before all PLE layers consume the shared buffer.
+ if self._clamp_input_ids:
+ for dp_rank, request in requests_by_dp.items():
+ self._input_bufs[dp_rank].input_ids_buf[
+ : request.num_tokens
+ ].clamp_min_(0)
+
+ for layer_name, layer in self._layers.items():
+ for dp_rank, request in requests_by_dp.items():
+ targets = self._worker_targets[dp_rank][layer_name]
+
+ input_bufs = self._input_bufs[dp_rank]
+ ngram_context = (
+ input_bufs.ngram_context_buf[: request.num_reqs]
+ if input_bufs.ngram_context_buf is not None
+ else None
+ )
+ result = layer.forward_impl(
+ input_bufs.input_ids_buf[: request.num_tokens],
+ input_bufs.input_ids_buf[: request.num_tokens],
+ input_bufs.query_start_loc_buf[: request.num_reqs + 1],
+ ngram_context,
+ output_buffer=self._scratch_bufs[dp_rank][layer_name],
+ )
+ if _rowdump.enabled():
+ _seq = getattr(self, "_rowdump_seq", 0)
+ self._rowdump_seq = _seq + 1
+ _rowdump.write(f"offload-dp{dp_rank}", 1, dp_rank, _seq,
+ input_bufs.input_ids_buf[: request.num_tokens], result)
+
+ # The result is identical on every TP rank in this DP group.
+ # A worker's host buffer may be overwritten only after its GPU
+ # released the previous result (``release_outputs`` stores
+ # flag=0 on the model stream after the forward, which is also
+ # after the in-graph DMA read it). Then write rows, then flag=1.
+ scratch = self._scratch_bufs[dp_rank][layer_name]
+ if result.dim() != 2 or result.shape[0] != request.num_tokens or (
+ result.shape[1] != scratch.shape[1]
+ ) or result.dtype != scratch.dtype or result.device.type != "cpu":
+ raise RuntimeError(
+ f"PLE result {tuple(result.shape)} {result.dtype} "
+ f"{result.device} for {layer_name} does not "
+ f"match the request ({request.num_tokens} rows x "
+ f"{scratch.shape[1]} {scratch.dtype} CPU); refusing to "
+ f"publish (step {self._step})"
+ )
+ slices = tuple(slice(0, size) for size in result.shape)
+ for target in targets:
+ self._wait_flag_released(target, layer_name, dp_rank)
+ target.host_output_buffer[slices].copy_(result[slices])
+ _publish_fence()
+ target.host_flag[0] = 1
+
+ def _wait_flag_released(
+ self, target: PleOffloadOutputTarget, layer_name: str, dp_rank: int
+ ) -> None:
+ """Spin until the GPU reset this target's flag (previous result used)."""
+ flag = target.host_flag
+ if int(flag[0]) == 0:
+ return
+ spins = 0
+ warned = False
+ start = time.monotonic()
+ while int(flag[0]) != 0:
+ spins += 1
+ if self._shutdown_event is not None and self._shutdown_event.is_set():
+ raise RuntimeError("PLE offload shutting down while waiting for GPU")
+ if spins > 200:
+ time.sleep(0.0001)
+ waited = time.monotonic() - start
+ if not warned and waited > 5.0:
+ warned = True
+ logger.warning(
+ "PLE offload waited >5 s for the GPU to release %s "
+ "(dp=%d, tp=%d, step %d); flag=%d",
+ layer_name,
+ dp_rank,
+ target.tp_rank,
+ self._step,
+ int(flag[0]),
+ )
+ if waited > _PLE_WAIT_TIMEOUT_S:
+ raise RuntimeError(
+ f"PLE offload: GPU did not release {layer_name} (dp={dp_rank}, "
+ f"tp={target.tp_rank}, step {self._step}) within "
+ f"{_PLE_WAIT_TIMEOUT_S:.0f} s; refusing to overwrite"
+ )
diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py
index 9657c28..0e132b6 100644
--- a/vllm/v1/worker/gpu/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -1,9 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
import gc
import itertools
from collections import defaultdict
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Iterator
from contextlib import ExitStack
from dataclasses import dataclass, replace
from itertools import groupby, product
@@ -30,6 +31,9 @@ from vllm.distributed.parallel_state import (
)
from vllm.forward_context import BatchDescriptor, set_forward_context
from vllm.logger import init_logger
+from vllm.model_executor.layers.ple_offload_layer import (
+ observe_ple_copy_geometry,
+)
from vllm.model_executor.offloader.base import get_offloader
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
@@ -170,6 +174,14 @@ class CudaGraphManager:
self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map()
self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}
+ # FULL graphs bake every PLE wait/copy into their native graph. Metadata
+ # comes from the actual custom-op implementation during native capture.
+ self._ple_full_copy_geometry: dict[
+ BatchExecutionDescriptor, tuple[tuple[str, int, int], ...]
+ ] = {}
+ self._expected_ple_copy_layers: tuple[str, ...] = ()
+ # Only the target-model manager is allowed to capture PLE operations.
+ self._allow_ple_full_graph = False
self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None
self._graphs_captured = False
@@ -469,15 +481,16 @@ class CudaGraphManager:
if self._capture_mem_samples is not None:
torch.accelerator.synchronize()
free_before = torch.accelerator.get_memory_info()[0]
- with torch.cuda.graph(
- graph, self.pool, stream=self._capture_stream(desc)
- ):
- forward_fn(CUDAGraphMode.NONE)
- # Join offloader's copy stream after forward to avoid
- # unjoined stream error. The last layer's start_prefetch
- # forks copy_stream, but wait_prefetch only happens in
- # the next forward pass.
- get_offloader().join_after_forward()
+ with self.capture_ple_copy_geometry(desc):
+ with torch.cuda.graph(
+ graph, self.pool, stream=self._capture_stream(desc)
+ ):
+ forward_fn(CUDAGraphMode.NONE)
+ # Join offloader's copy stream after forward to avoid
+ # unjoined stream error. The last layer's start_prefetch
+ # forks copy_stream, but wait_prefetch only happens in
+ # the next forward pass.
+ get_offloader().join_after_forward()
if self._capture_mem_samples is not None:
torch.accelerator.synchronize()
free_after = torch.accelerator.get_memory_info()[0]
@@ -493,6 +506,78 @@ class CudaGraphManager:
{desc.num_tokens for desc in self.graphs if desc.num_active_loras == 0}
)
+ @contextlib.contextmanager
+ def capture_ple_copy_geometry(
+ self, desc: BatchExecutionDescriptor
+ ) -> Iterator[None]:
+ """Production capture callback for actual PLE custom-op geometry."""
+ with observe_ple_copy_geometry() as geometry:
+ yield
+
+ if geometry and not self._allow_ple_full_graph:
+ raise RuntimeError(
+ "PLE wait-copy reached an unsupported FULL graph path; "
+ "speculator/draft FULL graphs may not contain PLE operations"
+ )
+ if geometry:
+ self.record_ple_copy_geometry(desc, geometry)
+ elif self._expected_ple_copy_layers:
+ raise RuntimeError(
+ f"PLE copy geometry was not observed while capturing FULL graph {desc}"
+ )
+
+ def record_ple_copy_geometry(
+ self,
+ desc: BatchExecutionDescriptor,
+ geometry: Iterable[tuple[str, int, int]],
+ ) -> None:
+ if desc.cg_mode != CUDAGraphMode.FULL:
+ raise ValueError("PLE capture geometry is recorded only for FULL graphs")
+ observed = tuple(
+ (str(layer), int(rows), int(nbytes))
+ for layer, rows, nbytes in geometry
+ )
+ if not observed:
+ raise RuntimeError(f"PLE FULL graph {desc} observed no wait-copy operations")
+ observed_layers = tuple(layer for layer, _, _ in observed)
+ expected_layers = self._expected_ple_copy_layers
+ if expected_layers and sorted(observed_layers) != sorted(expected_layers):
+ raise RuntimeError(
+ f"PLE FULL graph {desc} observed layers {observed_layers}, "
+ f"expected {expected_layers}"
+ )
+ inconsistent = [
+ (layer, rows, nbytes)
+ for layer, rows, nbytes in observed
+ if rows != desc.num_tokens or nbytes < 0
+ ]
+ if inconsistent:
+ raise RuntimeError(
+ f"PLE FULL graph {desc} actual PLE copy geometry "
+ f"{inconsistent} does not match descriptor rows={desc.num_tokens}"
+ )
+ previous = self._ple_full_copy_geometry.get(desc)
+ if previous is not None and previous != observed:
+ raise RuntimeError(
+ f"PLE FULL graph {desc} changed copy geometry from "
+ f"{previous} to {observed}"
+ )
+ self._ple_full_copy_geometry[desc] = observed
+
+ def ple_copy_rows(self, desc: BatchExecutionDescriptor) -> int:
+ try:
+ geometry = self._ple_full_copy_geometry[desc]
+ except KeyError as error:
+ raise RuntimeError(
+ f"PLE copy geometry was not recorded for FULL graph {desc}"
+ ) from error
+ rows = {num_rows for _, num_rows, _ in geometry}
+ if len(rows) != 1:
+ raise RuntimeError(
+ f"PLE FULL graph {desc} has inconsistent per-layer rows: {geometry}"
+ )
+ return rows.pop()
+
def dispatch(
self,
num_reqs: int,
@@ -532,6 +617,11 @@ class CudaGraphManager:
f"Expected FULL mode, got {desc.cg_mode}"
)
assert desc in self.graphs, f"No cudagraph for {desc}"
+ if (
+ desc in self._ple_full_copy_geometry
+ and not self._allow_ple_full_graph
+ ):
+ raise RuntimeError("refusing unsupported PLE speculator FULL replay")
# Sync offloader before replay - needed when transitioning from
# eager/piecewise to full cudagraph (e.g., prefill → decode).
# The previous eager iteration's start_prefetch may have queued
@@ -577,11 +667,32 @@ class ModelCudaGraphManager(CudaGraphManager):
varlen_decode=varlen_decode,
ubatch_runner=ubatch_runner,
)
+ self._allow_ple_full_graph = True
self.hidden_states: torch.Tensor | None = None
self.aux_hidden_states: list[torch.Tensor] = []
self.use_aux_hidden_state_outputs = False
self.intermediate_tensors: IntermediateTensors | None = None
+ def _prepare_ple_capture(
+ self, model: nn.Module, *, initialize_breakable: bool
+ ) -> None:
+ expected_ple_copy_layers = tuple(
+ str(getattr(module, "_ple_offload_name", name))
+ for name, module in model.named_modules()
+ if getattr(module, "_is_cpu_offloaded", False)
+ and getattr(module, "_host_output_ptr", None) is not None
+ )
+ if (
+ "<unbound>" in expected_ple_copy_layers
+ or len(set(expected_ple_copy_layers)) != len(expected_ple_copy_layers)
+ ):
+ raise RuntimeError(
+ "PLE FULL capture requires distinct, bound connector layer identities"
+ )
+ self._expected_ple_copy_layers = expected_ple_copy_layers
+ if initialize_breakable:
+ self.init_breakable_cg_runner(model)
+
def capture(
self,
model: nn.Module,
@@ -599,8 +710,9 @@ class ModelCudaGraphManager(CudaGraphManager):
) -> None:
"""Capture CUDA graphs for model forward pass."""
self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
- if self.use_breakable_cg:
- self.init_breakable_cg_runner(model)
+ self._prepare_ple_capture(
+ model, initialize_breakable=self.use_breakable_cg
+ )
if self.cudagraph_mode.has_piecewise_cudagraphs() and not (
self.use_breakable_cg or has_compiled_submodule(model)
@@ -901,7 +1013,9 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:
if manager.use_breakable_cg:
# Create the breakable runner before the wrapper pool swap so
# its pool is covered as well.
- manager.init_breakable_cg_runner(runner.model)
+ manager._prepare_ple_capture(
+ runner.model, initialize_breakable=True
+ )
all_wrappers = list(CUDAGraphWrapper._all_instances) + list(
BreakableCUDAGraphWrapper._all_instances
)
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index 744fe5b..e2ddc52 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -77,6 +77,9 @@ from vllm.v1.outputs import (
ModelRunnerOutput,
RoutedExpertsTensors,
)
+from vllm.v1.ple_offload.connector import PleOffloadConnector
+from vllm.v1.ple_offload.fatal import fail_ple_execute_model
+from vllm.v1.ple_offload.test_fault import _FAULT
from vllm.v1.watermarking import create_watermarker
from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
from vllm.v1.watermarking.spec_decode import (
@@ -1689,6 +1692,45 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.ec_connector.no_forward(scheduler_output).ec_connector_output,
)
+ def _prepare_ple_forward(
+ self,
+ input_batch: InputBatch,
+ batch_desc: BatchExecutionDescriptor,
+ dummy_run: bool,
+ *,
+ serving: bool = False,
+ transaction_id: str | None = None,
+ correlation_id: str | None = None,
+ ) -> None:
+ connector = self._ple_offload_connector
+ assert connector is not None
+ request_num_tokens = int(input_batch.num_tokens_after_padding)
+ if batch_desc.cg_mode == CUDAGraphMode.FULL:
+ manager = self.cudagraph_manager
+ if manager is None:
+ raise RuntimeError("PLE FULL replay has no CUDA graph manager")
+ copy_num_tokens = manager.ple_copy_rows(batch_desc)
+ else:
+ # PIECEWISE/eager execute the Python op with this actual padded view.
+ copy_num_tokens = int(input_batch.input_ids.shape[0])
+
+ if request_num_tokens != copy_num_tokens:
+ raise RuntimeError(
+ "PLE request/selected-forward copy-count mismatch: "
+ f"request={request_num_tokens}, copy={copy_num_tokens}, "
+ f"mode={batch_desc.cg_mode}"
+ )
+ connector.prepare_forward(
+ input_batch.num_reqs,
+ request_num_tokens,
+ dummy_run,
+ consumer_num_tokens=copy_num_tokens,
+ serving=serving,
+ transaction_id=transaction_id,
+ correlation_id=correlation_id,
+ )
+
+ @fail_ple_execute_model
@torch.inference_mode()
def execute_model(
self,
@@ -1956,6 +1998,54 @@ class GPUModelRunner(LoRAModelRunnerMixin):
)
self.step_timing.forward_start()
+ # prepare_inputs has finalized MRV2's GPU buffers. Stage D2H on this
+ # stream before the PLE placeholder can wait for CPU output.
+ connector = self._ple_offload_connector
+ ple_serving = False
+ ple_sequence = None
+ ple_transaction = None
+ ple_correlation = None
+ if connector is not None:
+ ple_serving = not dummy_run and not is_profile
+ ple_sequence = connector.begin_forward(
+ dummy_run=dummy_run,
+ num_tokens=int(input_batch.num_tokens_after_padding),
+ )
+ if _FAULT is not None and ple_serving:
+ ple_correlation = f"sequence:{ple_sequence}"
+ ple_transaction = connector.fault_transaction(ple_sequence)
+ if _FAULT is not None and _FAULT.fires(
+ "exc_before_prepare",
+ role="gpu",
+ unit="runner-forward",
+ eligible=ple_serving,
+ tp_rank=connector.tp_rank,
+ pp_rank=connector.pp_rank,
+ dp_rank=connector.dp_rank,
+ transaction_id=ple_transaction,
+ correlation_id=ple_correlation,
+ ):
+ raise RuntimeError("PLE test fault: exc_before_prepare")
+ self._prepare_ple_forward(
+ input_batch,
+ batch_desc,
+ dummy_run,
+ serving=ple_serving,
+ transaction_id=ple_transaction,
+ correlation_id=ple_correlation,
+ )
+ if _FAULT is not None and _FAULT.fires(
+ "exc_after_prepare",
+ role="gpu",
+ unit="runner-forward",
+ eligible=ple_serving,
+ tp_rank=connector.tp_rank,
+ pp_rank=connector.pp_rank,
+ dp_rank=connector.dp_rank,
+ transaction_id=ple_transaction,
+ correlation_id=ple_correlation,
+ ):
+ raise RuntimeError("PLE test fault: exc_after_prepare")
connector_kwargs = dict(
scheduler_output=scheduler_output,
request_state_indices=input_batch.idx_mapping,
@@ -2013,6 +2103,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
# Eager (NONE): call the raw model directly.
model_output = self.model(**model_inputs)
+ # The model has consumed every PLE output, so the CPU worker may reuse
+ # the registered buffers for the next request.
+ if connector is not None:
+ connector.release_outputs(
+ serving=ple_serving,
+ transaction_id=ple_transaction,
+ correlation_id=ple_correlation,
+ )
+ if _FAULT is not None and _FAULT.fires(
+ "exc_after_release",
+ role="gpu",
+ unit="runner-forward",
+ eligible=ple_serving,
+ tp_rank=connector.tp_rank,
+ pp_rank=connector.pp_rank,
+ dp_rank=connector.dp_rank,
+ transaction_id=ple_transaction,
+ correlation_id=ple_correlation,
+ ):
+ raise RuntimeError("PLE test fault: exc_after_release")
self.kv_connector.finish_forward()
if self.is_last_pp_rank:
@@ -2047,6 +2157,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
ec_connector_output=ec_connector_output,
routed_experts=routed_experts,
cudagraph_stats=cudagraph_stats,
+ ple_serving=ple_serving,
+ ple_sequence=ple_sequence,
+ ple_transaction=ple_transaction,
)
if not self.is_last_pp_rank:
@@ -2058,6 +2171,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
)
return None
+ @fail_ple_execute_model
@torch.inference_mode()
@step_eplb_after()
def sample_tokens(
@@ -2077,6 +2191,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
ec_connector_output = self.execute_model_state.ec_connector_output
routed_experts = self.execute_model_state.routed_experts
cudagraph_stats = self.execute_model_state.cudagraph_stats
+ ple_serving = self.execute_model_state.ple_serving
+ ple_sequence = self.execute_model_state.ple_sequence
+ ple_transaction = self.execute_model_state.ple_transaction
self.execute_model_state = None
if not self.is_last_pp_rank:
@@ -2084,6 +2201,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
# IntermediateTensors instead of final hidden states. Receive the
# sampled tokens broadcast from the last rank and update local state.
assert self.pp_handler is not None
+ connector = self._ple_offload_connector
+ if (
+ connector is not None
+ and _FAULT is not None
+ and ple_serving
+ and ple_sequence is not None
+ and connector.completion_is_pending(ple_sequence)
+ and _FAULT.fires(
+ "exc_async_output",
+ role="gpu",
+ unit="runner-output-relay",
+ eligible=True,
+ tp_rank=connector.tp_rank,
+ pp_rank=connector.pp_rank,
+ dp_rank=connector.dp_rank,
+ transaction_id=ple_transaction,
+ correlation_id=f"sequence:{ple_sequence}",
+ )
+ ):
+ raise RuntimeError("PLE test fault: exc_async_output")
all_decode_next = self.pp_handler.receive(input_batch)
# Optimistically update num_computed_tokens for entire batch here.
# Will be adjusted for rejections if necessary in update_requests.
@@ -2379,6 +2516,9 @@ class ExecuteModelState(NamedTuple):
ec_connector_output: ECConnectorOutput | None
routed_experts: RoutedExpertsTensors | None
cudagraph_stats: CUDAGraphStat | None
+ ple_serving: bool
+ ple_sequence: int | None
+ ple_transaction: str | None
class BatchReqState(NamedTuple):
diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
index 09f91f9..43bc705 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -87,6 +87,7 @@ from vllm.v1.outputs import (
DraftTokenIds,
ModelRunnerOutput,
)
+from vllm.v1.ple_offload.fatal import fail_ple_execute_model
from vllm.v1.utils import compute_iteration_details, report_usage_stats
from vllm.v1.worker import guard_telemetry
from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel
@@ -231,6 +232,148 @@ class Worker(WorkerBase):
# Resolved lazily on first sleep/wake; persists worker-process state.
self._sleep_mode_backend: SleepModeBackend | None = None
+ def _has_ple_layers(self) -> bool:
+ """Return whether this model configuration constructs PLE layers."""
+ if not envs.VLLM_PLE_CPU_OFFLOAD:
+ return False
+ text_config = self.model_config.hf_text_config
+ ple_layer_ids = getattr(text_config, "ple_layer_ids", None)
+ if not ple_layer_ids:
+ return False
+ # PP (vllm#54709, carried from the earlier overlay): only the pipeline rank
+ # that owns a PLE layer takes part in the offload (rank 0 with the
+ # default partition). Other ranks build PPMissingLayer there and must
+ # skip the connector entirely.
+ pp_size = self.parallel_config.pipeline_parallel_size
+ if pp_size > 1:
+ from vllm.distributed.utils import get_pp_indices
+
+ tp_size = self.parallel_config.tensor_parallel_size
+ pp_rank = (self.rank // tp_size) % pp_size
+ start, end = get_pp_indices(
+ int(text_config.num_hidden_layers), pp_rank, pp_size
+ )
+ return any(start <= int(i) - 1 < end for i in ple_layer_ids)
+ return True
+
+ def _validate_ple_offload_config(self) -> None:
+ """Reject unsupported PLE offload execution modes."""
+ parallel_config = self.parallel_config
+ unsupported = []
+ if not self.use_v2_model_runner:
+ unsupported.append("model runner V1")
+ if not current_platform.is_cuda():
+ unsupported.append(f"device={current_platform.device_type}")
+ if parallel_config.nnodes != 1:
+ unsupported.append(f"nnodes={parallel_config.nnodes}")
+ if parallel_config.data_parallel_backend != "mp":
+ unsupported.append(f"DP backend={parallel_config.data_parallel_backend}")
+ if (
+ parallel_config.data_parallel_size_local
+ != parallel_config.data_parallel_size
+ ):
+ unsupported.append(
+ "non-local DP "
+ f"({parallel_config.data_parallel_size_local}/"
+ f"{parallel_config.data_parallel_size} local ranks)"
+ )
+ # PP allowed (vllm#54709): PLE layers must live on pipeline rank 0
+ # (enforced placement-based in models/config.py); num_workers stays
+ # dp_size * tp_size = the ranks of the owning stage.
+ if parallel_config.prefill_context_parallel_size != 1:
+ unsupported.append(f"PCP={parallel_config.prefill_context_parallel_size}")
+ if parallel_config.decode_context_parallel_size != 1:
+ unsupported.append(f"DCP={parallel_config.decode_context_parallel_size}")
+ if parallel_config.use_ubatching:
+ unsupported.append("ubatching/DBO")
+ if self.model_config.architecture not in {
+ "Qwen4ExpForCausalLM",
+ "Qwen4ExpForConditionalGeneration",
+ }:
+ unsupported.append("architecture")
+ if self.vllm_config.weight_transfer_config is not None:
+ unsupported.append("weight transfer")
+
+ if unsupported:
+ raise ValueError(
+ "VLLM_PLE_CPU_OFFLOAD does not support the requested "
+ "configuration. Unsupported settings: "
+ f"{', '.join(unsupported)}"
+ )
+
+ def spawn_ple_offload(self) -> None:
+ """Spawn one node-local PLE CPU worker from DP0/TP0."""
+ if (
+ not self._ple_offload_enabled
+ or self.rank != 0
+ or self.parallel_config.data_parallel_rank != 0
+ ):
+ return
+
+ from vllm.v1.ple_offload.worker import PleOffloadWorker
+
+ ipc_addr = self.parallel_config._ple_offload_ipc_path
+ if not ipc_addr:
+ raise RuntimeError("PLE offload IPC address was not initialized")
+ dp_size = self.parallel_config.data_parallel_size
+ tp_size = self.parallel_config.tensor_parallel_size
+ num_workers = dp_size * tp_size
+ logger.info(
+ "PleOffload: spawning worker "
+ "(rank=%d, local_rank=%d, dp_size=%d, tp_size=%d, "
+ "num_workers=%d, ipc_addr=%s).",
+ self.rank,
+ self.local_rank,
+ dp_size,
+ tp_size,
+ num_workers,
+ ipc_addr,
+ )
+ self._ple_offload_worker_handle = PleOffloadWorker.make_process(
+ self.vllm_config,
+ num_workers,
+ ipc_addr,
+ )
+ self._start_ple_offload_watchdog()
+
+ def _start_ple_offload_watchdog(self) -> None:
+ """Fail this worker if the PLE offload process exits unexpectedly.
+
+ Host-mapped pull: a dead producer would leave the captured forward
+ waiting on its host flag forever; turn that into a fast, visible
+ engine failure instead of a hang. ``shutdown`` sets the stop marker
+ before closing the handle, so a normal close is not a failure.
+ """
+ import threading
+
+ handle = self._ple_offload_worker_handle
+ self._ple_offload_stopping = False
+
+ def watch() -> None:
+ returncode = handle.proc.wait()
+ if getattr(self, "_ple_offload_stopping", False):
+ return
+ logger.error(
+ "PLE offload process %s exited unexpectedly (exitcode %s); "
+ "failing this worker so the engine does not hang on a PLE wait",
+ handle.proc.pid,
+ returncode,
+ )
+ os._exit(1)
+
+ threading.Thread(
+ target=watch, name="ple-offload-watchdog", daemon=True
+ ).start()
+
+ def wait_ple_offload_ready(self) -> None:
+ """Wait until weights and every TP registration are ready."""
+ if self._ple_offload_worker_handle is None:
+ return
+
+ from vllm.v1.ple_offload.worker import PleOffloadWorker
+
+ PleOffloadWorker.wait_for_ready(self._ple_offload_worker_handle)
+
@property
def sleep_mode_backend(self) -> "SleepModeBackend":
if self._sleep_mode_backend is None:
@@ -1152,6 +1295,7 @@ class Worker(WorkerBase):
) -> ModelRunnerOutput | AsyncModelRunnerOutput:
return self.model_runner.sample_tokens(grammar_output)
+ @fail_ple_execute_model
@torch.inference_mode()
@with_gpu_sync_check
def execute_model(
@@ -1487,6 +1631,7 @@ class Worker(WorkerBase):
self.model_runner.reset_lora_state()
def shutdown(self) -> None:
+ self._ple_offload_stopping = True
gc.unfreeze()
# Bounds-guard instrumentation: synchronous final read of the guard counters, so
diff --git a/vllm_ple_offload_bootstrap/__init__.py b/vllm_ple_offload_bootstrap/__init__.py
new file mode 100644
index 0000000..60ddda4
--- /dev/null
+++ b/vllm_ple_offload_bootstrap/__init__.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Minimal PLE child bootstrap; intentionally imports no torch or vLLM."""
+
+import importlib.abc
+import importlib.machinery
+import os
+import pickle
+import struct
+import sys
+import threading
+from multiprocessing.connection import Connection
+from types import ModuleType
+from typing import Any
+
+_FAULT = None
+if os.environ.get("VLLM_PLE_TEST_FAULT") is not None:
+ from vllm_ple_test_fault import (
+ _FAULT as _FAULT,
+ _block_on_harness_gate,
+ )
+
+
+def _send_partial_ready_frame(fd: int) -> tuple[int, int]:
+ """Write the complete header and an incomplete body in one syscall."""
+ body = pickle.dumps(
+ {"status": "READY", "layer_names": []},
+ protocol=pickle.HIGHEST_PROTOCOL,
+ )
+ frame = struct.pack("!i", len(body)) + body
+ fragment = frame[: 4 + max(1, len(body) // 2)]
+ try:
+ return os.write(fd, fragment), len(frame)
+ except OSError:
+ return 0, len(frame)
+
+
+def _send_partial_ready_header(fd: int) -> tuple[int, int]:
+ """Write fewer than the four required header bytes in one syscall."""
+ header = struct.pack("!i", 1)
+ try:
+ return os.write(fd, header[:2]), len(header)
+ except OSError:
+ return 0, len(header)
+
+
+def _read_exact(fd: int, size: int) -> bytes:
+ chunks = bytearray()
+ while len(chunks) < size:
+ chunk = os.read(fd, size - len(chunks))
+ if not chunk:
+ raise EOFError("PLE launch payload was truncated")
+ chunks.extend(chunk)
+ return bytes(chunks)
+
+
+def _recv_launch_payload(fd: int) -> dict:
+ size = struct.unpack("!Q", _read_exact(fd, 8))[0]
+ payload = pickle.loads(_read_exact(fd, size))
+ if not isinstance(payload, dict):
+ raise RuntimeError("invalid PLE bootstrap payload")
+ return payload
+
+
+def _fires(
+ cut: str,
+ *,
+ unit: str = "bootstrap-launch",
+ transaction: str = "launch:1",
+ bytes_written: int | None = None,
+ frame_size: int | None = None,
+) -> bool:
+ return _FAULT is not None and _FAULT.fires(
+ cut,
+ role="bootstrap",
+ unit=unit,
+ eligible=True,
+ transaction_id=transaction,
+ correlation_id="launch:1",
+ bytes_written=bytes_written,
+ frame_size=frame_size,
+ )
+
+
+def _write_and_ack_partial(
+ fd: int,
+ cut: str,
+ sender: Any,
+) -> bool:
+ """A short or failed prefix write cannot acknowledge or enter the park."""
+ if _FAULT is None or _FAULT.cut != cut:
+ # Disarmed or armed for another cut: the control stream is untouched.
+ return False
+ written, frame_size = sender(fd)
+ if written <= 0 or written >= frame_size:
+ return False
+ expected = 2 if cut == "launch_partial_ready_header" else None
+ if expected is not None and written != expected:
+ return False
+ if expected is None and written <= 4:
+ return False
+ return _fires(
+ cut,
+ bytes_written=written,
+ frame_size=frame_size,
+ )
+
+
+class _FaultImportLoader(importlib.abc.Loader):
+ def __init__(self, delegate: importlib.abc.Loader) -> None:
+ self._delegate = delegate
+
+ def create_module(self, spec):
+ create = getattr(self._delegate, "create_module", None)
+ return None if create is None else create(spec)
+
+ def exec_module(self, module: ModuleType) -> None:
+ if _fires(
+ "launch_block_worker_import",
+ unit="bootstrap-import",
+ transaction="module:vllm.v1.ple_offload.worker",
+ ):
+ _block_on_harness_gate()
+ self._delegate.exec_module(module)
+
+
+class _FaultImportFinder(importlib.abc.MetaPathFinder):
+ _TARGET = "vllm.v1.ple_offload.worker"
+
+ def find_spec(self, fullname: str, path, target=None):
+ if fullname != self._TARGET:
+ return None
+ spec = importlib.machinery.PathFinder.find_spec(fullname, path, target)
+ if spec is None or spec.loader is None:
+ raise ImportError(f"cannot find {fullname}")
+ if not isinstance(spec.loader, importlib.abc.Loader):
+ raise ImportError(f"{fullname} has no executable loader")
+ spec.loader = _FaultImportLoader(spec.loader)
+ return spec
+
+
+def _install_import_gate() -> None:
+ target = _FaultImportFinder._TARGET
+ if target in sys.modules:
+ raise RuntimeError("PLE worker import gate was a cache hit")
+ sys.meta_path.insert(0, _FaultImportFinder())
+
+
+def main() -> None:
+ """Receive launch state before importing torch or vLLM."""
+ if len(sys.argv) != 3:
+ raise SystemExit("expected inherited PLE control and parent-death fds")
+ control_fd = int(sys.argv[1])
+ death_fd = int(sys.argv[2])
+ if control_fd == death_fd:
+ raise RuntimeError("PLE control and parent-death fds must be distinct")
+
+ os.set_inheritable(control_fd, False)
+ os.set_inheritable(death_fd, False)
+ death_pipe = Connection(death_fd)
+
+ def monitor_parent() -> None:
+ try:
+ death_pipe.recv()
+ except (EOFError, OSError):
+ pass
+ os._exit(1)
+
+ threading.Thread(
+ target=monitor_parent,
+ daemon=True,
+ name="PleOffloadBootstrapDeathMonitor",
+ ).start()
+
+ if _fires("launch_park_before_payload_read"):
+ threading.Event().wait()
+
+ # An armed reconstruction payload invokes its reducer inside this loads().
+ payload = _recv_launch_payload(control_fd)
+
+ if _write_and_ack_partial(
+ control_fd,
+ "launch_partial_ready",
+ _send_partial_ready_frame,
+ ):
+ threading.Event().wait()
+
+ if _write_and_ack_partial(
+ control_fd,
+ "launch_partial_ready_header",
+ _send_partial_ready_header,
+ ):
+ threading.Event().wait()
+
+ if _fires("launch_silent_ready"):
+ threading.Event().wait()
+
+ ready_pipe = Connection(control_fd)
+ vllm_config = pickle.loads(payload["vllm_config_payload"])
+
+ if _FAULT is not None and _FAULT.cut == "launch_block_worker_import":
+ _install_import_gate()
+ from vllm.v1.ple_offload.worker import PleOffloadWorker
+
+ PleOffloadWorker.proc_main(
+ vllm_config=vllm_config,
+ num_workers=int(payload["num_workers"]),
+ ipc_addr=str(payload["ipc_addr"]),
+ ready_pipe=ready_pipe,
+ )
diff --git a/vllm_ple_offload_bootstrap/__main__.py b/vllm_ple_offload_bootstrap/__main__.py
new file mode 100644
index 0000000..5b3ca57
--- /dev/null
+++ b/vllm_ple_offload_bootstrap/__main__.py
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fresh-interpreter entry point for the PLE CPU-offload worker."""
+
+from vllm_ple_offload_bootstrap import main
+
+main()
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink