Parallel stage initialization (admission + SH/EX device locks) (Worker refactor 1/N) (#5224)

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -205,6 +205,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
     engine._omni_lb_policy = "random"
     engine.request_queue = types.SimpleNamespace()
     engine._log_stats = True
+    engine._parallel_stage_init = False
 
     captured: dict[str, object] = {}
     runtime = types.SimpleNamespace(stage_pools=[], initialize=lambda: None)

tests/engine/test_parallel_stage_init.py

@@ -0,0 +1,842 @@
+"""Unit tests for parallel stage initialization (CPU-only).
+
+Covers the pure/lock logic added for ``parallel_stage_init``:
+  * init-group device keying (resolved physical set; parallel unique keys),
+  * admission control arithmetic + graph reserve,
+  * SH/EX device phase locks (real flock on a temp dir),
+  * the phase-locked executor wrapper ordering,
+  * a regression guard that the NVML process-scope path stayed removed.
+
+The GPU-only same-device concurrency (actual overlapping init) requires a
+hardware soak and is out of scope for these unit tests.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import inspect
+import os
+import subprocess
+import threading
+import types
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.engine.stage_admission import (
+    ADMISSION_EXEMPT,
+    DeviceLedger,
+    StageAdmissionError,
+    StageDemand,
+    check_admission,
+    evaluate,
+    graph_reserve_bytes,
+)
+from vllm_omni.engine.stage_engine_core_proc import _install_phase_locks
+from vllm_omni.engine.stage_init_utils import (
+    LogicalStageInitPlan,
+    ReplicaInitPlan,
+    device_init_lock_path,
+    open_device_lock_file,
+)
+from vllm_omni.engine.stage_phase_lock import (
+    DeviceLockTimeoutError,
+    DevicePhaseLock,
+    resolve_driver_device_ids,
+    wrap_executor_with_phase_locks,
+)
+from vllm_omni.engine.stage_runtime import StageRuntime
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# --------------------------------------------------------------------------- #
+# helpers
+# --------------------------------------------------------------------------- #
+def _llm_replica(stage_id: int, replica_id: int, devices: str | None, *, vllm_config=None) -> ReplicaInitPlan:
+    metadata = types.SimpleNamespace(
+        stage_id=stage_id,
+        stage_type="llm",
+        runtime_cfg={"devices": devices} if devices is not None else {},
+        replica_id=replica_id,
+    )
+    stage_cfg = types.SimpleNamespace(
+        stage_id=stage_id,
+        stage_type="llm",
+        runtime=types.SimpleNamespace(devices=devices),
+        engine_args={},
+    )
+    return ReplicaInitPlan(
+        replica_id=replica_id,
+        num_replicas=1,
+        launch_mode="local",
+        stage_cfg=stage_cfg,
+        metadata=metadata,
+        stage_connector_spec={},
+        omni_kv_connector=(None, None, None),
+        stage_vllm_config=vllm_config,
+        executor_class=object,
+    )
+
+
+def _runtime(parallel_stage_init: bool = False) -> StageRuntime:
+    return StageRuntime(
+        stage_configs=[],
+        model="dummy",
+        config_path="dummy",
+        stage_init_timeout=5,
+        diffusion_batch_size=1,
+        async_chunk=False,
+        parallel_stage_init=parallel_stage_init,
+    )
+
+
+def _fake_vllm_config(util: float, *, model="m", dtype="bf16", cudagraph_mode="FULL", capture_sizes=(1, 2, 4)):
+    return types.SimpleNamespace(
+        cache_config=types.SimpleNamespace(gpu_memory_utilization=util),
+        model_config=types.SimpleNamespace(model=model, dtype=dtype, enforce_eager=False),
+        compilation_config=types.SimpleNamespace(
+            cudagraph_mode=cudagraph_mode,
+            cudagraph_capture_sizes=list(capture_sizes),
+        ),
+    )
+
+
+# --------------------------------------------------------------------------- #
+# A1: init-group device key
+# --------------------------------------------------------------------------- #
+def test_init_group_key_resolves_physical_devices_serial():
+    runtime = _runtime(parallel_stage_init=False)
+    runtime._init_visible_devices_baseline = None  # no CUDA_VISIBLE_DEVICES mapping
+    r0 = _llm_replica(0, 0, "0")
+    r1 = _llm_replica(1, 0, "1")
+    r0b = _llm_replica(2, 0, "0")
+
+    k0 = runtime._replica_init_group_key(r0)
+    k1 = runtime._replica_init_group_key(r1)
+    k0b = runtime._replica_init_group_key(r0b)
+
+    assert k0 == "device:0"
+    assert k1 == "device:1"
+    # Different physical devices -> different groups (parallel).
+    assert k0 != k1
+    # Same physical device -> same group (serialized under LOCK_EX).
+    assert k0 == k0b
+
+
+def test_init_group_key_parallel_is_unique_per_replica():
+    runtime = _runtime(parallel_stage_init=True)
+    runtime._init_visible_devices_baseline = None
+    # Both on the SAME device, but parallel mode must give distinct keys so they
+    # run concurrently, coordinated by the child SH/EX locks + admission.
+    keys = [
+        runtime._replica_init_group_key(_llm_replica(0, 0, "0")),
+        runtime._replica_init_group_key(_llm_replica(0, 1, "0")),
+        runtime._replica_init_group_key(_llm_replica(1, 0, "0")),
+    ]
+    assert keys == ["parallel:0:0", "parallel:0:1", "parallel:1:0"]
+    assert len(set(keys)) == 3
+
+
+def test_init_group_key_remote_and_diffusion_unchanged():
+    runtime = _runtime(parallel_stage_init=True)
+    remote = _llm_replica(1, 0, "0")
+    remote.launch_mode = "remote"
+    remote.metadata.runtime_cfg = None
+    assert runtime._replica_init_group_key(remote) == "remote:1:0"
+
+    diffusion = _llm_replica(2, 0, "0")
+    diffusion.metadata.stage_type = "diffusion"
+    assert runtime._replica_init_group_key(diffusion) == "inline:diffusion"
+
+
+# --------------------------------------------------------------------------- #
+# B1: admission
+# --------------------------------------------------------------------------- #
+def test_evaluate_fits():
+    cap = 40 * 1024**3
+    demands = [StageDemand(0, 0, [0], utilization=0.5, graph_reserve_bytes=1 * 1024**3)]
+    ledgers = evaluate(demands, {0: cap}, external_reserve_bytes=0, safety_margin_bytes=0)
+    assert ledgers[0].fits
+    assert ledgers[0].kv_budget_bytes == int(cap * 0.5)
+
+
+def test_evaluate_oversubscription_raises():
+    cap = 40 * 1024**3
+    demands = [
+        StageDemand(0, 0, [0], utilization=0.7, graph_reserve_bytes=0),
+        StageDemand(1, 0, [0], utilization=0.7, graph_reserve_bytes=0),
+    ]
+    with pytest.raises(StageAdmissionError):
+        evaluate(demands, {0: cap}, external_reserve_bytes=0, safety_margin_bytes=0)
+
+
+def test_evaluate_unknown_capacity_raises():
+    with pytest.raises(StageAdmissionError):
+        evaluate([StageDemand(0, 0, [3], 0.5, 0)], {0: 1})
+
+
+def test_graph_reserve_positive_when_enabled_zero_when_disabled():
+    # cudagraph enabled -> conservative positive reserve
+    reserve = graph_reserve_bytes(_fake_vllm_config(0.9, cudagraph_mode="FULL"))
+    assert reserve > 0
+    # cudagraph disabled -> exactly 0
+    assert graph_reserve_bytes(_fake_vllm_config(0.9, cudagraph_mode="NONE")) == 0
+    eager = _fake_vllm_config(0.9)
+    eager.model_config.enforce_eager = True
+    assert graph_reserve_bytes(eager) == 0
+
+
+def test_check_admission_multi_device_and_grouping():
+    cap = 80 * 1024**3
+    plan_a = LogicalStageInitPlan(
+        stage_idx=0,
+        stage_id=0,
+        replicas=[_llm_replica(0, 0, "0,1", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))],
+    )
+    plan_b = LogicalStageInitPlan(
+        stage_idx=1,
+        stage_id=1,
+        replicas=[_llm_replica(1, 0, "1", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))],
+    )
+    ledgers = check_admission(
+        [plan_a, plan_b],
+        resolve_physical_devices=lambda r: [int(x) for x in r.metadata.runtime_cfg["devices"].split(",")],
+        device_total_memory=lambda _d: cap,
+        external_reserve_bytes=0,
+        safety_margin_bytes=0,
+    )
+    # device 0: only stage 0 (0.4). device 1: stage 0 + stage 1 (0.8).
+    assert ledgers[0].kv_budget_bytes == int(cap * 0.4)
+    assert ledgers[1].kv_budget_bytes == int(cap * 0.4) + int(cap * 0.4)
+    assert ledgers[1].fits
+
+
+def _local_diffusion_replica(stage_id: int, replica_id: int, devices: str, *, util: float | None) -> ReplicaInitPlan:
+    replica = _llm_replica(stage_id, replica_id, devices)
+    replica.metadata.stage_type = "diffusion"
+    replica.stage_vllm_config = None
+    replica.stage_cfg.engine_args = {} if util is None else {"gpu_memory_utilization": util}
+    return replica
+
+
+def test_check_admission_admits_diffusion_with_explicit_util():
+    cap = 80 * 1024**3
+    diff = _local_diffusion_replica(0, 0, "0", util=0.3)
+    plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff])
+    ledgers = check_admission(
+        [plan],
+        resolve_physical_devices=lambda _r: [0],
+        device_total_memory=lambda _d: cap,
+        external_reserve_bytes=0,
+        safety_margin_bytes=0,
+    )
+    assert set(ledgers) == {0}
+    assert ledgers[0].kv_budget_bytes == int(cap * 0.3)
+
+
+def test_check_admission_unresolved_local_raises():
+    """A local replica invisible to the ledger must fail admission, not skip it:
+    it would otherwise initialize concurrently with unbounded demand."""
+    unresolved = _llm_replica(1, 0, None, vllm_config=_fake_vllm_config(0.9, cudagraph_mode="NONE"))
+    plan = LogicalStageInitPlan(stage_idx=0, stage_id=1, replicas=[unresolved])
+    with pytest.raises(StageAdmissionError, match="stage1/replica0.*unresolved"):
+        check_admission(
+            [plan],
+            resolve_physical_devices=lambda _r: None,
+            device_total_memory=lambda _d: 80 * 1024**3,
+        )
+
+
+def test_check_admission_local_diffusion_without_util_raises():
+    """The reviewer's repro: local LLM + local diffusion with no explicit util
+    used to leave the diffusion stage out of the ledger silently."""
+    llm = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.45, cudagraph_mode="NONE"))
+    diff = _local_diffusion_replica(1, 0, "0", util=None)
+    plans = [
+        LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[llm]),
+        LogicalStageInitPlan(stage_idx=1, stage_id=1, replicas=[diff]),
+    ]
+    with pytest.raises(StageAdmissionError, match="diffusion, no gpu_memory_utilization"):
+        check_admission(
+            plans,
+            resolve_physical_devices=lambda _r: [0],
+            device_total_memory=lambda _d: 80 * 1024**3,
+        )
+
+
+def test_check_admission_exempt_replicas_skip_without_error():
+    """Only ADMISSION_EXEMPT (remote/operator-isolated) may skip the ledger."""
+    cap = 80 * 1024**3
+    local = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))
+    remote = _llm_replica(1, 0, None, vllm_config=_fake_vllm_config(0.9, cudagraph_mode="NONE"))
+    remote.launch_mode = "remote"
+    plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[local, remote])
+    ledgers = check_admission(
+        [plan],
+        resolve_physical_devices=lambda r: ADMISSION_EXEMPT if r.launch_mode == "remote" else [0],
+        device_total_memory=lambda _d: cap,
+        external_reserve_bytes=0,
+        safety_margin_bytes=0,
+    )
+    assert set(ledgers) == {0}
+    assert ledgers[0].contributors == ["stage0/replica0"]
+    assert ledgers[0].kv_budget_bytes == int(cap * 0.4)
+
+
+def test_run_stage_admission_excludes_remote_replicas(monkeypatch):
+    """Remote replicas consume a remote node's memory — never the local ledger.
+
+    Remote LLM plans already carry runtime_cfg=None, but remote diffusion
+    replicas keep their cfg + utilization, so without an explicit launch_mode
+    guard they would be counted against local devices. They must resolve to
+    ADMISSION_EXEMPT (not None): None now means an unaccountable LOCAL replica
+    and fails admission.
+    """
+    import vllm_omni.engine.stage_admission as admission_mod
+
+    runtime = _runtime(parallel_stage_init=True)
+    runtime._init_visible_devices_baseline = None
+
+    local = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))
+    remote_diff = _llm_replica(1, 0, "0")
+    remote_diff.launch_mode = "remote"
+    remote_diff.metadata.stage_type = "diffusion"
+    remote_diff.stage_vllm_config = None
+    remote_diff.stage_cfg.engine_args = {"gpu_memory_utilization": 0.3}
+
+    captured = {}
+
+    def _fake_check_admission(stage_plans, *, resolve_physical_devices, device_total_memory, **kw):
+        captured["resolve"] = resolve_physical_devices
+        return {}
+
+    monkeypatch.setattr(admission_mod, "check_admission", _fake_check_admission)
+    runtime._run_stage_admission([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[local, remote_diff])])
+
+    assert captured["resolve"](local) == [0]
+    assert captured["resolve"](remote_diff) is ADMISSION_EXEMPT
+
+
+# --------------------------------------------------------------------------- #
+# B1b: admission capacity is queried in the visible-ordinal namespace
+# --------------------------------------------------------------------------- #
+def _capture_admission(monkeypatch, runtime, replicas):
+    """Run _run_stage_admission with check_admission stubbed; return its callables."""
+    import vllm_omni.engine.stage_admission as admission_mod
+
+    captured: dict = {}
+
+    def _fake_check_admission(stage_plans, *, resolve_physical_devices, device_total_memory, **kw):
+        captured["resolve"] = resolve_physical_devices
+        captured["capacity"] = device_total_memory
+        return {}
+
+    monkeypatch.setattr(admission_mod, "check_admission", _fake_check_admission)
+    runtime._run_stage_admission([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=replicas)])
+    return captured
+
+
+def _stub_platform(monkeypatch, seen: list[int], total: int = 80 * 1024**3):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    monkeypatch.setattr(
+        runtime_mod,
+        "current_omni_platform",
+        types.SimpleNamespace(
+            device_control_env_var="CUDA_VISIBLE_DEVICES",
+            get_device_total_memory=lambda d: (seen.append(d), total)[1],
+        ),
+    )
+
+
+@pytest.mark.parametrize(
+    ("baseline", "devices", "physical", "ordinal"),
+    [
+        ("3,4,5", "1", 4, 1),  # restricted visibility
+        ("1,0", "0", 1, 0),  # reordered visibility
+        ("5,GPU-deadbeef", "0", 5, 0),  # mixed identifiers, numeric first
+        ("GPU-deadbeef,5", "1", 5, 1),  # mixed identifiers, numeric second
+        ("0,1,2,3", "2", 2, 2),  # identity
+    ],
+    ids=["restricted", "reordered", "mixed-numeric-first", "mixed-numeric-second", "identity"],
+)
+def test_admission_capacity_uses_visible_ordinal(monkeypatch, baseline, devices, physical, ordinal):
+    """The ledger stays keyed by physical id, but the capacity query is translated.
+
+    ``get_device_total_memory``/``torch.cuda.get_device_properties`` read their
+    argument as an ordinal in the *current* process's visible-device namespace.
+    With CUDA_VISIBLE_DEVICES=3,4,5 the orchestrator only has ordinals 0..2, so
+    passing physical id 4 straight through would raise — or, when visibility is
+    reordered, silently measure a different GPU.
+    """
+    runtime = _runtime(parallel_stage_init=True)
+    runtime._init_visible_devices_baseline = baseline
+    replica = _llm_replica(0, 0, devices, vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))
+
+    seen: list[int] = []
+    _stub_platform(monkeypatch, seen)
+    captured = _capture_admission(monkeypatch, runtime, [replica])
+
+    assert captured["resolve"](replica) == [physical], "ledger key must stay physical"
+    assert captured["capacity"](physical) == 80 * 1024**3
+    assert seen == [ordinal], f"expected the capacity query on ordinal {ordinal}, got {seen}"
+
+
+def test_admission_capacity_rejects_device_outside_baseline(monkeypatch):
+    """A physical id absent from the captured baseline cannot be measured here."""
+    runtime = _runtime(parallel_stage_init=True)
+    runtime._init_visible_devices_baseline = "3,4,5"
+    replica = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))
+
+    seen: list[int] = []
+    _stub_platform(monkeypatch, seen)
+    captured = _capture_admission(monkeypatch, runtime, [replica])
+
+    with pytest.raises(StageAdmissionError, match="not in the visibility baseline"):
+        captured["capacity"](7)
+    assert seen == [], "must not query the platform with an unmappable id"
+
+
+def test_admission_capacity_identity_when_no_baseline(monkeypatch):
+    """With no device-control env captured, ordinals and physical ids coincide."""
+    runtime = _runtime(parallel_stage_init=True)
+    runtime._init_visible_devices_baseline = None
+    replica = _llm_replica(0, 0, "2", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))
+
+    seen: list[int] = []
+    _stub_platform(monkeypatch, seen)
+    captured = _capture_admission(monkeypatch, runtime, [replica])
+
+    assert captured["capacity"](2) == 80 * 1024**3
+    assert seen == [2]
+
+
+# --------------------------------------------------------------------------- #
+# B1c: fail closed for executors the phase locks cannot guard
+# --------------------------------------------------------------------------- #
+class _RayBackedExecutor:
+    uses_ray = True
+
+
+class _LocalExecutor:
+    uses_ray = False
+
+
+def test_parallel_init_refuses_ray_backed_executor():
+    """Phase locks are node-local flocks and the ledger only sees this host.
+
+    A Ray executor may place the workers that load weights, profile, allocate KV
+    and capture graphs on other nodes, which participate in neither mechanism —
+    and the parent skips its own device lock when parallel_stage_init is on.
+    """
+    runtime = _runtime(parallel_stage_init=True)
+    replica = _llm_replica(0, 0, "0")
+    replica.executor_class = _RayBackedExecutor
+    plans = [LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[replica])]
+
+    with pytest.raises(RuntimeError, match="cannot guard Ray-backed or multi-node executors"):
+        runtime._reject_unguardable_executors(plans)
+
+
+@pytest.mark.parametrize(
+    ("use_ray", "nnodes", "reason"),
+    [
+        (True, 1, "Ray-backed"),
+        (False, 2, r"multi-node \(nnodes=2\)"),
+    ],
+    ids=["parallel-config-ray", "multi-node-mp"],
+)
+def test_parallel_init_refuses_unguardable_parallel_config(use_ray, nnodes, reason):
+    """The config guard also covers non-Ray executors spread across nodes."""
+    runtime = _runtime(parallel_stage_init=True)
+    vllm_config = types.SimpleNamespace(
+        parallel_config=types.SimpleNamespace(use_ray=use_ray, nnodes=nnodes),
+    )
+    replica = _llm_replica(0, 0, "0", vllm_config=vllm_config)
+    replica.executor_class = _LocalExecutor
+    plans = [LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[replica])]
+
+    with pytest.raises(RuntimeError, match=reason):
+        runtime._reject_unguardable_executors(plans)
+
+
+def test_parallel_init_allows_non_ray_executor():
+    runtime = _runtime(parallel_stage_init=True)
+    replicas = [_llm_replica(0, 0, "0"), _llm_replica(1, 0, "1")]
+    replicas[0].executor_class = _LocalExecutor
+    replicas[1].executor_class = object  # no uses_ray attribute at all
+    runtime._reject_unguardable_executors([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=replicas)])
+
+
+def test_parallel_init_ignores_remote_ray_replicas():
+    """Remote replicas are launched by the runtime owning their node."""
+    runtime = _runtime(parallel_stage_init=True)
+    replica = _llm_replica(0, 0, "0")
+    replica.executor_class = _RayBackedExecutor
+    replica.launch_mode = "remote"
+    runtime._reject_unguardable_executors([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[replica])])
+
+
+# --------------------------------------------------------------------------- #
+# B2: SH/EX device phase locks (real flock)
+# --------------------------------------------------------------------------- #
+def _lock(tmp: Path, **kw) -> DevicePhaseLock:
+    return DevicePhaseLock([0], lock_dir=str(tmp), timeout_s=kw.pop("timeout_s", 0.3), **kw)
+
+
+def test_shared_locks_coexist(tmp_path):
+    a, b = _lock(tmp_path), _lock(tmp_path)
+    with a.shared():
+        with b.shared():
+            pass  # no timeout
+
+
+def test_exclusive_excludes_shared_fail_closed(tmp_path):
+    a, b = _lock(tmp_path, timeout_s=0.2), _lock(tmp_path, timeout_s=0.2)
+    with a.exclusive():
+        with pytest.raises(DeviceLockTimeoutError):
+            with b.shared():
+                pass
+
+
+def test_exclusive_excludes_exclusive(tmp_path):
+    a, b = _lock(tmp_path, timeout_s=0.2), _lock(tmp_path, timeout_s=0.2)
+    with a.exclusive():
+        with pytest.raises(DeviceLockTimeoutError):
+            with b.exclusive():
+                pass
+
+
+def test_empty_device_set_is_noop(tmp_path):
+    lock = DevicePhaseLock([], lock_dir=str(tmp_path))
+    with lock.shared():
+        pass
+    with lock.exclusive():
+        pass
+
+
+def _dead_pid() -> int:
+    """A PID that has certainly exited and been reaped."""
+    proc = subprocess.Popen(["/bin/true"])
+    proc.wait()
+    return proc.pid
+
+
+def test_legacy_device_lock_never_unlinks_the_path():
+    """The stale-PID unlink path is gone from the legacy acquirer.
+
+    ``flock`` is released by the kernel when the holder exits (SIGKILL included),
+    so a dead holder never keeps the lock and unlink-based cleanup is never
+    needed. It was also unsafe: the cleanup only ran on ``BlockingIOError``, i.e.
+    only while the lock *was* held, so every case it fired in was a case where
+    unlinking broke a live holder.
+    """
+    import vllm_omni.engine.stage_init_utils as siu
+
+    assert not hasattr(siu, "_cleanup_stale_lock_if_dead")
+    assert "unlink" not in inspect.getsource(siu.acquire_device_locks)
+
+    # Test cleanup runs in many autouse/runtime fixtures and must not unlink the
+    # shared production lock path either. Persistent unlocked flock files are
+    # harmless; deleting one can split live holders across different inodes.
+    helper_source = (Path(__file__).parents[1] / "helpers" / "clean.py").read_text()
+    assert "_cleanup_stale_device_locks" not in helper_source
+    assert "vllm_omni_device_*_init.lock" not in helper_source
+
+
+def test_legacy_lock_cannot_break_a_live_shared_holder(tmp_path, monkeypatch):
+    """A live LOCK_SH holder must keep excluding the legacy LOCK_EX acquirer
+    even when the PID recorded in the lock file belongs to a dead process.
+
+    Several LOCK_SH holders overwrite that PID, so the recorded value can name a
+    process that has exited while another holder still owns the inode. Unlinking
+    on that signal let the legacy acquirer create a fresh file under the same
+    name and take LOCK_EX on a *different* inode — conflicting with nobody.
+    """
+    import vllm_omni.engine.stage_init_utils as siu
+
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
+    monkeypatch.setattr(
+        siu,
+        "device_init_lock_path",
+        lambda device_id, lock_dir="/tmp": device_init_lock_path(device_id, str(tmp_path)),
+    )
+
+    path = Path(device_init_lock_path(0, str(tmp_path)))
+    holder = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=5)
+
+    with holder.shared():
+        path.write_text(f"{_dead_pid()}\n")  # peer SH holder exited after recording its PID
+        inode_before = path.stat().st_ino
+
+        lock_fds = siu.acquire_device_locks(0, {"tensor_parallel_size": 1}, stage_init_timeout=1)
+
+        assert lock_fds == [], "legacy LOCK_EX must not be granted while a shared holder is live"
+        assert path.exists(), "the lock path must never be unlinked"
+        assert path.stat().st_ino == inode_before, "the lock inode must stay stable"
+
+    siu.release_device_locks(lock_fds)
+
+
+def test_phase_lock_reuses_the_existing_lock_inode(tmp_path):
+    """Repeated acquisitions coordinate on one stable inode per device."""
+    path = Path(device_init_lock_path(0, str(tmp_path)))
+    lock = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=1)
+
+    with lock.shared():
+        first = path.stat().st_ino
+    with lock.exclusive():
+        assert path.stat().st_ino == first
+    assert os.path.exists(path)
+
+
+def test_device_lock_file_never_visible_at_restrictive_mode(tmp_path):
+    """Concurrent creators must never observe the lock at a umask-filtered mode.
+
+    Creating in place with O_EXCL and widening afterwards publishes the file at
+    0600 under a restrictive umask; a second user opening it in that window is
+    locked out and parallel init aborts. Publishing atomically closes the window,
+    so every concurrent creator either wins the link or opens a readable file.
+    """
+    path = device_init_lock_path(0, str(tmp_path))
+    results: list[tuple[int, bool]] = []
+    errors: list[BaseException] = []
+    modes: list[int] = []
+    barrier = threading.Barrier(8)
+
+    def _worker() -> None:
+        barrier.wait()
+        try:
+            fd, writable = open_device_lock_file(path)
+        except BaseException as exc:  # noqa: BLE001 - surfaced by the assert below
+            errors.append(exc)
+            return
+        try:
+            modes.append(os.stat(path).st_mode & 0o777)
+            results.append((fd, writable))
+        finally:
+            os.close(fd)
+
+    previous = os.umask(0o077)
+    try:
+        threads = [threading.Thread(target=_worker) for _ in range(8)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+    finally:
+        os.umask(previous)
+
+    assert not errors, f"concurrent creators were locked out: {errors!r}"
+    assert len(results) == 8
+    assert all(m & 0o044 == 0o044 for m in modes), f"lock was visible at {[oct(m) for m in modes]}"
+    assert len({os.stat(path).st_ino}) == 1
+
+
+@pytest.mark.parametrize("umask_value", [0o022, 0o027, 0o077], ids=["022", "027", "077"])
+def test_device_lock_file_is_readable_by_other_users(tmp_path, umask_value):
+    """The lock file must stay group/other readable whatever the creator's umask.
+
+    These files coordinate GPU init *across* users, so a second user has to be
+    able to open the first user's file. The mode passed to ``os.open`` is filtered
+    by umask, so a restrictive umask (0027, 0077) would produce 0640 or 0600 and
+    lock every other user out -- fatal for parallel_stage_init, and a silent loss
+    of serialization in the legacy path.
+    """
+    path = device_init_lock_path(0, str(tmp_path))
+    previous = os.umask(umask_value)
+    try:
+        fd, writable = open_device_lock_file(path)
+    finally:
+        os.umask(previous)
+    try:
+        assert writable is True
+        mode = os.stat(path).st_mode & 0o777
+        assert mode & 0o044 == 0o044, f"lock file mode {mode:04o} is not readable by other users"
+    finally:
+        os.close(fd)
+
+
+@pytest.mark.skipif(os.geteuid() == 0, reason="root bypasses file permission bits")
+def test_phase_lock_tolerates_lock_file_owned_by_another_user(tmp_path):
+    """A lock file created by a different user must not break locking.
+
+    These files coordinate *across* users -- two people on the same physical GPU
+    must contend on the same path -- so whoever creates one owns it at 0644 and
+    every later user's O_RDWR open fails with EACCES. That made parallel_stage_init
+    die at startup on any shared machine. flock attaches to the open file
+    description and needs no write access, so the read-only fallback keeps full
+    mutual exclusion and only forfeits the diagnostic PID write.
+    """
+    path = Path(device_init_lock_path(0, str(tmp_path)))
+    path.write_text("999999\n")
+    path.chmod(0o444)
+
+    fd, writable = open_device_lock_file(str(path))
+    try:
+        assert writable is False, "expected the read-only fallback on an unwritable file"
+    finally:
+        os.close(fd)
+
+    holder = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=0.3)
+    contender = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=0.2)
+    with holder.exclusive():
+        with pytest.raises(DeviceLockTimeoutError):
+            with contender.shared():
+                pass
+
+    assert path.read_text() == "999999\n", "must not clobber a file it cannot own"
+
+
+def test_resolve_driver_device_ids_dp_slice(monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    env_var = current_omni_platform.device_control_env_var
+    monkeypatch.setenv(env_var, "4,5,6,7")
+    pc = types.SimpleNamespace(
+        tensor_parallel_size=2,
+        pipeline_parallel_size=1,
+        prefill_context_parallel_size=1,
+        sequence_parallel_size=1,
+        cfg_parallel_size=1,
+    )
+    vc = types.SimpleNamespace(parallel_config=pc)
+    assert resolve_driver_device_ids(vc, local_dp_rank=0) == [4, 5]
+    assert resolve_driver_device_ids(vc, local_dp_rank=1) == [6, 7]
+
+
+# --------------------------------------------------------------------------- #
+# B2: executor wrapper ordering + retry
+# --------------------------------------------------------------------------- #
+class _RecordingLocker:
+    def __init__(self):
+        self.events: list[str] = []
+
+    @contextlib.contextmanager
+    def shared(self):
+        self.events.append("SH-enter")
+        try:
+            yield
+        finally:
+            self.events.append("SH-exit")
+
+    @contextlib.contextmanager
+    def exclusive(self):
+        self.events.append("EX-enter")
+        try:
+            yield
+        finally:
+            self.events.append("EX-exit")
+
+
+class _FakeExecutor:
+    instances: list[object] = []
+
+    def __init__(self, vllm_config):
+        self.vllm_config = vllm_config
+        self.dam_calls = 0
+        _FakeExecutor.instances.append(self)
+
+    def determine_available_memory(self):
+        self.dam_calls += 1
+        return 123
+
+    def initialize_from_config(self, kv_cache_configs):
+        return "ok"
+
+    def compile_or_warm_up_model(self):
+        return "warm"
+
+
+def test_wrapper_phase_ordering():
+    locker = _RecordingLocker()
+    wrapped_cls = wrap_executor_with_phase_locks(_FakeExecutor, locker)
+    exec_ = wrapped_cls(vllm_config=object())
+    assert exec_.determine_available_memory() == 123
+    assert exec_.initialize_from_config([]) == "ok"
+    assert exec_.compile_or_warm_up_model() == "warm"
+    assert locker.events == [
+        "SH-enter",
+        "SH-exit",  # load
+        "EX-enter",
+        "EX-exit",  # profile
+        "SH-enter",
+        "SH-exit",  # KV cache allocation
+        "SH-enter",
+        "SH-exit",  # kernel warmup + CUDA-graph capture
+    ]
+
+
+def test_wrapper_locks_warmup_and_capture_separately():
+    """Regression guard for the upstream split.
+
+    ``Executor.initialize_from_config`` used to fuse ``compile_or_warm_up_model``
+    into itself; since vLLM v0.20.0 ``EngineCore._initialize_kv_caches`` calls the
+    two separately. If the wrapper only brackets ``initialize_from_config``, one
+    stage captures CUDA graphs — mutating device memory — while a peer holds the
+    exclusive lock for its profiling measurement.
+    """
+    locker = _RecordingLocker()
+    wrapped_cls = wrap_executor_with_phase_locks(_FakeExecutor, locker)
+    exec_ = wrapped_cls(vllm_config=object())
+    locker.events.clear()
+
+    assert exec_.compile_or_warm_up_model() == "warm"
+
+    assert locker.events == ["SH-enter", "SH-exit"], "compile_or_warm_up_model must run under the shared device lock"
+    assert "compile_or_warm_up_model" in vars(wrapped_cls), (
+        "the wrapper must override compile_or_warm_up_model, not inherit it unguarded"
+    )
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [{}, {"vllm_config": object()}, {"executor_class": object}],
+    ids=["both-missing", "executor-missing", "config-missing"],
+)
+def test_install_phase_locks_fails_closed_on_missing_kwargs(kwargs):
+    """parallel_stage_init must never proceed with the phase-lock guard
+    uninstalled — a missing vllm_config/executor_class aborts the child."""
+    with pytest.raises(RuntimeError, match="phase-lock guard cannot be installed"):
+        _install_phase_locks(kwargs, local_dp_rank=0)
+
+
+def test_install_phase_locks_wraps_executor(monkeypatch):
+    import vllm_omni.engine.stage_phase_lock as phase_lock_mod
+
+    locker = _RecordingLocker()
+    locker.device_ids = [0]
+    monkeypatch.setattr(
+        phase_lock_mod.DevicePhaseLock,
+        "from_child",
+        classmethod(lambda _cls, _cfg, _rank: locker),
+    )
+    kwargs = {"vllm_config": object(), "executor_class": _FakeExecutor}
+
+    _install_phase_locks(kwargs, local_dp_rank=0)
+
+    exec_ = kwargs["executor_class"](vllm_config=object())
+    assert exec_.determine_available_memory() == 123
+    assert locker.events[:2] == ["SH-enter", "SH-exit"]
+
+
+# --------------------------------------------------------------------------- #
+# A2: NVML process-scope removal regression guard
+# --------------------------------------------------------------------------- #
+def test_base_worker_dropped_nvml_process_scope():
+    import vllm_omni.worker.base as base
+
+    src = Path(base.__file__).read_text()
+    assert "get_process_gpu_memory" not in src
+    assert "is_process_scoped_memory_available" not in src
+
+
+def test_device_ledger_required_and_fits():
+    led = DeviceLedger(device_id=0, capacity_bytes=100, kv_budget_bytes=40, graph_reserve_bytes=20)
+    led.external_reserve_bytes = 10
+    led.safety_margin_bytes = 10
+    assert led.required_bytes == 80
+    assert led.fits
+    led.kv_budget_bytes = 90
+    assert not led.fits

tests/helpers/clean.py

@@ -170,36 +170,6 @@ def _print_device_processes() -> None:
         print("=" * 80)
 
 
-def _cleanup_stale_device_locks() -> None:
-    """Remove stale device-initialization lock files whose recorded PID is dead.
-
-    Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a
-    crashed / killed test run and block subsequent orchestrator startups.
-    """
-    import glob as _glob
-
-    for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"):
-        try:
-            with open(lock_file) as fh:
-                content = fh.read().strip()
-            if not content:
-                continue
-            pid = int(content)
-        except (OSError, ValueError):
-            continue
-
-        try:
-            os.kill(pid, 0)
-        except ProcessLookupError:
-            print(f"  Removing stale device lock {lock_file} (PID {pid} is dead)")
-            try:
-                os.unlink(lock_file)
-            except OSError:
-                pass
-        except PermissionError:
-            pass
-
-
 def cleanup_test_environment(*, shutdown_ray: bool = False) -> None:
     """Tear down distributed state and reset device memory for tests."""
     from vllm import envs
@@ -224,8 +194,6 @@ def cleanup_test_environment(*, shutdown_ray: bool = False) -> None:
         ray.shutdown()
 
     print("Pre-test device status:")
-    _cleanup_stale_device_locks()
-
     num_devices = current_omni_platform.device_count()
     if num_devices > 0:
         try:

tests/worker/test_worker_base.py

@@ -2,9 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Characterization tests for ``vllm_omni.worker.base.OmniGPUWorkerBase``.
 
-Pins the CURRENT behaviour of ``determine_available_memory`` (BOTH the NVML /
-process-scoped arm AND the profiling-fallback arm — this base still carries the
-NVML branch), plus the
+Pins the CURRENT behaviour of ``determine_available_memory`` (device-level
+profiling is the only path — the NVML / process-scoped arm was removed with
+parallel stage init, which coordinates concurrent same-device measurement via
+admission + SH/EX device locks instead), plus the
 memory-pool / sleep / wake-up plumbing that a later change may touch.
 
 Workers are constructed via ``object.__new__`` with only the attributes each
@@ -57,27 +58,20 @@ def _make_worker(*, requested_memory: int, kv_cache_memory_bytes: int = 0, model
 # --------------------------------------------------------------------------- #
 # determine_available_memory                                                  #
 # --------------------------------------------------------------------------- #
-def test_determine_available_memory_nvml_arm(monkeypatch):
-    """NVML available -> available = requested - per-process memory."""
+def test_no_process_scoped_collaborators_remain():
+    """The NVML arm is gone from base.py, so the method CANNOT take a
+    process-scoped path on any host: its collaborators are not even imported.
+    (The removal itself happened in the parallel-stage-init feature commit;
+    tests/engine/test_parallel_stage_init.py pins the same invariant.)"""
+    assert not hasattr(base, "is_process_scoped_memory_available")
+    assert not hasattr(base, "detect_pid_host")
+    assert not hasattr(base, "get_process_gpu_memory")
+
+
+def test_determine_available_memory_profiling_path(monkeypatch):
+    """available = requested - (weights + peak + non_torch); the only path."""
     worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB)
     monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB))
-    monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: True)
-    monkeypatch.setattr(base, "detect_pid_host", lambda: True)
-    monkeypatch.setattr(base, "get_process_gpu_memory", lambda local_rank: 8 * GIB)
-
-    out = OmniGPUWorkerBase.determine_available_memory(worker)
-
-    # NVML arm ignores profiled usage; uses process_memory directly.
-    assert out == 30 * GIB - 8 * GIB
-
-
-def test_determine_available_memory_profiling_fallback_arm(monkeypatch):
-    """NVML unavailable -> available = requested - (weights + peak + non_torch)."""
-    worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB)
-    monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB))
-    monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: False)
-    # detect_pid_host is short-circuited by the `and`, but pin it anyway.
-    monkeypatch.setattr(base, "detect_pid_host", lambda: True)
 
     out = OmniGPUWorkerBase.determine_available_memory(worker)
 
@@ -85,25 +79,20 @@ def test_determine_available_memory_profiling_fallback_arm(monkeypatch):
     assert out == 30 * GIB - profiled
 
 
-def test_determine_available_memory_falls_back_when_not_pid_host(monkeypatch):
-    """detect_pid_host False also routes to the profiling fallback."""
-    worker = _make_worker(requested_memory=20 * GIB, model_memory_usage=5 * GIB)
-    monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=0, torch_peak=0))
-    monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: True)
-    monkeypatch.setattr(base, "detect_pid_host", lambda: False)
+def test_determine_available_memory_populates_total_consumed(monkeypatch):
+    """total_consumed mirrors upstream's MemoryProfilingResult field."""
+    worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB)
+    monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB))
 
-    out = OmniGPUWorkerBase.determine_available_memory(worker)
+    OmniGPUWorkerBase.determine_available_memory(worker)
 
-    assert out == 20 * GIB - 5 * GIB
+    assert worker.total_consumed == 3 * GIB
 
 
 def test_determine_available_memory_clamps_to_zero(monkeypatch):
     """Over-subscription clamps the KV budget to 0, never negative."""
-    worker = _make_worker(requested_memory=5 * GIB, model_memory_usage=0)
+    worker = _make_worker(requested_memory=5 * GIB, model_memory_usage=8 * GIB)
     monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=0, torch_peak=0))
-    monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: True)
-    monkeypatch.setattr(base, "detect_pid_host", lambda: True)
-    monkeypatch.setattr(base, "get_process_gpu_memory", lambda local_rank: 8 * GIB)
 
     out = OmniGPUWorkerBase.determine_available_memory(worker)
 

vllm_omni/config/omni_config.py

@@ -1299,6 +1299,12 @@ class VllmOmniOrchestratorConfig:
     omni_lb_policy: str = "random"
     omni_heartbeat_timeout: float = Field(default=30.0, gt=0.0)
     batch_timeout: int = Field(default=10, ge=0)
+    # When True, stages sharing a physical GPU initialize concurrently, guarded
+    # by pre-launch admission control + engine-core-held SH/EX device locks
+    # (see stage_admission / stage_phase_lock). Default False keeps the legacy
+    # per-device LOCK_EX serialization. Enable only when the GPU is dedicated to
+    # this deployment (see rfc_parallel_stage_init).
+    parallel_stage_init: bool = False
 
 
 @config(config=ConfigDict(arbitrary_types_allowed=True))

vllm_omni/engine/arg_utils.py

@@ -504,6 +504,10 @@ class OrchestratorArgs:
     # === Lifecycle ===
     stage_init_timeout: int = 300
     init_timeout: int = 600
+    # Initialize stages sharing a physical GPU concurrently, guarded by
+    # pre-launch admission + engine-core SH/EX device locks. Off by default;
+    # enable only when the GPU is dedicated to this deployment.
+    parallel_stage_init: bool = False
 
     # === Cross-stage Communication ===
     batch_timeout: int = 10

vllm_omni/engine/async_omni_engine.py

@@ -185,6 +185,10 @@ class AsyncOmniEngine:
         self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0)
         if self._omni_heartbeat_timeout <= 0:
             raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}")
+        # Concurrent same-device stage init (admission + SH/EX phase locks).
+        # Sourced from the parallel_stage_init orchestrator/CLI arg (config,
+        # not an env var); default False preserves serial init.
+        self._parallel_stage_init: bool = bool(kwargs.get("parallel_stage_init") or False)
 
         if single_stage_mode:
             logger.info(
@@ -322,6 +326,7 @@ class AsyncOmniEngine:
             stage_init_timeout=stage_init_timeout,
             async_chunk=self.async_chunk,
             tokenizer=self.tokenizer,
+            parallel_stage_init=self._parallel_stage_init,
             single_stage_id_filter=self._single_stage_id_filter,
             omni_master_address=self._omni_master_address,
             omni_master_port=self._omni_master_port,

vllm_omni/engine/stage_admission.py

@@ -0,0 +1,297 @@
+"""Pre-launch admission control for parallel stage initialization.
+
+When ``VllmOmniOrchestratorConfig.parallel_stage_init`` is enabled, several stage
+engine cores initialize concurrently and allocate KV cache / capture CUDA graphs
+independently. The SH/EX phase locks (``stage_phase_lock``) keep each *measurement*
+clean, but they cannot bound the *sum* of what all stages will allocate — two
+stages can each profile against a mostly-empty GPU, each compute a large KV
+budget, and then both allocate → OOM.
+
+Admission is the hard backstop for that: **before any stage launches**, prove for
+every physical device *g*::
+
+    Σ_{s on g} capacity(g)·utilization(s) + Σ graph_reserve(s,g)
+        + external_reserve + safety_margin  ≤  capacity(g)
+
+If it does not hold, fail fast (``StageAdmissionError``) rather than OOM at
+runtime. Because the budgets are proven to fit before anyone allocates, any
+profile/allocate interleaving is safe.
+
+The arithmetic (``evaluate``) is pure and unit-testable. Plan-walking
+(``check_admission``) takes injectable callables for device resolution and total
+memory so it can also be exercised without a GPU.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+from vllm.logger import init_logger
+from vllm.utils.mem_utils import format_gib
+
+logger = init_logger(__name__)
+
+# Tunable reserves (bytes). Conservative defaults; operators can pass overrides
+# to ``check_admission``. Kept as constants (not env vars) per project policy.
+_DEFAULT_EXTERNAL_RESERVE_BYTES = 1 * 1024**3  # headroom for unmanaged consumers
+_DEFAULT_SAFETY_MARGIN_BYTES = 1 * 1024**3  # fragmentation / allocator slack
+
+# Graph-pool reserve. A single conservative constant used whenever CUDA-graph
+# capture is enabled (never 0 in that case); zero when capture is disabled.
+# Kept as a constant (not env var / not a measured table) per project policy.
+_DEFAULT_GRAPH_RESERVE_BYTES = 2 * 1024**3
+
+
+class StageAdmissionError(RuntimeError):
+    """Raised when the planned stage budgets do not fit the available devices."""
+
+
+class AdmissionExempt:
+    """Type of :data:`ADMISSION_EXEMPT`; not meant to be instantiated elsewhere."""
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "ADMISSION_EXEMPT"
+
+
+ADMISSION_EXEMPT = AdmissionExempt()
+"""Resolver sentinel: the replica is deliberately outside local admission
+(e.g. it runs on a remote node and consumes that node's memory). Only replicas
+marked exempt may skip the ledger — a local replica the resolver cannot account
+for fails admission instead (see :func:`check_admission`)."""
+
+
+@dataclass
+class StageDemand:
+    """One replica's claim on a set of physical devices."""
+
+    stage_id: int
+    replica_id: int
+    device_ids: list[int]
+    utilization: float
+    graph_reserve_bytes: int
+    is_diffusion: bool = False
+
+
+@dataclass
+class DeviceLedger:
+    """Per-device admission accounting (also the integration-run record)."""
+
+    device_id: int
+    capacity_bytes: int
+    kv_budget_bytes: int = 0
+    graph_reserve_bytes: int = 0
+    external_reserve_bytes: int = 0
+    safety_margin_bytes: int = 0
+    contributors: list[str] = field(default_factory=list)
+
+    @property
+    def required_bytes(self) -> int:
+        return self.kv_budget_bytes + self.graph_reserve_bytes + self.external_reserve_bytes + self.safety_margin_bytes
+
+    @property
+    def fits(self) -> bool:
+        return self.required_bytes <= self.capacity_bytes
+
+
+# ---- graph reserve ---------------------------------------------------------
+
+
+def _cudagraph_disabled(vllm_config: Any) -> bool:
+    model_config = getattr(vllm_config, "model_config", None)
+    if getattr(model_config, "enforce_eager", False):
+        return True
+    comp = getattr(vllm_config, "compilation_config", None)
+    mode = getattr(comp, "cudagraph_mode", None)
+    if mode is not None and str(mode).upper().endswith("NONE"):
+        return True
+    return False
+
+
+def graph_reserve_bytes(vllm_config: Any) -> int:
+    """Reserve for this stage's CUDA-graph pool (bytes).
+
+    A single conservative constant when capture is enabled (never 0 in that
+    case), 0 when capture is disabled. This intentionally over-reserves rather
+    than measure per-model — admission stays a safe upper bound.
+    """
+    if _cudagraph_disabled(vllm_config):
+        return 0
+    return _DEFAULT_GRAPH_RESERVE_BYTES
+
+
+# ---- pure evaluation -------------------------------------------------------
+
+
+def evaluate(
+    demands: Sequence[StageDemand],
+    capacities: dict[int, int],
+    *,
+    external_reserve_bytes: int = _DEFAULT_EXTERNAL_RESERVE_BYTES,
+    safety_margin_bytes: int = _DEFAULT_SAFETY_MARGIN_BYTES,
+) -> dict[int, DeviceLedger]:
+    """Build the per-device ledger and raise ``StageAdmissionError`` if any device
+    is over-subscribed. Pure: no config objects, no GPU."""
+    ledgers: dict[int, DeviceLedger] = {}
+
+    def _ledger(dev: int) -> DeviceLedger:
+        if dev not in ledgers:
+            if dev not in capacities:
+                raise StageAdmissionError(f"No capacity known for physical device {dev}")
+            ledgers[dev] = DeviceLedger(
+                device_id=dev,
+                capacity_bytes=capacities[dev],
+                external_reserve_bytes=external_reserve_bytes,
+                safety_margin_bytes=safety_margin_bytes,
+            )
+        return ledgers[dev]
+
+    for d in demands:
+        for dev in d.device_ids:
+            ledger = _ledger(dev)
+            ledger.kv_budget_bytes += int(ledger.capacity_bytes * d.utilization)
+            ledger.graph_reserve_bytes += d.graph_reserve_bytes
+            ledger.contributors.append(f"stage{d.stage_id}/replica{d.replica_id}")
+
+    over = [led for led in ledgers.values() if not led.fits]
+    for led in ledgers.values():
+        logger.info(
+            "[admission] device %d: capacity=%s required=%s (kv=%s graph=%s ext=%s margin=%s) "
+            "headroom=%s contributors=%s%s",
+            led.device_id,
+            format_gib(led.capacity_bytes),
+            format_gib(led.required_bytes),
+            format_gib(led.kv_budget_bytes),
+            format_gib(led.graph_reserve_bytes),
+            format_gib(led.external_reserve_bytes),
+            format_gib(led.safety_margin_bytes),
+            format_gib(led.capacity_bytes - led.required_bytes),
+            led.contributors,
+            " <<< OVER" if not led.fits else "",
+        )
+    if over:
+        detail = "; ".join(
+            f"device {led.device_id}: need {format_gib(led.required_bytes)} GiB > "
+            f"{format_gib(led.capacity_bytes)} GiB (contributors {led.contributors})"
+            for led in over
+        )
+        raise StageAdmissionError(
+            "parallel_stage_init admission failed — per-device budget exceeds capacity. "
+            "Lower gpu_memory_utilization, reduce co-located stages, or disable "
+            f"parallel_stage_init. {detail}"
+        )
+    return ledgers
+
+
+# ---- plan walking ----------------------------------------------------------
+
+
+def check_admission(
+    stage_plans: Sequence[Any],
+    *,
+    resolve_physical_devices: Callable[[Any], list[int] | AdmissionExempt | None],
+    device_total_memory: Callable[[int], int],
+    external_reserve_bytes: int = _DEFAULT_EXTERNAL_RESERVE_BYTES,
+    safety_margin_bytes: int = _DEFAULT_SAFETY_MARGIN_BYTES,
+    graph_reserve: Callable[[Any], int] = graph_reserve_bytes,
+) -> dict[int, DeviceLedger]:
+    """Extract demands from the orchestrator's stage plans and admit them.
+
+    ``resolve_physical_devices(replica)`` returns the physical GPU ids a replica
+    occupies, or :data:`ADMISSION_EXEMPT` for replicas deliberately outside
+    local admission (e.g. running on a remote node). Any *other* replica that
+    cannot be accounted — unresolved or non-integer devices, or a diffusion
+    stage without an explicit ``gpu_memory_utilization`` — raises
+    ``StageAdmissionError``: under parallel init every local replica gets its
+    own init group and no parent holds a whole-init exclusive lock, so a
+    consumer invisible to the ledger would initialize concurrently with
+    unbounded demand, defeating admission (fail-closed).
+    ``device_total_memory(id)`` returns a device's total bytes.
+    """
+    demands: list[StageDemand] = []
+    exempt: list[str] = []
+    unaccounted: list[str] = []
+    for plan in stage_plans:
+        for replica in getattr(plan, "replicas", []):
+            metadata = replica.metadata
+            label = f"stage{metadata.stage_id}/replica{replica.replica_id}"
+            device_ids = resolve_physical_devices(replica)
+            if isinstance(device_ids, AdmissionExempt):
+                exempt.append(label)
+                continue
+            if not device_ids:
+                unaccounted.append(f"{label} (unresolved devices)")
+                continue
+            vllm_config = replica.stage_vllm_config
+            if vllm_config is None:
+                # Diffusion stages have no resolved vllm_config pre-launch and
+                # skip profile/capture; admit them with their raw util. A local
+                # diffusion stage without one has unknowable demand and must
+                # fail admission below, not bypass it.
+                util = _diffusion_utilization(replica)
+                if util is None:
+                    unaccounted.append(f"{label} (diffusion, no gpu_memory_utilization)")
+                    continue
+                demands.append(
+                    StageDemand(
+                        stage_id=metadata.stage_id,
+                        replica_id=replica.replica_id,
+                        device_ids=list(device_ids),
+                        utilization=util,
+                        graph_reserve_bytes=0,
+                        is_diffusion=True,
+                    )
+                )
+                continue
+            demands.append(
+                StageDemand(
+                    stage_id=metadata.stage_id,
+                    replica_id=replica.replica_id,
+                    device_ids=list(device_ids),
+                    utilization=float(vllm_config.cache_config.gpu_memory_utilization),
+                    graph_reserve_bytes=int(graph_reserve(vllm_config)),
+                )
+            )
+
+    if unaccounted:
+        raise StageAdmissionError(
+            "parallel_stage_init admission cannot account for local replicas: "
+            f"{unaccounted}. Every local stage must resolve to integer physical "
+            "device ids and declare a memory budget (diffusion stages: set "
+            "gpu_memory_utilization in engine_args); otherwise it would "
+            "initialize concurrently without bounding its demand. Fix the "
+            "stage config or disable parallel_stage_init."
+        )
+    if exempt:
+        logger.info(
+            "[admission] exempt from local admission (remote/operator-isolated): %s. "
+            "These replicas consume no local device memory.",
+            exempt,
+        )
+
+    capacities: dict[int, int] = {}
+    for d in demands:
+        for dev in d.device_ids:
+            if dev not in capacities:
+                capacities[dev] = int(device_total_memory(dev))
+
+    return evaluate(
+        demands,
+        capacities,
+        external_reserve_bytes=external_reserve_bytes,
+        safety_margin_bytes=safety_margin_bytes,
+    )
+
+
+def _diffusion_utilization(replica: Any) -> float | None:
+    """Best-effort gpu_memory_utilization for a diffusion replica from raw args."""
+    stage_cfg = getattr(replica, "stage_cfg", None)
+    engine_args = getattr(stage_cfg, "engine_args", None)
+    if isinstance(engine_args, dict):
+        util = engine_args.get("gpu_memory_utilization")
+        if util is not None:
+            return float(util)
+    return None

vllm_omni/engine/stage_engine_core_proc.py

@@ -41,6 +41,35 @@ logger = init_logger(__name__)
 _SIGNAL_EXIT_BASE = 128
 
 
+def _install_phase_locks(kwargs: dict[str, Any], local_dp_rank: int) -> None:
+    """Wrap ``kwargs["executor_class"]`` with the SH/EX phase-lock guard.
+
+    Fail-closed: ``parallel_stage_init`` promises that every memory-mutating
+    init phase in this child runs under the per-device locks, so a missing
+    ``vllm_config`` or ``executor_class`` must abort the launch rather than
+    silently proceed with an unguarded parallel initialization.
+    """
+    from vllm_omni.engine.stage_phase_lock import (
+        DevicePhaseLock,
+        wrap_executor_with_phase_locks,
+    )
+
+    missing = [key for key in ("vllm_config", "executor_class") if kwargs.get(key) is None]
+    if missing:
+        raise RuntimeError(
+            f"parallel_stage_init is enabled but EngineCore kwargs are missing {missing}, "
+            "so the SH/EX phase-lock guard cannot be installed. Refusing to run an "
+            "unguarded parallel initialization; fix the launch plumbing or disable "
+            "parallel_stage_init."
+        )
+    locker = DevicePhaseLock.from_child(kwargs["vllm_config"], local_dp_rank)
+    kwargs["executor_class"] = wrap_executor_with_phase_locks(kwargs["executor_class"], locker)
+    logger.info(
+        "[StageEngineCoreProc] parallel_stage_init: SH/EX phase locks on devices %s",
+        locker.device_ids,
+    )
+
+
 def _signal_exit_code(signum: int) -> int:
     """Return the conventional process exit code for signal-driven exits."""
     return _SIGNAL_EXIT_BASE + signum
@@ -69,6 +98,7 @@ class StageEngineCoreProc(EngineCoreProc):
         omni_coordinator_address: str | None = None,
         omni_stage_id: int | None = None,
         omni_replica_id: int = 0,
+        omni_parallel_stage_init: bool = False,
         **kwargs: Any,
     ) -> None:
         """Launch StageEngineCoreProc busy loop in background process.
@@ -139,6 +169,15 @@ class StageEngineCoreProc(EngineCoreProc):
             # default sampling extra_args.
             maybe_apply_cfg_scheduler_patches(kwargs.get("vllm_config"))
 
+            # When parallel stage init is enabled, wrap this driver's executor
+            # so its memory-mutating phases (load / KV alloc / capture) hold a
+            # per-device LOCK_SH and its profiling measurement holds LOCK_EX.
+            # The wrapper must be installed here (in the engine-core child):
+            # phase boundaries live inside EngineCore.__init__, invisible to the
+            # orchestrator. See vllm_omni.engine.stage_phase_lock.
+            if omni_parallel_stage_init:
+                _install_phase_locks(kwargs, local_dp_rank)
+
             engine_core = StageEngineCoreProc(
                 *args,
                 engine_index=dp_rank,

vllm_omni/engine/stage_engine_core_proc_manager.py

@@ -75,6 +75,7 @@ class StageEngineCoreProcManager(CoreEngineProcManager):
         omni_replica_base_id: int = 0,
         client_handshake_address: str | None = None,
         tensor_queue: Queue | None = None,
+        omni_parallel_stage_init: bool = False,
     ) -> None:
         # NOTE: we intentionally do not call ``super().__init__`` — the
         # parent's body hardcodes the wrong target. We re-implement it here
@@ -92,6 +93,7 @@ class StageEngineCoreProcManager(CoreEngineProcManager):
             "tensor_queue": tensor_queue,
             "omni_stage_id": int(omni_stage_id),
             "omni_coordinator_address": omni_coordinator_address,
+            "omni_parallel_stage_init": bool(omni_parallel_stage_init),
         }
 
         if client_handshake_address:

vllm_omni/engine/stage_engine_startup.py

@@ -974,6 +974,7 @@ def _launch_omni_core_engines(
     omni_coordinator_address: str | None = None,
     stage_visible_devices: str | None = None,
     spawn_device_lock: threading.Lock | None = None,
+    omni_parallel_stage_init: bool = False,
 ) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]:
     """Launch local engine cores using the omni registration flow.
 
@@ -1068,8 +1069,20 @@ def _launch_omni_core_engines(
                         omni_stage_id=stage_id,
                         omni_coordinator_address=omni_coordinator_address,
                         omni_replica_base_id=replica_id,
+                        omni_parallel_stage_init=omni_parallel_stage_init,
                     )
             else:
+                if omni_parallel_stage_init:
+                    # The upstream CoreEngineProcManager targets EngineCoreProc.run_engine_core,
+                    # which does NOT install the SH/EX phase-lock executor wrapper. Since the
+                    # orchestrator skips its own device lock when parallel_stage_init is on,
+                    # running here would leave same-device init unguarded. Fail closed.
+                    raise RuntimeError(
+                        f"parallel_stage_init is enabled but stage {stage_id} would launch via the "
+                        "upstream CoreEngineProcManager (no omni coordinator), which cannot install "
+                        "the phase-lock guard. Disable parallel_stage_init for this deployment or "
+                        "run with the omni coordinator."
+                    )
                 with scoped_spawn_device_env(
                     stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id
                 ):
@@ -1120,6 +1133,7 @@ def launch_stage_replica(
     omni_coordinator_address: str | None = None,
     stage_visible_devices: str | None = None,
     spawn_device_lock: threading.Lock | None = None,
+    omni_parallel_stage_init: bool = False,
 ) -> Iterator[StageReplicaResources]:
     """Launch a local LLM stage replica.
 
@@ -1141,6 +1155,7 @@ def launch_stage_replica(
             omni_coordinator_address=omni_coordinator_address,
             stage_visible_devices=stage_visible_devices,
             spawn_device_lock=spawn_device_lock,
+            omni_parallel_stage_init=omni_parallel_stage_init,
         ) as resources:
             engine_manager, coordinator, addresses = resources
             yield StageReplicaResources(
@@ -1170,6 +1185,7 @@ def launch_stage_replica(
             omni_stage_id=stage_id,
             omni_coordinator_address=omni_coordinator_address,
             omni_replica_base_id=replica_id,
+            omni_parallel_stage_init=omni_parallel_stage_init,
         )
 
     with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket:

vllm_omni/engine/stage_init_utils.py

@@ -17,6 +17,7 @@ import importlib
 import json
 import multiprocessing as mp
 import os
+import tempfile
 import time
 from collections.abc import Callable, Collection, Generator, Mapping, Sequence
 from contextlib import contextmanager
@@ -1562,43 +1563,108 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
     return input_processor
 
 
-def _cleanup_stale_lock_if_dead(lock_file: str) -> bool:
-    """If *lock_file* exists and its recorded PID is dead, unlink the file.
+def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str:
+    """Return the per-physical-device initialization lock file path.
 
-    Returns ``True`` if the stale lock was cleaned up (caller should retry),
-    ``False`` otherwise (lock holder appears alive, or file could not be read).
+    Shared by the orchestrator-side ``acquire_device_locks`` (legacy full-init
+    ``LOCK_EX``) and the engine-core-side ``DevicePhaseLock`` (parallel-stage-init
+    SH/EX phase locks) so both coordinate on the *same* file per device.
+
+    That coordination only holds while the **inode** is stable, so nothing may
+    unlink this path: two holders of the same pathname on different inodes do not
+    conflict. The PID written into the file is diagnostic only.
+    """
+    return os.path.join(lock_dir, f"vllm_omni_device_{device_id}_init.lock")
+
+
+def _open_existing_lock_file(lock_file: str) -> tuple[int, bool] | None:
+    """Open an already-existing lock file; ``None`` if it does not exist yet.
+
+    Prefers ``O_RDWR`` so the holder can stamp its PID, but falls back to
+    ``O_RDONLY``: ``flock`` locks attach to the open file description and, unlike
+    POSIX ``fcntl`` record locks, require no write access, so a read-only
+    descriptor still provides full mutual exclusion.
     """
     try:
-        with open(lock_file) as fh:
-            content = fh.read().strip()
-        if not content:
-            return False
-        pid = int(content)
-    except (OSError, ValueError):
-        return False
-
-    # Check whether the PID is still alive.
-    try:
-        os.kill(pid, 0)
-    except ProcessLookupError:
-        # PID does not exist — stale lock.
-        logger.info(
-            "Removing stale device lock %s (PID %s is dead)",
-            lock_file,
-            pid,
-        )
-        try:
-            os.unlink(lock_file)
-            return True
-        except OSError:
-            logger.debug("Failed to unlink stale lock %s", lock_file)
-            return False
+        return os.open(lock_file, os.O_RDWR), True
+    except FileNotFoundError:
+        return None
     except PermissionError:
-        # PID exists but we cannot signal it (different user) — treat as alive.
-        return False
+        # Created by another user on this shared machine; read-only still locks.
+        pass
+    try:
+        return os.open(lock_file, os.O_RDONLY), False
+    except FileNotFoundError:
+        return None
+    except PermissionError as exc:
+        raise PermissionError(
+            f"Device init lock {lock_file} exists but is not readable by this user "
+            f"({exc}). It coordinates GPU initialization across users, so it must stay "
+            "readable by all of them; have its owner remove it or chmod it to 0644."
+        ) from exc
 
-    # PID is alive — legitimate lock holder.
-    return False
+
+def open_device_lock_file(lock_file: str) -> tuple[int, bool]:
+    """Open a per-device lock file, tolerating one created by another user.
+
+    Returns ``(fd, writable)``.
+
+    These lock files coordinate *across* users -- two people running on the same
+    physical GPU must contend on the same file -- so every user has to be able to
+    open whichever one exists. Two things make that awkward:
+
+    * whoever creates it owns it, so later users may only get read access, which
+      ``flock`` is perfectly happy with; and
+    * the mode passed to ``os.open`` is filtered by the creator's ``umask``, so
+      under 0027 or 0077 the file would land 0640 or 0600 and lock everyone else
+      out entirely.
+
+    Creating it therefore stages a temporary file, widens it with ``fchmod``, and
+    publishes it with an atomic ``os.link``. Creating in place with ``O_EXCL`` and
+    widening afterwards would briefly expose the file at the umask-filtered mode,
+    and a concurrent user opening it in that window would be locked out -- the
+    very failure this avoids. ``link`` also fails cleanly if another process wins
+    the race, which keeps the "first creator wins" inode stable.
+    """
+    existing = _open_existing_lock_file(lock_file)
+    if existing is not None:
+        return existing
+
+    lock_dir = os.path.dirname(lock_file) or "."
+    tmp_fd, tmp_path = tempfile.mkstemp(dir=lock_dir, prefix=".vllm_omni_device_lock_")
+    published = False
+    try:
+        os.fchmod(tmp_fd, 0o644)
+        try:
+            os.link(tmp_path, lock_file)
+        except FileExistsError:
+            pass  # another process published first; fall through and open theirs
+        else:
+            published = True
+            return tmp_fd, True
+    finally:
+        try:
+            os.unlink(tmp_path)
+        except OSError:
+            logger.debug("Could not remove staged device lock %s", tmp_path)
+        if not published:
+            os.close(tmp_fd)
+
+    existing = _open_existing_lock_file(lock_file)
+    if existing is not None:
+        return existing
+    raise PermissionError(f"Device init lock {lock_file} could not be created or opened")
+
+
+def record_lock_holder_pid(fd: int, writable: bool) -> None:
+    """Stamp the holder's PID into an open lock file (diagnostic only)."""
+    if not writable:
+        return
+    try:
+        os.ftruncate(fd, 0)
+        os.write(fd, f"{os.getpid()}\n".encode())
+    except OSError:
+        pass
 
 
 def acquire_device_locks(
@@ -1677,28 +1743,31 @@ def acquire_device_locks(
         # Acquire locks
         wait_start = time.time()
         for device_id in devices_to_lock:
-            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+            lock_file = device_init_lock_path(device_id)
             lock_acquired = False
-            already_cleaned_stale = False  # only try stale cleanup once per device
 
             while not lock_acquired:
                 try:
-                    lock_fd = os.open(lock_file, os.O_CREAT | os.O_RDWR, 0o644)
+                    lock_fd, lock_writable = open_device_lock_file(lock_file)
                     try:
                         fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                        os.ftruncate(lock_fd, 0)
-                        os.write(lock_fd, f"{os.getpid()}\n".encode())
-                        os.fsync(lock_fd)
+                        record_lock_holder_pid(lock_fd, lock_writable)
                         lock_acquired = True
                         lock_fds.append(lock_fd)
                         logger.debug("Acquired exclusive lock for device %s", device_id)
                     except BlockingIOError:
                         os.close(lock_fd)
-                        # Detect and clean stale locks from dead processes.
-                        if not already_cleaned_stale:
-                            already_cleaned_stale = True
-                            if _cleanup_stale_lock_if_dead(lock_file):
-                                continue  # retry flock immediately
+                        # NOTE: no stale-lock cleanup here. ``flock`` is released
+                        # by the kernel when the holder exits (SIGKILL included),
+                        # so a dead holder never keeps this lock. Unlinking the
+                        # path on a dead *recorded* PID was actively unsafe: the
+                        # PID is written by every holder, and under the
+                        # parallel-stage-init SH/EX protocol several LOCK_SH
+                        # holders overwrite it. Unlinking while one still held the
+                        # old inode let a contender create a fresh file with the
+                        # same name and take LOCK_EX on it, which then conflicted
+                        # with nobody -- breaking the interoperability between the
+                        # legacy full-init lock and the phase locks.
                         if time.time() - wait_start > stage_init_timeout:
                             logger.warning(
                                 "Timeout waiting for device %s initialization lock, proceeding anyway",
@@ -1707,8 +1776,10 @@ def acquire_device_locks(
                             break
                         time.sleep(0.01)
                 except OSError as e:
-                    logger.debug(
-                        "Failed to acquire lock for device %s: %s, continuing anyway",
+                    logger.warning(
+                        "Failed to acquire lock for device %s: %s. Continuing WITHOUT "
+                        "device-init serialization; concurrent initialization on this "
+                        "device may incorrectly measure available memory.",
                         device_id,
                         e,
                     )

vllm_omni/engine/stage_phase_lock.py

@@ -0,0 +1,242 @@
+"""Per-device shared/exclusive init phase locks for parallel stage init.
+
+When ``VllmOmniOrchestratorConfig.parallel_stage_init`` is enabled, several
+stage engine-core processes may initialize on the same physical GPU at the same
+time. Correctness then rests on two mechanisms (see
+``rfc_parallel_stage_init.zh.md``):
+
+  * **Admission** (``stage_admission``) proves, before any stage launches, that
+    every physical device's summed budget fits — this is the hard OOM backstop.
+  * **SH/EX phase locks** (this module) make each memory *measurement* quiescent:
+    memory-mutating phases (weight load, KV allocation, CUDA-graph capture) take
+    a per-device shared (``LOCK_SH``) lock; the profiling measurement inside
+    ``determine_available_memory`` takes an exclusive (``LOCK_EX``) lock. The
+    exclusive holder waits for in-flight mutations to drain and blocks new ones
+    for its (seconds-long) window, so the device-level profiling number each
+    stage computes stays within its admitted budget.
+
+The lock is held by the engine-core **driver** process (one per stage / per local
+DP rank) and wraps the driver-side collective calls, so a single acquisition
+covers every worker rank with no per-rank locking (which would deadlock on
+inverted lock order).
+Acquisition is always over the driver's whole device slice in **sorted global
+order** so different drivers cannot deadlock, and it is **fail-closed**: a lock
+timeout raises rather than silently proceeding.
+
+The core locking logic here intentionally depends only on ``os``/``fcntl`` so it
+is importable and unit-testable without a GPU; torch/platform imports are lazy.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import fcntl
+import os
+import time
+from collections.abc import Iterator
+
+from vllm.logger import init_logger
+
+from vllm_omni.engine.stage_init_utils import (
+    device_init_lock_path,
+    open_device_lock_file,
+    record_lock_holder_pid,
+)
+
+logger = init_logger(__name__)
+
+# Default upper bound on how long a single phase will wait for its device locks
+# before failing closed. Generous: a peer's exclusive profile window is seconds,
+# but a peer's whole shared load/capture span can be much longer.
+_DEFAULT_LOCK_TIMEOUT_S = 900.0
+
+
+class DeviceLockTimeoutError(RuntimeError):
+    """Raised (fail-closed) when a phase cannot acquire its device locks in time."""
+
+
+def resolve_driver_device_ids(vllm_config, local_dp_rank: int = 0) -> list[int]:
+    """Physical GPU ids this engine-core driver owns, as a sorted list.
+
+    A driver (one local vLLM DP rank) owns a contiguous slice of the visible
+    devices of size ``TP*PP*PCP*SP*CFG`` (DP excluded — each DP rank is its own
+    driver). Visible ids come from the device-control env var the parent set on
+    the child (e.g. ``CUDA_VISIBLE_DEVICES``); when unset, all devices are visible.
+    """
+    from vllm_omni.platforms import current_omni_platform
+
+    env_var = current_omni_platform.device_control_env_var
+    visible = os.environ.get(env_var) if env_var else None
+    physical: list[int]
+    if visible:
+        physical = []
+        for tok in visible.split(","):
+            tok = tok.strip()
+            if not tok:
+                continue
+            try:
+                physical.append(int(tok))
+            except ValueError:
+                # Non-integer visibility (UUID / MIG) can't be mapped to a lock
+                # id; fall back to positional ordinals below.
+                physical = []
+                break
+    else:
+        physical = []
+    if not physical:
+        physical = list(range(current_omni_platform.get_device_count()))
+
+    pc = vllm_config.parallel_config
+    devices_per_dp = (
+        int(getattr(pc, "tensor_parallel_size", 1))
+        * int(getattr(pc, "pipeline_parallel_size", 1))
+        * int(getattr(pc, "prefill_context_parallel_size", 1))
+        * int(getattr(pc, "sequence_parallel_size", 1))
+        * int(getattr(pc, "cfg_parallel_size", 1))
+    )
+    devices_per_dp = max(1, devices_per_dp)
+    start = int(local_dp_rank) * devices_per_dp
+    slice_ = physical[start : start + devices_per_dp]
+    if not slice_:
+        # Defensive: fall back to the whole visible set rather than lock nothing.
+        slice_ = physical
+    return sorted(slice_)
+
+
+class DevicePhaseLock:
+    """Sorted, fail-closed, whole-set SH/EX flock over a driver's physical devices."""
+
+    def __init__(
+        self,
+        device_ids: list[int],
+        *,
+        timeout_s: float = _DEFAULT_LOCK_TIMEOUT_S,
+        lock_dir: str = "/tmp",
+    ) -> None:
+        self._device_ids = sorted(device_ids)
+        self._timeout_s = timeout_s
+        self._lock_dir = lock_dir
+
+    @classmethod
+    def from_child(
+        cls,
+        vllm_config,
+        local_dp_rank: int = 0,
+        *,
+        timeout_s: float = _DEFAULT_LOCK_TIMEOUT_S,
+    ) -> DevicePhaseLock:
+        """Build the lock for the current engine-core child from its config/env."""
+        device_ids = resolve_driver_device_ids(vllm_config, local_dp_rank)
+        return cls(device_ids, timeout_s=timeout_s)
+
+    @property
+    def device_ids(self) -> list[int]:
+        return list(self._device_ids)
+
+    # ---- acquisition -------------------------------------------------------
+
+    def _acquire(self, mode: int) -> list[int]:
+        """Acquire *mode* (LOCK_SH|LOCK_EX) on every device in sorted order.
+
+        Fail-closed: on timeout, releases whatever was acquired and raises
+        ``DeviceLockTimeoutError``. Sorted global order makes cross-driver deadlock
+        impossible.
+        """
+        acquired: list[int] = []
+        deadline = time.monotonic() + self._timeout_s
+        try:
+            for device_id in self._device_ids:
+                fd, writable = open_device_lock_file(device_init_lock_path(device_id, self._lock_dir))
+                while True:
+                    try:
+                        fcntl.flock(fd, mode | fcntl.LOCK_NB)
+                        break
+                    except BlockingIOError:
+                        if time.monotonic() > deadline:
+                            os.close(fd)
+                            raise DeviceLockTimeoutError(
+                                f"Timed out after {self._timeout_s:.0f}s acquiring "
+                                f"{'EX' if mode == fcntl.LOCK_EX else 'SH'} lock on device {device_id}"
+                            )
+                        time.sleep(0.01)
+                record_lock_holder_pid(fd, writable)
+                acquired.append(fd)
+            return acquired
+        except BaseException:
+            _release_fds(acquired)
+            raise
+
+    @contextlib.contextmanager
+    def shared(self) -> Iterator[None]:
+        """Hold ``LOCK_SH`` on the device slice for a memory-mutating phase."""
+        if not self._device_ids:
+            yield
+            return
+        fds = self._acquire(fcntl.LOCK_SH)
+        try:
+            yield
+        finally:
+            _release_fds(fds)
+
+    @contextlib.contextmanager
+    def exclusive(self) -> Iterator[None]:
+        """Hold ``LOCK_EX`` on the device slice for the profiling measurement."""
+        if not self._device_ids:
+            yield
+            return
+        fds = self._acquire(fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            _release_fds(fds)
+
+
+def _release_fds(fds: list[int]) -> None:
+    for fd in fds:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_UN)
+            os.close(fd)
+        except (OSError, ValueError):
+            pass
+
+
+def wrap_executor_with_phase_locks(executor_class, locker: DevicePhaseLock):
+    """Return a subclass of *executor_class* that brackets init phases with locks.
+
+    * ``__init__`` (executor construction → weight load) under ``shared()``.
+    * ``determine_available_memory`` (profiling) under ``exclusive()``.
+    * ``initialize_from_config`` (KV cache allocation) under ``shared()``.
+    * ``compile_or_warm_up_model`` (kernel warmup + CUDA-graph capture) under
+      ``shared()``.
+
+    ``compile_or_warm_up_model`` needs its own bracket: upstream used to fuse it
+    into ``Executor.initialize_from_config``, but since vLLM v0.20.0 (vllm#39240)
+    ``EngineCore._initialize_kv_caches`` calls the two separately. Leaving it
+    unwrapped would let one stage capture graphs — mutating device memory — while
+    a peer holds the exclusive lock for its profiling measurement, which is
+    exactly the interference this protocol exists to prevent.
+
+    Works uniformly for any executor (UniProc/MultiProc): the lock is taken in
+    the driver process that instantiates the executor and issues its collectives.
+    """
+
+    class _PhaseLockedExecutor(executor_class):
+        def __init__(self, vllm_config):
+            with locker.shared():
+                super().__init__(vllm_config)
+
+        def determine_available_memory(self):
+            with locker.exclusive():
+                return super().determine_available_memory()
+
+        def initialize_from_config(self, kv_cache_configs):
+            with locker.shared():
+                return super().initialize_from_config(kv_cache_configs)
+
+        def compile_or_warm_up_model(self):
+            with locker.shared():
+                return super().compile_or_warm_up_model()
+
+    _PhaseLockedExecutor.__name__ = f"PhaseLocked{executor_class.__name__}"
+    _PhaseLockedExecutor.__qualname__ = _PhaseLockedExecutor.__name__
+    return _PhaseLockedExecutor

vllm_omni/engine/stage_runtime.py

@@ -9,6 +9,7 @@ import concurrent.futures
 import copy
 import os
 import threading
+import time
 from collections.abc import Callable, Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -125,6 +126,7 @@ class StageRuntime:
         stage_init_timeout: int,
         async_chunk: bool,
         tokenizer: str | None = None,
+        parallel_stage_init: bool = False,
         log_stats: bool = False,
     ) -> None:
         self._stage_configs = stage_configs
@@ -133,6 +135,11 @@ class StageRuntime:
         self._stage_init_timeout = stage_init_timeout
         self._async_chunk = async_chunk
         self._tokenizer = tokenizer
+        # When True, same-device stages initialize concurrently, coordinated by
+        # pre-launch admission control + the engine-core-held SH/EX device locks
+        # (see VllmOmniOrchestratorConfig.parallel_stage_init). Default False
+        # keeps the legacy per-device LOCK_EX serialization.
+        self._parallel_stage_init = parallel_stage_init
         self._log_stats = log_stats
         self._num_stages = len(stage_configs)
 
@@ -450,6 +457,13 @@ class StageRuntime:
         init_state_lock = threading.Lock()
         self._init_visible_devices_baseline = os.environ.get(current_omni_platform.device_control_env_var)
 
+        if self._parallel_stage_init:
+            # Refuse executor backends the phase locks cannot cover, then prove
+            # every physical device's summed budget fits BEFORE any stage spawns.
+            # Both fail fast here rather than OOM at allocation time.
+            self._reject_unguardable_executors(stage_plans)
+            self._run_stage_admission(stage_plans)
+
         init_groups: dict[str, list[tuple[int, ReplicaInitPlan]]] = {}
         for plan in stage_plans:
             for replica in plan.replicas:
@@ -506,8 +520,154 @@ class StageRuntime:
 
         return initialized_clients_by_stage
 
+    def _reject_unguardable_executors(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None:
+        """Fail closed for executors the SH/EX phase locks cannot actually guard.
+
+        ``DevicePhaseLock`` is a node-local ``flock`` over files on the engine-core
+        driver's host, and the admission ledger is built only from this host's
+        devices. A Ray executor may place the workers that load weights, profile,
+        allocate KV cache and capture graphs on **other** nodes, which participate
+        in neither mechanism. Running there would take the parallel path — the
+        parent skips its own device lock when ``parallel_stage_init`` is on —
+        while only advertising the protection, so refuse before anything spawns.
+
+        Ray detection uses both the executor's ``uses_ray`` class attribute and
+        ``ParallelConfig.use_ray``. Multi-node multiprocessing/external-launcher
+        configurations are detected through ``ParallelConfig.nnodes``. Both are
+        refused because some workers can live outside the driver's host.
+
+        Remote replicas are skipped: they are launched by the runtime owning their
+        node and are already ``ADMISSION_EXEMPT`` in this host's ledger.
+        """
+        offenders: list[str] = []
+        for plan in stage_plans:
+            for replica in plan.replicas:
+                if replica.launch_mode == "remote":
+                    continue
+                executor_class = replica.executor_class
+                parallel_config = getattr(replica.stage_vllm_config, "parallel_config", None)
+                uses_ray = bool(getattr(executor_class, "uses_ray", False)) or bool(
+                    getattr(parallel_config, "use_ray", False)
+                )
+                nnodes = int(getattr(parallel_config, "nnodes", 1) or 1)
+                reasons: list[str] = []
+                if uses_ray:
+                    reasons.append("Ray-backed")
+                if nnodes > 1:
+                    reasons.append(f"multi-node (nnodes={nnodes})")
+                if reasons:
+                    executor_name = getattr(executor_class, "__name__", "unknown executor")
+                    offenders.append(
+                        f"stage{replica.metadata.stage_id}/replica{replica.replica_id} "
+                        f"({executor_name}; {', '.join(reasons)})"
+                    )
+        if offenders:
+            raise RuntimeError(
+                "parallel_stage_init cannot guard Ray-backed or multi-node executors: the SH/EX "
+                "phase locks are node-local file locks and the admission ledger only accounts "
+                "for this host's devices, so workers placed on other nodes would initialize unprotected. "
+                f"Offending replicas: {', '.join(offenders)}. Disable parallel_stage_init for "
+                "this deployment, or use a single-node non-Ray executor backend."
+            )
+
+    def _run_stage_admission(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None:
+        """Pre-launch per-device admission for parallel stage init (fail-fast)."""
+        from vllm_omni.engine.stage_admission import (
+            ADMISSION_EXEMPT,
+            AdmissionExempt,
+            StageAdmissionError,
+            check_admission,
+        )
+
+        def _resolve(replica: ReplicaInitPlan) -> list[int] | AdmissionExempt | None:
+            if replica.launch_mode == "remote":
+                # Remote replicas consume a remote node's memory and the ledger
+                # is per-LOCAL-device. Remote LLM plans already carry
+                # runtime_cfg=None, but remote diffusion replicas keep their
+                # cfg — mark both exempt explicitly. Only this exemption may
+                # skip the ledger: a LOCAL replica that resolves to None below
+                # fails admission (fail-closed in check_admission).
+                return ADMISSION_EXEMPT
+            resolved = self._resolve_replica_physical_devices(replica.metadata.stage_id, replica.metadata.runtime_cfg)
+            if not resolved:
+                return None
+            ids: list[int] = []
+            for tok in str(resolved).split(","):
+                tok = tok.strip()
+                if not tok:
+                    continue
+                try:
+                    ids.append(int(tok))
+                except ValueError:
+                    # Non-integer visibility (UUID / MIG) can't be reasoned about.
+                    return None
+            return ids or None
+
+        def _visible_ordinal(physical_id: int) -> int:
+            """Translate a physical device id to this process's visible ordinal.
+
+            The ledger is keyed by *physical* id (resolved against the visibility
+            baseline captured before any stage narrowed the env), but
+            ``get_device_total_memory`` and ``torch.cuda.get_device_properties``
+            interpret their argument as an ordinal in the **current** process's
+            visible-device namespace. Passing the physical id straight through
+            raises an invalid-ordinal error under a restricted
+            ``CUDA_VISIBLE_DEVICES`` and, when visibility is reordered, silently
+            reads a different GPU — which yields a wrong capacity, and therefore
+            an unsafe admission result, on heterogeneous devices.
+            """
+            baseline = self._init_visible_devices_baseline
+            if not baseline:
+                # No device-control env captured: ordinals are physical ids.
+                return physical_id
+            for ordinal, tok in enumerate(baseline.split(",")):
+                tok = tok.strip()
+                if not tok:
+                    continue
+                try:
+                    candidate = int(tok)
+                except ValueError:
+                    # UUID/MIG entries cannot be converted to physical integer
+                    # ids, but they still occupy an ordinal in the visible list.
+                    continue
+                if candidate == physical_id:
+                    return ordinal
+            raise StageAdmissionError(
+                f"Physical device {physical_id} is not in the visibility baseline "
+                f"{baseline!r} captured at init, so its capacity cannot be measured "
+                "from this process. Refusing to admit against an unknown device."
+            )
+
+        def _total_memory(device_id: int) -> int:
+            ordinal = _visible_ordinal(device_id)
+            try:
+                mem = current_omni_platform.get_device_total_memory(ordinal)
+                if mem:
+                    return int(mem)
+            except (NotImplementedError, AttributeError, TypeError):
+                pass
+            import torch
+
+            return int(torch.cuda.get_device_properties(ordinal).total_memory)
+
+        check_admission(
+            stage_plans,
+            resolve_physical_devices=_resolve,
+            device_total_memory=_total_memory,
+        )
+
     def _replica_init_group_key(self, replica: ReplicaInitPlan) -> str:
-        """Return the scheduling group used during replica initialization."""
+        """Return the scheduling group used during replica initialization.
+
+        Replicas sharing a group initialize sequentially; different groups
+        initialize in parallel threads. Local LLM replicas are keyed by their
+        **resolved canonical physical device set** so stages on different
+        physical GPUs land in different groups (parallel) while stages sharing a
+        device stay in one group (serialized, then further guarded by the
+        per-device ``LOCK_EX`` file lock). Keying on the raw ``runtime.devices``
+        config value instead would fail to parallelize logically-distinct stages
+        that map to different physical GPUs.
+        """
         if replica.launch_mode == "local" and replica.metadata.stage_type == "diffusion":
             # Local diffusion process spawning must stay on the orchestrator
             # thread. Keep all local diffusion replicas in one sequential group.
@@ -515,9 +675,28 @@ class StageRuntime:
         if replica.launch_mode == "remote":
             return f"remote:{replica.metadata.stage_id}:{replica.replica_id}"
 
+        physical_devices = self._resolve_replica_physical_devices(
+            replica.metadata.stage_id,
+            replica.metadata.runtime_cfg,
+        )
+        if self._parallel_stage_init:
+            # Same-device concurrency is coordinated by the engine-core SH/EX
+            # device locks + pre-launch admission, so give every replica its own
+            # group to let them all initialize in parallel.
+            return f"parallel:{replica.metadata.stage_id}:{replica.replica_id}"
+
         runtime_cfg = replica.metadata.runtime_cfg or {}
-        devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
-        return f"device:{devices}"
+        raw_devices = (
+            runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
+        )
+        if str(raw_devices) != str(physical_devices):
+            logger.debug(
+                "[stage_init] Stage-%s init-group key: raw devices=%s -> resolved physical=%s",
+                replica.metadata.stage_id,
+                raw_devices,
+                physical_devices,
+            )
+        return f"device:{physical_devices}"
 
     def _initialize_replica(
         self,
@@ -566,28 +745,76 @@ class StageRuntime:
                 raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing executor_class")
             if plan.engine_args_dict is None:
                 raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing engine args")
-            with self._scoped_spawn_device_env(physical_devices):
-                lock_fds = acquire_device_locks(
+            # G3 device locks: in the default (serial) path the parent holds a
+            # per-device LOCK_EX across the whole child init. When parallel stage
+            # init is enabled the engine-core child takes phased SH/EX locks
+            # itself (see stage_phase_lock), so the parent must NOT also hold the
+            # lock here — it would self-deadlock while waiting for the child's
+            # READY handshake.
+            if not self._parallel_stage_init:
+                g3_start = time.perf_counter()
+                with self._scoped_spawn_device_env(physical_devices):
+                    lock_fds = acquire_device_locks(
+                        plan.metadata.stage_id,
+                        plan.engine_args_dict,
+                        stage_init_timeout,
+                    )
+                logger.debug(
+                    "[stage_init] Stage-%s G3 device-lock acquire took %.3fs",
                     plan.metadata.stage_id,
-                    plan.engine_args_dict,
-                    stage_init_timeout,
+                    time.perf_counter() - g3_start,
+                )
+
+            launch_cm = launch_stage_replica(
+                vllm_config=vllm_config,
+                executor_class=executor_class,
+                log_stats=self._log_stats,
+                stage_id=plan.metadata.stage_id,
+                replica_id=plan.replica_id,
+                stage_config=plan.stage_cfg,
+                omni_master_server=self._get_omni_master_server(),
+                omni_coordinator_address=self._get_coordinator_address(),
+                stage_visible_devices=physical_devices,
+                spawn_device_lock=self._spawn_device_lock,
+                omni_parallel_stage_init=self._parallel_stage_init,
+            )
+            # G2 launch lock serializes engine-core *spawning* across replicas
+            # (ZMQ port-allocation races + simultaneous CUDA context init).
+            #   * Default path: hold it across the whole launch context manager,
+            #     whose __exit__ waits for READY — this serializes the full child
+            #     init (spawn + load + profile + KV + capture).
+            #   * Parallel path: hold it only around the spawn (__enter__); run
+            #     the READY-wait (__exit__) outside the lock so replicas init
+            #     concurrently, coordinated by the child SH/EX device locks.
+            if self._parallel_stage_init:
+                # The per-stage runtime.env overlay mutates os.environ
+                # (process-global), so it must be applied under the launch lock
+                # and only needs to cover the spawn (__enter__) — children
+                # inherit the env at spawn time; the READY-wait needs no env.
+                g2_start = time.perf_counter()
+                with self._replica_launch_lock:
+                    with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
+                        resources = launch_cm.__enter__()
+                g2_spawned = time.perf_counter()
+                launch_cm.__exit__(None, None, None)
+                logger.debug(
+                    "[stage_init] Stage-%s G2 spawn(locked)=%.3fs, READY(unlocked)=%.3fs",
+                    plan.metadata.stage_id,
+                    g2_spawned - g2_start,
+                    time.perf_counter() - g2_spawned,
+                )
+            else:
+                g2_start = time.perf_counter()
+                with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
+                    g2_locked = time.perf_counter()
+                    with launch_cm as resources:
+                        pass
+                logger.debug(
+                    "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn+READY=%.3fs",
+                    plan.metadata.stage_id,
+                    g2_locked - g2_start,
+                    time.perf_counter() - g2_locked,
                 )
-            # Serialize engine-core spawning across all LLM replicas to avoid
-            # ZMQ port-allocation races and simultaneous CUDA context init.
-            with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
-                with launch_stage_replica(
-                    vllm_config=vllm_config,
-                    executor_class=executor_class,
-                    log_stats=self._log_stats,
-                    stage_id=plan.metadata.stage_id,
-                    replica_id=plan.replica_id,
-                    stage_config=plan.stage_cfg,
-                    omni_master_server=self._get_omni_master_server(),
-                    omni_coordinator_address=self._get_coordinator_address(),
-                    stage_visible_devices=physical_devices,
-                    spawn_device_lock=self._spawn_device_lock,
-                ) as resources:
-                    pass
 
             logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id)
             if resources is None:
@@ -773,6 +1000,7 @@ class DistStageRuntime(StageRuntime):
         omni_heartbeat_timeout: float = 30.0,
         omni_lb_policy: str = "random",
         request_queue: janus.Queue[EngineQueueMessage] | None = None,
+        parallel_stage_init: bool = False,
     ) -> None:
         super().__init__(
             stage_configs=stage_configs,
@@ -781,6 +1009,7 @@ class DistStageRuntime(StageRuntime):
             stage_init_timeout=stage_init_timeout,
             async_chunk=async_chunk,
             tokenizer=tokenizer,
+            parallel_stage_init=parallel_stage_init,
             log_stats=log_stats,
         )
         self._single_stage_id_filter = single_stage_id_filter
@@ -1101,6 +1330,7 @@ def create_stage_runtime(
     stage_init_timeout: int,
     async_chunk: bool,
     tokenizer: str | None = None,
+    parallel_stage_init: bool = False,
     # Distributed-only params:
     single_stage_id_filter: int | None = None,
     omni_master_address: str | None = None,
@@ -1122,6 +1352,7 @@ def create_stage_runtime(
             stage_init_timeout=stage_init_timeout,
             async_chunk=async_chunk,
             tokenizer=tokenizer,
+            parallel_stage_init=parallel_stage_init,
             log_stats=log_stats,
             single_stage_id_filter=single_stage_id_filter,
             omni_master_address=omni_master_address,
@@ -1138,5 +1369,6 @@ def create_stage_runtime(
         stage_init_timeout=stage_init_timeout,
         async_chunk=async_chunk,
         tokenizer=tokenizer,
+        parallel_stage_init=parallel_stage_init,
         log_stats=log_stats,
     )

vllm_omni/worker/base.py

@@ -1,4 +1,4 @@
-"""Base worker class for vLLM-Omni with process-scoped GPU memory accounting."""
+"""Base worker class for vLLM-Omni with device-level GPU memory profiling."""
 
 from __future__ import annotations
 
@@ -17,22 +17,20 @@ from vllm_omni.diffusion.data import (
     OmniSleepTask,
     OmniWakeTask,
 )
-from vllm_omni.entrypoints.utils import detect_pid_host
 from vllm_omni.platforms import current_omni_platform
-from vllm_omni.worker.gpu_memory_utils import (
-    get_process_gpu_memory,
-    is_process_scoped_memory_available,
-)
 
 logger = init_logger(__name__)
 
 
 class OmniGPUWorkerBase(GPUWorker):
-    """Base GPU worker for vLLM-Omni with process-scoped memory accounting.
+    """Base GPU worker for vLLM-Omni.
 
-    This class overrides determine_available_memory() to use per-process GPU
-    memory tracking via pynvml, allowing multiple stages to initialize
-    concurrently on the same GPU without memory accounting interference.
+    This class overrides determine_available_memory() to compute the KV cache
+    budget from device-level profiling data. Concurrent same-device stage
+    initialization is coordinated externally by the parallel-stage-init
+    admission control + SH/EX device locks (see
+    ``VllmOmniOrchestratorConfig.parallel_stage_init``), which keep each
+    measurement quiescent instead of relying on per-process NVML accounting.
 
     It also replaces vLLM's TorchProfilerWrapper with OmniTorchProfilerWrapper
     for custom trace naming, background gzip, and trace path collection.
@@ -91,21 +89,27 @@ class OmniGPUWorkerBase(GPUWorker):
 
     @torch.inference_mode()
     def determine_available_memory(self) -> int:
-        """Process-scoped GPU memory profiling for concurrent stage initialization.
+        """Device-level GPU memory profiling for the KV cache budget.
 
         Algorithm:
             1. requested_memory = total_gpu_memory * gpu_memory_utilization
                (computed in init_device from cache_config)
 
-            2. process_memory = memory used by THIS process only (via pynvml)
-               - Uses nvmlDeviceGetComputeRunningProcesses to get per-PID memory
-               - Supports CUDA_VISIBLE_DEVICES with indices, UUIDs, or MIG IDs
+            2. profiled_usage = weights + peak_activation + non_torch_increase
+               (measured by ``memory_profiling`` around ``profile_run()``;
+               ``non_torch_increase`` is device-level, so it reflects whatever
+               else is resident on the GPU at profiling time)
 
-            3. available_kv_cache = requested_memory - process_memory
+            3. available_kv_cache = requested_memory - profiled_usage
 
-        Fallback:
-            If NVML is unavailable, falls back to profiling data:
-            available = requested - (weights + activations + non_torch)
+        Note:
+            Process-scoped NVML estimation was removed in favour of the
+            device-level profiling fallback plus the parallel-stage-init
+            admission + SH/EX device locks (see
+            ``VllmOmniOrchestratorConfig.parallel_stage_init``): the EX lock
+            makes each measurement quiescent, so the device-level number is the
+            correct, conservative budget. The NVML helpers in
+            ``gpu_memory_utils`` are retained for diffusion memory reporting.
         """
         if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
             self.model_runner.profile_run()
@@ -127,49 +131,25 @@ class OmniGPUWorkerBase(GPUWorker):
         # true here). Mirror upstream so the omni override keeps it populated.
         self.total_consumed = profile_result.total_consumed
 
-        process_memory = (
-            get_process_gpu_memory(self.local_rank)
-            if is_process_scoped_memory_available() and detect_pid_host()
-            else None
+        profiled_usage = (
+            int(self.model_runner.model_memory_usage)
+            + profile_result.torch_peak_increase
+            + profile_result.non_torch_increase
+        )
+        self.available_kv_cache_memory_bytes = max(0, self.requested_memory - profiled_usage)
+        logger.debug(
+            "Profiling KV budget (PID %d, GPU %d): requested=%s, profiled=%s, available=%s",
+            os.getpid(),
+            self.local_rank,
+            format_gib(self.requested_memory),
+            format_gib(profiled_usage),
+            format_gib(self.available_kv_cache_memory_bytes),
+        )
+        logger.info_once(
+            "Available KV cache memory: %s GiB (device-level profiling)",
+            format_gib(self.available_kv_cache_memory_bytes),
+            scope="local",
         )
-
-        if process_memory is not None:
-            # NVML available: use per-process memory
-            self.available_kv_cache_memory_bytes = max(0, self.requested_memory - process_memory)
-            logger.debug(
-                "Process-scoped memory (PID %d, GPU %d): requested=%s, used=%s, available=%s",
-                os.getpid(),
-                self.local_rank,
-                format_gib(self.requested_memory),
-                format_gib(process_memory),
-                format_gib(self.available_kv_cache_memory_bytes),
-            )
-            logger.info_once(
-                "Available KV cache memory: %s GiB (process-scoped)",
-                format_gib(self.available_kv_cache_memory_bytes),
-                scope="local",
-            )
-        else:
-            # NVML unavailable: use profiling data as conservative fallback
-            profiled_usage = (
-                int(self.model_runner.model_memory_usage)
-                + profile_result.torch_peak_increase
-                + profile_result.non_torch_increase
-            )
-            self.available_kv_cache_memory_bytes = max(0, self.requested_memory - profiled_usage)
-            logger.debug(
-                "Profiling fallback (PID %d, GPU %d): requested=%s, profiled=%s, available=%s",
-                os.getpid(),
-                self.local_rank,
-                format_gib(self.requested_memory),
-                format_gib(profiled_usage),
-                format_gib(self.available_kv_cache_memory_bytes),
-            )
-            logger.info_once(
-                "Available KV cache memory: %s GiB (profiling fallback)",
-                format_gib(self.available_kv_cache_memory_bytes),
-                scope="local",
-            )
 
         return int(self.available_kv_cache_memory_bytes)