Integrate portable edge graph backends with Omni runtime

Extract versioned stage contracts, model shared CPU/iGPU/NPU memory constraints, and add explicit graph routes, reservations, output acknowledgements, and worker retirement. Fix native Windows shared-buffer lifetime and host paths. Document qualified hardware paths and remaining rollout gates.

.gitignore

@@ -295,3 +295,6 @@ benchmarks/edge_harness/results/**/*.npz
 benchmarks/edge_harness/results/**/*.onnx
 benchmarks/edge_harness/results/**/*.npy
 benchmarks/edge_harness/results/**/*.wav
+
+# Portable wire fixtures are source inputs, not generated measurements.
+!packages/omni-stage-contracts/conformance/*.json

benchmarks/tts/stream_latency_bench.py

@@ -408,4 +408,9 @@ async def main() -> int:
 
 
 if __name__ == "__main__":
+    # Install the Windows selector policy before creating the caller's loop.
+    # Importing AsyncOmni inside main() is too late for that existing loop.
+    from vllm_omni.windows.aio import install_selector_policy
+
+    install_selector_policy()
     raise SystemExit(asyncio.run(main()))

docs/design/omni_edge_backends.md

@@ -0,0 +1,235 @@
+# Local edge backends in Omni
+
+The first backend is `external.graph.v1`: one complete, stateless graph operation
+per request, with a fixed input bucket and copied host tensors. It uses the normal
+`Omni` / `AsyncOmni` APIs, `PipelineConfig`, `StageRuntime`, `StagePool` and
+orchestrator. Native vLLM stages retain their model execution, KV cache, sampling
+and batching. No accelerator is mandatory.
+
+## Boundaries
+
+| Owner | Responsibility |
+|---|---|
+| `omni_stage_contracts` | Version/features, buffer descriptors, request/event identities, opaque state, physical device descriptors, artifact manifests. Root import uses only the standard library. |
+| `vllm_omni.host` | Explicit interpreter/OS routes, process bootstrap, process identity and bounded tree retirement. |
+| `vllm_omni.engine.backends` | Adapt backend operations to `StageClientBase`; observe placement; fence epochs and worker generations. |
+| `StageRuntime` | Resolve all replicas, reserve the whole plan before loading, create clients, unwind failures. |
+| `ResourceLedger` | Atomically charge every affected memory constraint; keep uncertain allocations quarantined. |
+| `StagePool` / orchestrator | Existing stage dispatch, handoff, polling, metrics, cancellation and request cleanup. |
+| Model adapter | Tokenization, media preparation, stage transforms, sampling semantics, artifact export and numerical/task validation. |
+
+Native Windows device-init locks use the host temporary directory, independent
+of a local or UNC checkout. This coordinates same-user native processes. It does
+not establish a cross-user or Windows/WSL global lock authority; cross-domain
+workers in one graph plan are coordinated by that controller's reservations.
+
+Native Windows SHM transport also uses a host-local directory for locks and
+acknowledgements. A producer retains each named mapping until its single consumer
+copies the payload and acknowledges its generation, or the producer cancels/closes
+it. A length header handles Windows page-rounded mapping sizes. Retained mappings
+are limited to 64 MiB and 256 buffers per producer process; exceeding either limit
+fails the write explicitly. This host transport fix does not enable cross-OS or
+accelerator zero-copy. POSIX keeps its existing unlink-based lifetime.
+
+`edge.local.session` and `edge.local.external.protocol` retain compatibility imports.
+The old direct `ExternalStage` API remains available; new pipelines should use
+the runtime-owned graph backend. New state handles default to non-replayable and
+non-migratable. Legacy M0 defaults are explicitly preserved in `contracts.legacy`.
+
+## Hardware composition
+
+The abstraction supports these single-machine compositions. This table separates
+topology/admission support from successful execution of a particular model; it
+does not promise that every listed accelerator can run every model.
+
+| Composition | Resource accounting | Current execution evidence |
+|---|---|---|
+| CPU only | Host RAM; WSL quota when applicable | ORT CPU graphs through public Omni on Windows and WSL. Native Windows CPU text remains unqualified. |
+| CPU + integrated GPU | Shared host RAM | Radeon 890M DirectML graph through Omni. This extraction's probe is not a complete vision/language pipeline. |
+| CPU + integrated NPU | Shared host RAM | AMD NPU VitisAI graph through Omni; permitted CPU partitions are explicit in the placement report. |
+| CPU + integrated GPU + integrated NPU | Shared RAM plus optional observed package/power/bandwidth relationships | Topology and admission contracts tested; individual laptop routes execute. Simultaneous full-model overlap is not qualified. |
+| CPU + integrated GPU + integrated NPU + discrete GPU | Shared RAM and separate discrete VRAM | Current Ryzen AI 9 HX370 / Radeon 890M / AMD NPU / RTX 5090 Laptop machine. Spark CUDA text and native two-stage TTS run on Windows and WSL; using all accelerators together is not qualified. |
+| CPU + discrete GPU, with integrated devices unused or absent | Host RAM plus discrete VRAM | Native CUDA text/TTS execution uses this subset. NVIDIA is optional in the abstraction. |
+
+Mobile/embedded SoCs fit the CPU + integrated GPU + NPU composition, but their
+local controller and model paths still require device-specific qualification.
+Galaxy S25 AI Hub inference/profiling passed for one Spark attention component;
+this is component evidence, not full local mobile generation or streaming.
+AI Hub is a test facility, never a deployment dependency. Intel, Apple and other
+vendor combinations are not certified by the AMD/NVIDIA/Qualcomm results.
+
+`DeviceDescriptor` separates physical device identity, execution domain and
+integration. A CPU, integrated GPU and NPU can all reference `machine:ram`.
+An additional discrete GPU references its own `machine:vram:<physical-id>`.
+A WSL allocation can consume both `machine:ram` and `wsl:quota` constraints;
+those are two limits on the same allocation, not two amounts of physical RAM.
+Use the same physical GPU ID when Windows and WSL expose the same adapter.
+
+Package IDs, bandwidth groups and power domains are optional observed relationships.
+Missing values mean unknown. Fitting memory does not authorize overlap on a shared
+package: selection and overlap still require model quality and complete workload
+measurements. Synthetic topology tests cover CPU-only and all four integrated /
+discrete combinations; they are not hardware certification.
+
+## Deploy a graph
+
+Register a `PipelineConfig` with `StageExecutionType.GRAPH`, `model_stage="graph"`
+and an appropriate final output type. Set `async_chunk=False`. A graph following
+another stage must have a model-owned `custom_process_input_func` that returns
+`{"tensors": {name: numpy_array}}`. Names, dtypes and shapes must match the
+manifest's placement examples. Graph stages do not tokenize prompts or apply
+model-specific transformations.
+
+Graph-containing v1 pipelines must be a linear chain with exactly one terminal
+output at the last stage. Fan-out, joins and multiple final outputs are refused
+before loading; those require reference-counted consumer leases.
+
+Each stage in a graph-containing plan declares the same capacities and its own
+demands, including native stages:
+
+```yaml
+pipeline: my_registered_pipeline
+async_chunk: false
+stages:
+  - stage_id: 0
+    backend:
+      name: external.graph.v1
+      manifest: /absolute/path/to/bundle/manifest.json
+      route:
+        name: local-cpu
+        interpreter: /absolute/path/to/worker/python
+        worker: /absolute/path/to/worker_ort.py
+        ep: cpu
+        os_domain: posix
+      max_io_bytes: 8388608
+      min_fraction_on_target: 1.0
+      allowed_providers: [CPUExecutionProvider]
+    resource_budget:
+      capacities: {machine:ram: 2147483648, wsl:quota: 1073741824}
+      demands: {machine:ram: 536870912, wsl:quota: 536870912}
+```
+
+Quote pool keys or use block mappings if your YAML producer treats colon-containing
+keys specially. On native Windows use Windows paths and `os_domain: windows`.
+From WSL use controller-visible paths; host services translate worker paths. Named
+legacy routes still honor `VLLM_OMNI_EXTERNAL_PYTHON_<ROUTE>`, but explicit routes
+have no developer-machine defaults.
+
+Capacities are **pre-load controller ceilings**, below observed free host memory /
+guest limits with safety margin. They are not continuously sampled available RAM.
+Demands must include loading peaks, weights, state, activations, workspace,
+transport copies, retained output and headroom. The ledger does not claim to be
+an OS memory limiter or to attribute shared vendor allocations from RSS.
+Driver/cache memory outside the worker remains part of the deployment estimate.
+
+The graph worker environment needs NumPy, ONNX Runtime and ONNX (for checking
+external weight references against the manifest). Vendor providers need their
+own qualified environment. The controller need not install their runtimes.
+
+## Artifact and placement gate
+
+```json
+{
+  "schema_version": 1,
+  "component": "model-owned-component-name",
+  "files": {"graph.onnx": "<sha256>", "weights.bin": "<sha256>", "inputs.npz": "<sha256>"},
+  "metadata": {
+    "graph_file": "graph.onnx", "example_inputs_file": "inputs.npz",
+    "checkpoint_revision": "<revision>", "precision": "<precision>",
+    "layout": "<layout>", "exporter": "<version>",
+    "calibration": "<record or not applicable>", "validation": "<record>"
+  }
+}
+```
+
+All payload paths are relative to the manifest. Hashes are checked before startup;
+ONNX external weight references must also belong to that verified set. Keep bundles
+immutable during execution. Checkpoint, exporter, precision and validation metadata
+are model-owned: supplying a manifest does not establish model quality.
+
+Warmup records actual provider node assignments. Unverified placement, absent
+providers and undeclared CPU partitions refuse startup. The default target fraction
+is 1.0. A deliberately partitioned NPU graph can lower it and explicitly allow CPU,
+but reports always retain the real split. `execution_plan` records the manifest
+identity, worker generation, route, reservations and placement. It labels this B
+(backend execution), not complete-model P.
+
+## Flow, cancellation and ownership
+
+* Graph pipelines admit one request before the engine ingress queue. Extra work
+  receives `ResourceUnavailable`; await consumption or cancellation and retry.
+  The current API does not queue a batch of graph requests.
+* Each graph client has one in-flight call and one retained result. Delivery does
+  not free the slot. Public synchronous/asynchronous generators acknowledge when
+  iteration advances or closes. Direct engine consumers call
+  `OmniRequestOutput.release_stage_buffers()` in `finally`.
+* After acknowledgement, retained result arrays belong to the application. Keeping
+  arbitrarily many consumed results is application memory, outside engine admission.
+* Intermediate graph handoffs copy their next-stage inputs before acknowledgement.
+  Buffer descriptors carry owner, generation, exact byte size, dtype and shape.
+  Framing rejects unknown required features and malformed lengths before payload allocation.
+* Cancellation fences the request epoch immediately. An active non-preemptible graph
+  call retires its worker. It does not replay work or restart transparently; recreate
+  the pipeline to use that backend again. Completed calls can cancel without killing
+  an otherwise reusable worker.
+* Windows workers announce PID plus process creation time. Cross-OS termination
+  verifies that identity and retires the native process tree before stopping the WSL
+  launcher. An absent/unverified tree keeps its reservation quarantined. Python
+  call references stay charged until the call unwinds.
+
+Native APIs and their admission logic remain compatible. This change adds a shared
+explicit reservation transaction for graph-containing plans; it does not yet replace
+every native memory observation/ledger or qualify mixed-model fan-out and streaming.
+
+## Packaging, validation and scope
+
+Build `packages/omni-stage-contracts` with `uv build`. Its wheel imports without
+torch/vLLM; the optional `wire` extra installs NumPy. The main Omni wheel bundles
+the same source. Python/native implementers can use `conformance/v1.json` from
+the source distribution. A native mobile implementation has not yet run these fixtures.
+
+`examples/edge/run_graph_stage.py` drives the public API, writes a deployment profile,
+records the resolved plan and raw request samples, and saves outputs for numerical
+comparison. It supports explicit CPU, DirectML and VitisAI worker routes.
+
+The implementation has tests for real runtime startup, two-stage handoff on both
+orchestrator modes, public APIs, ingress/consumer backpressure, malformed framing,
+placement refusal, shared-pool overcommit, stale acknowledgements, cancellation,
+worker failure and artifact membership. Native Spark acceptance and small actual
+Windows/WSL graph probes are recorded separately in the workspace experiment report.
+
+Validation snapshot (2026-09-22, base `00510ce38bbb96ef2e4ce462618d06ab142fae7c`):
+
+* WSL focused regression: 415 passed, 1 skipped, 20 deselected. Windows focused
+  regression: 78 passed, 2 skipped; the controller lacks ONNX for those two tests,
+  while actual graph probes use their separate worker environments.
+* Spark 4B BF16 CUDA M0: 11 tests passed on each host. These runs check the
+  specified text/state/cancellation workload, not every model or device.
+* CPU and DirectML Add graph outputs are exact. AMD NPU versus the same A16W8
+  CPU graph has maximum absolute error `4.458427429199219e-05`, normalized
+  maximum `3.111236521857185e-05` against a `0.002` bound. This is backend
+  validation, not original-checkpoint model-quality validation.
+* Real Qwen3-TTS 0.6B two-stage streaming completes on both hosts with checkpoint
+  revision `85e237c12c027371202489a0ec509ded67b5e4b5`. One measured utterance per
+  host still has playback stalls; these diagnostic samples do not satisfy the
+  full streaming, interruption, quality or sustained-thermal acceptance gates.
+* The built Omni wheel runs 20 public-API CPU graph requests outside the checkout
+  with exact outputs and zero retained reservation after shutdown. The standalone
+  contracts source archive imports in a dependency-free environment.
+* Fresh Galaxy S25 [inference](https://workbench.aihub.qualcomm.com/jobs/jp0mvkj6g/)
+  and [NPU profiling](https://workbench.aihub.qualcomm.com/jobs/jpyo7nx75/) passed
+  for the existing W8A16 attention component. Three outputs match the same
+  artifact/input reference exactly; the `full1024` artifact label denotes one
+  attention component's cache bucket, not the complete Spark model.
+
+Raw records, machine/package inventories, failed attempts and reproduction commands
+are preserved in the surrounding edge-infer workspace under
+`analysis/experiments/omni_abstraction_20260922/`. They are not bundled as runtime
+dependencies. Tests of reduced-device configurations are not measurements on a
+different physical hardware SKU.
+
+Still outside the completed v1 scope: arbitrary device/model certification, mobile
+C++ control integration, persistent external AR sessions, shared-device zero-copy,
+state migration/replay, mixed-model streaming/fan-out qualification, joint thermal /
+bandwidth policy, native Windows CPU text qualification, and the full M1–M4 device
+matrix. These remain explicit rollout gates, not implied supported features.

examples/edge/README.md

@@ -0,0 +1,18 @@
+# Local graph example
+
+See [the backend contract and deployment guide](../../docs/design/omni_edge_backends.md).
+
+From a configured Omni environment:
+
+```bash
+python examples/edge/run_graph_stage.py /path/to/manifest.json \
+  --python /path/to/ort-venv/bin/python --report /path/to/run.json
+```
+
+The worker environment needs `numpy`, `onnx` and `onnxruntime`. For Windows,
+use `--windows-worker` and the controller-visible path to `python.exe`.
+`--ep vitisai` or `--ep dml` selects a vendor provider; declare every permitted
+provider using `--allowed-provider`. No provider substitution is hidden.
+
+The example reads fixed-bucket inputs from the bundle. It reports backend
+execution only; use your own model adapter and quality fixtures for a model claim.

examples/edge/run_graph_stage.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run a verified graph bundle through the public Omni API on one machine.
+
+The bundle manifest identifies graph_file and example_inputs_file. This example
+uses those fixed-bucket examples as inputs; real pipelines supply model-owned
+input/handoff functions. All memory arguments are controller ceilings/estimates.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+from omni_stage_contracts import ArtifactManifest
+from vllm_omni.config.pipeline_registry import register_pipeline
+from vllm_omni.config.stage_config import PipelineConfig, StageExecutionType, StagePipelineConfig
+from vllm_omni.entrypoints.omni import Omni
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("manifest", type=Path)
+    parser.add_argument("--route", default="ort-cpu")
+    parser.add_argument("--python", dest="interpreter")
+    parser.add_argument("--windows-worker", action="store_true")
+    parser.add_argument("--ep", default="cpu", choices=["cpu", "dml", "vitisai"])
+    parser.add_argument("--allowed-provider", action="append", default=[])
+    parser.add_argument("--min-fraction", type=float, default=1.0)
+    parser.add_argument("--capacity-mib", type=int, default=1024)
+    parser.add_argument("--stage-mib", type=int, default=512)
+    parser.add_argument("--max-io-mib", type=int, default=8)
+    parser.add_argument("--repeats", type=int, default=20)
+    parser.add_argument("--report", type=Path, required=True)
+    args = parser.parse_args()
+    if args.repeats < 1:
+        parser.error("--repeats must be positive")
+    manifest_path = args.manifest.resolve()
+    manifest = ArtifactManifest.read(manifest_path)
+    with np.load(manifest_path.parent / manifest.metadata["example_inputs_file"], allow_pickle=False) as samples:
+        inputs = {name: samples[name] for name in samples.files}
+    register_pipeline(
+        PipelineConfig(
+            model_type="omni_external_graph",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="graph",
+                    execution_type=StageExecutionType.GRAPH,
+                    final_output=True,
+                    final_output_type="latent",
+                ),
+            ),
+        )
+    )
+    route = args.route
+    if args.interpreter:
+        import vllm_omni.edge.local.external.worker_ort as worker_module
+
+        route = {
+            "name": args.route,
+            "interpreter": args.interpreter,
+            "worker": worker_module.__file__,
+            "os_domain": "windows" if args.windows_worker else "posix",
+            "ep": args.ep,
+        }
+    backend = {
+        "name": "external.graph.v1",
+        "route": route,
+        "manifest": str(manifest_path),
+        "min_fraction_on_target": args.min_fraction,
+        "max_io_bytes": args.max_io_mib << 20,
+        "profile_prefix": str(args.report.resolve().with_suffix(".placement")),
+    }
+    if args.allowed_provider:
+        backend["allowed_providers"] = args.allowed_provider
+    result = {
+        "schema_version": 1,
+        "evidence": "B",
+        "platform": platform.platform(),
+        "python": sys.executable,
+        "workload": {"repeats": args.repeats, "concurrency": 1},
+        "artifact": str(manifest_path),
+        "samples_s": [],
+        "started_unix": time.time(),
+    }
+    engine = None
+    try:
+        with tempfile.TemporaryDirectory(prefix="omni-graph-") as tmp:
+            deploy = Path(tmp) / "deploy.yaml"
+            deploy.write_text(
+                yaml.safe_dump(
+                    {
+                        "pipeline": "omni_external_graph",
+                        "async_chunk": False,
+                        "stages": [
+                            {
+                                "stage_id": 0,
+                                "backend": backend,
+                                "resource_budget": {
+                                    "capacities": {"machine:ram": args.capacity_mib << 20},
+                                    "demands": {"machine:ram": args.stage_mib << 20},
+                                },
+                            }
+                        ],
+                    }
+                )
+            )
+            started = time.perf_counter()
+            engine = Omni(model=str(manifest_path.parent), deploy_config=str(deploy))
+            result["startup_s"] = time.perf_counter() - started
+            client = engine.engine.stage_pools[0].stage_client
+            result["execution_plan"] = client.execution_plan
+            outputs = None
+            for _ in range(args.repeats):
+                started = time.perf_counter()
+                (output,) = engine.generate({"tensors": inputs}, use_tqdm=False)
+                result["samples_s"].append(time.perf_counter() - started)
+                outputs = output.custom_output["tensors"]
+            result["worker_stats"] = client._worker.stats()
+            result["finite"] = all(np.isfinite(v).all().item() for v in outputs.values())
+            args.report.parent.mkdir(parents=True, exist_ok=True)
+            np.savez(args.report.with_suffix(".npz"), **outputs)
+            result["status"] = "passed"
+    except Exception as exc:
+        result.update(status="failed", error=f"{type(exc).__name__}: {exc}")
+        raise
+    finally:
+        if engine is not None:
+            engine.shutdown()
+            result["ledger_after_shutdown"] = engine.engine._runtime.resource_ledger.snapshot()
+        result["ended_unix"] = time.time()
+        args.report.parent.mkdir(parents=True, exist_ok=True)
+        args.report.write_text(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

omni_stage_contracts/__init__.py

@@ -0,0 +1,27 @@
+"""Source-checkout bridge; wheels use the canonical package directly."""
+
+from pathlib import Path
+
+__path__ = [str(Path(__file__).resolve().parents[1] / "packages/omni-stage-contracts/omni_stage_contracts")]
+
+from .types import (
+    PROTOCOL_VERSION,
+    ArtifactManifest,
+    BufferRef,
+    DeviceDescriptor,
+    StageEvent,
+    StageRequest,
+    StateHandle,
+    negotiate,
+)
+
+__all__ = [
+    "PROTOCOL_VERSION",
+    "ArtifactManifest",
+    "BufferRef",
+    "DeviceDescriptor",
+    "StageEvent",
+    "StageRequest",
+    "StateHandle",
+    "negotiate",
+]

packages/omni-stage-contracts/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include conformance *.json

packages/omni-stage-contracts/README.md

@@ -0,0 +1,9 @@
+# Omni stage contracts
+
+Versioned, dependency-free data contracts for local Omni stage adapters.
+The optional `wire` extra provides the NumPy copied-host tensor transport.
+Importing the contracts never imports torch, vLLM or vLLM-Omni.
+
+Build from this directory with `uv build`; both the wheel and source archive
+contain the canonical package. The parent Omni distribution bundles the same
+source for compatibility. Install matching releases when using both distributions.

packages/omni-stage-contracts/conformance/v1.json

@@ -0,0 +1,24 @@
+{
+  "protocol_version": 1,
+  "required_features": ["host-copy", "request-epochs"],
+  "buffer": {
+    "object_id": "output-1", "owner": "stage-0", "generation": "worker-A",
+    "dtype": "float32", "shape": [2, 3], "nbytes": 24,
+    "memory_domain": "host-copy", "layout": "C", "ready": true
+  },
+  "request": {
+    "request_id": "request-1", "stage_id": 0, "epoch": 2,
+    "worker_generation": "worker-A", "operation": "run", "inputs": []
+  },
+  "event": {
+    "request_id": "request-1", "stage_id": 0, "epoch": 2, "seq": 1,
+    "kind": "tensor", "worker_generation": "worker-A", "buffers": [], "terminal": true
+  },
+  "state": {
+    "session_id": "session-1", "backend": "backend-1", "artifact_id": "sha256:fixture",
+    "backend_instance_id": "backend-instance-A", "worker_generation": "worker-A",
+    "layout_version": 1, "epoch": 2, "replayable": false, "migratable": false
+  },
+  "reject_buffer_overrides": [{"nbytes": 23}, {"shape": [-1, 3]}, {"memory_domain": "cuda-ipc"}],
+  "stale_state_overrides": [{"epoch": 1}, {"worker_generation": "worker-B"}, {"backend_instance_id": "backend-instance-B"}]
+}

packages/omni-stage-contracts/omni_stage_contracts/__init__.py

@@ -0,0 +1,23 @@
+"""Portable Omni data contracts. Importing this module needs only the stdlib."""
+
+from .types import (
+    PROTOCOL_VERSION,
+    ArtifactManifest,
+    BufferRef,
+    DeviceDescriptor,
+    StageEvent,
+    StageRequest,
+    StateHandle,
+    negotiate,
+)
+
+__all__ = [
+    "PROTOCOL_VERSION",
+    "ArtifactManifest",
+    "BufferRef",
+    "DeviceDescriptor",
+    "StageEvent",
+    "StageRequest",
+    "StateHandle",
+    "negotiate",
+]

packages/omni-stage-contracts/omni_stage_contracts/legacy.py

@@ -0,0 +1,87 @@
+"""Compatible M0 contracts; new backend state uses types.StateHandle."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+STATE_LAYOUT_VERSION = 1
+
+
+@dataclass(frozen=True)
+class StateHandle:
+    """A reference to session state the backend owns.
+
+    The engine never reads the state. It reads this, to decide whether a
+    request may continue a session or must start a new one.
+    """
+
+    session_id: str
+    backend: str
+    artifact_id: str
+    """Ties the state to the checkpoint that produced it. The same shapes from
+    a different quantization are not the same state."""
+    layout_version: int = STATE_LAYOUT_VERSION
+    epoch: int = 0
+    """Incremented on every cancel. Events from a lower epoch are stale."""
+    replayable: bool = True
+    """A text session can be rebuilt by re-submitting its committed prompt, so
+    recovery after a backend crash is possible at an output boundary."""
+    migratable: bool = False
+    """Never true in M0. Moving KV between backends needs a measured layout
+    conversion for the specific pair; assuming it is the failure the proposal
+    warns about."""
+    created_unix: float = field(default_factory=time.time)
+
+    def next_epoch(self) -> StateHandle:
+        from dataclasses import replace
+
+        return replace(self, epoch=self.epoch + 1)
+
+    backend_instance_id: str = ""
+    worker_generation: str = ""
+
+    def accepts(self, other: StateHandle) -> bool:
+        """Whether ``other`` may continue this session's state."""
+        return (
+            other.backend == self.backend
+            and other.backend_instance_id == self.backend_instance_id
+            and other.worker_generation == self.worker_generation
+            and other.session_id == self.session_id
+            and other.artifact_id == self.artifact_id
+            and other.layout_version == self.layout_version
+            and other.epoch == self.epoch
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass(frozen=True)
+class ChunkEvent:
+    """One unit of output, with everything needed to order and retire it."""
+
+    request_id: str
+    stage_id: int
+    seq: int
+    """Monotonic within (request, epoch). Gaps mean loss; repeats mean a bug."""
+    epoch: int
+    kind: str
+    """``token`` | ``done`` | ``error`` | ``cancelled``."""
+    payload: Any
+    """For ``token``: the incremental text and token ids. Never the whole
+    output so far -- a cumulative payload makes a dropped chunk invisible."""
+    started_unix: float
+    emitted_unix: float
+    final: bool = False
+    error: str | None = None
+    input_watermark: int = 0
+    """How much of the input the producer had consumed when this was emitted.
+    One number in M0 (prompt tokens); it is what a duplex turn will align on."""
+    release_token: str = ""
+    """Returned by the consumer to give credit back. Empty when the event
+    carries no resource the producer is holding on its behalf."""
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)

packages/omni-stage-contracts/omni_stage_contracts/types.py

@@ -0,0 +1,194 @@
+"""Versioned host-copy contracts shared by Python and native stage adapters."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+PROTOCOL_VERSION = 1
+FEATURES = frozenset({"host-copy", "request-epochs"})
+_ITEM_SIZES = {
+    "bool": 1,
+    "uint8": 1,
+    "int8": 1,
+    "int16": 2,
+    "uint16": 2,
+    "int32": 4,
+    "uint32": 4,
+    "int64": 8,
+    "uint64": 8,
+    "float16": 2,
+    "float32": 4,
+    "float64": 8,
+}
+
+
+def negotiate(version: int, required: list[str] | tuple[str, ...] = ()) -> None:
+    if type(version) is not int or version != PROTOCOL_VERSION:
+        raise ValueError(f"unsupported stage protocol version: {version!r}")
+    missing = set(required) - FEATURES
+    if missing:
+        raise ValueError(f"unsupported required stage features: {sorted(missing)}")
+
+
+@dataclass(frozen=True)
+class BufferRef:
+    object_id: str
+    owner: str
+    generation: str
+    dtype: str
+    shape: tuple[int, ...]
+    nbytes: int
+    memory_domain: str = "host-copy"
+    layout: str = "C"
+    ready: bool = True
+
+    def __post_init__(self) -> None:
+        if not self.object_id or not self.owner or not self.generation:
+            raise ValueError("buffer identity/owner/generation are required")
+        if self.memory_domain != "host-copy" or self.layout != "C" or not self.ready:
+            raise ValueError("v1 only supports ready, contiguous copied host buffers")
+        if self.dtype not in _ITEM_SIZES:
+            raise ValueError(f"unsupported dtype: {self.dtype}")
+        if len(self.shape) > 32 or any(type(x) is not int or x < 0 for x in self.shape):
+            raise ValueError("invalid tensor dimensions")
+        expected = math.prod(self.shape) * _ITEM_SIZES[self.dtype]
+        if type(self.nbytes) is not int or self.nbytes != expected or self.nbytes > 2 << 30:
+            raise ValueError("tensor shape/dtype/byte length mismatch or payload exceeds limit")
+
+
+@dataclass(frozen=True)
+class StateHandle:
+    session_id: str
+    backend: str
+    artifact_id: str
+    layout_version: int = 1
+    epoch: int = 0
+    replayable: bool = False
+    migratable: bool = False
+    backend_instance_id: str = ""
+    worker_generation: str = ""
+
+    def accepts(self, other: StateHandle) -> bool:
+        return all(
+            getattr(self, k) == getattr(other, k)
+            for k in (
+                "session_id",
+                "backend",
+                "artifact_id",
+                "layout_version",
+                "epoch",
+                "backend_instance_id",
+                "worker_generation",
+            )
+        )
+
+    def next_epoch(self) -> StateHandle:
+        from dataclasses import replace
+
+        return replace(self, epoch=self.epoch + 1)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass(frozen=True)
+class StageRequest:
+    request_id: str
+    stage_id: int
+    epoch: int = 0
+    worker_generation: str = ""
+    operation: str = "run"
+    inputs: tuple[BufferRef, ...] = ()
+
+    def __post_init__(self):
+        if not self.request_id or not self.worker_generation or not self.operation:
+            raise ValueError("request identity, worker generation and operation are required")
+        if any(type(n) is not int or n < 0 for n in (self.stage_id, self.epoch)):
+            raise ValueError("stage_id and epoch must be nonnegative integers")
+
+
+@dataclass(frozen=True)
+class StageEvent:
+    request_id: str
+    stage_id: int
+    epoch: int
+    seq: int
+    kind: str
+    worker_generation: str
+    buffers: tuple[BufferRef, ...] = ()
+    terminal: bool = False
+    error: str | None = None
+
+    def __post_init__(self):
+        if not self.request_id or not self.worker_generation or not self.kind:
+            raise ValueError("event identity, worker generation and kind are required")
+        if any(type(n) is not int or n < 0 for n in (self.stage_id, self.epoch, self.seq)):
+            raise ValueError("stage_id, epoch and seq must be nonnegative integers")
+
+
+@dataclass(frozen=True)
+class DeviceDescriptor:
+    """Physical identity is independent from per-domain execution routes.
+
+    Unknown package/power/bandwidth relationships are empty, never inferred
+    from an integrated/discrete label. Pool IDs identify backing constraints.
+    """
+
+    device_id: str
+    kind: str
+    domain_id: str
+    memory_pool_ids: tuple[str, ...]
+    integration: str = "unknown"
+    package_id: str | None = None
+    bandwidth_group_ids: tuple[str, ...] = ()
+    power_domain_ids: tuple[str, ...] = ()
+
+    def __post_init__(self) -> None:
+        if self.kind not in {"cpu", "gpu", "npu"}:
+            raise ValueError(f"unknown device kind: {self.kind}")
+        if self.integration not in {"integrated", "discrete", "unknown"}:
+            raise ValueError("unknown integration class")
+        if not self.device_id or not self.domain_id or not self.memory_pool_ids:
+            raise ValueError("device, domain and physical memory constraints are required")
+        if len(set(self.memory_pool_ids)) != len(self.memory_pool_ids):
+            raise ValueError("duplicate physical memory constraint")
+
+
+@dataclass(frozen=True)
+class ArtifactManifest:
+    """A complete artifact payload set; paths are relative to the manifest."""
+
+    schema_version: int
+    component: str
+    files: dict[str, str]
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def read(cls, path: str | Path) -> ArtifactManifest:
+        path = Path(path)
+        manifest = cls(**json.loads(path.read_text(encoding="utf-8")))
+        if manifest.schema_version != 1 or not manifest.files:
+            raise ValueError("unsupported or empty artifact manifest")
+        root = path.parent.resolve()
+        for name, digest in manifest.files.items():
+            if not isinstance(name, str) or "\\" in name or ":" in name:
+                raise ValueError("artifact payloads must use portable relative POSIX paths")
+            target = (root / name).resolve()
+            if not target.is_relative_to(root):
+                raise ValueError(f"artifact path leaves bundle: {name}")
+            if not target.is_file() or file_digest(target) != digest:
+                raise ValueError(f"artifact payload missing or digest mismatch: {name}")
+        return manifest
+
+
+def file_digest(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as stream:
+        for block in iter(lambda: stream.read(1 << 20), b""):
+            h.update(block)
+    return h.hexdigest()

packages/omni-stage-contracts/omni_stage_contracts/wire.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The wire between the engine and a stage that runs in another process.
+
+Deliberately tiny, and deliberately **stdlib + numpy only**. The worker on the
+other end of this socket may be a native-Windows interpreter that has no torch,
+no vllm and no vllm_omni on its path -- ``C:\\Users\\zhout\\npu-ep`` is exactly
+that: Python 3.12, onnxruntime, onnx, numpy, nothing else. Anything imported
+here has to be installable there, so this module imports neither.
+
+**Why a socket at all.** Two devices on this laptop cannot be reached from the
+engine's own interpreter, and for different reasons:
+
+* the **XDNA2 NPU** is an MCDM device. WSL2's GPU-PV forwards WDDM display
+  adapters only, so there is no ``/dev/accel`` inside WSL and no amount of
+  installing changes that. VitisAI is a Windows DLL.
+* the **Radeon 890M** is reachable from WSL through DirectML, but
+  ``torch-directml`` pins torch 2.4.1 and Omni runs 2.13.
+
+So out-of-process is forced, not preferred. One protocol covers both, plus the
+native-Windows DirectML route, which needs its own venv anyway because
+``onnxruntime-directml`` is pinned at 1.24.4 and cannot load the VitisAI EP
+(that needs ORT >= 1.25) -- the two AMD devices cannot share an interpreter.
+
+**Framing.** ``[4-byte big-endian header length][JSON header][raw tensor bytes]``.
+Tensors are described in the header (name, dtype, shape, nbytes) and their
+buffers are concatenated after it in the order listed, so a receive is two
+reads and one ``frombuffer`` per tensor -- no per-element work and no pickle.
+Measured on this machine, WSL <-> native Windows: 11 us half round trip,
+330 MB/s. That is ample for a stage boundary (a 2 s Code2Wav chunk is ~192 KB,
+about 0.6 ms) and nowhere near enough for per-layer traffic inside a decode
+loop, which is the arithmetic behind "coarse stages only".
+
+``pickle`` is not used anywhere in here on purpose: the header is JSON, the
+payload is raw buffers, so a hostile or merely broken peer cannot execute code
+by replying.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import socket
+import struct
+from typing import Any
+
+import numpy as np
+
+from .types import negotiate
+
+# Ops the client sends.
+OP_LOAD = "load"
+"""Open a session on a graph. Replies with the placement report."""
+OP_RUN = "run"
+"""One forward. Inputs in, outputs out."""
+OP_STATS = "stats"
+"""Counters and RSS. The worker's memory is invisible to the engine's sampler
+when the worker is a Windows process, so it has to report its own."""
+OP_CLOSE = "close"
+
+# Ops the worker sends.
+OP_HELLO = "hello"
+"""First message after connect-back: interpreter, versions, available EPs."""
+OP_OK = "ok"
+OP_ERR = "err"
+
+MAX_HEADER_BYTES = 1 << 20
+"""A header is metadata; a megabyte of it is a bug or an attack, not a graph."""
+
+MAX_PAYLOAD_BYTES = 2 << 30
+"""Refuse rather than allocate on a corrupt length. Two gibibytes is far above
+any legitimate stage payload here -- the largest is a vision tower's patch
+tensor, single-digit MB."""
+
+
+class ProtocolError(RuntimeError):
+    """The peer sent something this protocol cannot represent."""
+
+
+class WorkerError(RuntimeError):
+    """The peer ran and failed. Carries the remote traceback when there is one."""
+
+    def __init__(self, message: str, *, remote_traceback: str = "", code: str = "") -> None:
+        # The remote traceback goes in the message, not only in an attribute:
+        # a bare "AssertionError:" surfacing here with the actual stack one
+        # attribute away is a debugging loop nobody should have to discover.
+        detail = f"{message}\n--- worker traceback ---\n{remote_traceback}" if remote_traceback else message
+        super().__init__(detail)
+        self.message = message
+        self.remote_traceback = remote_traceback
+        self.code = code
+
+
+def _recv_exactly(sock: socket.socket, n: int) -> bytes:
+    """Read exactly ``n`` bytes or raise.
+
+    ``recv`` is allowed to return short reads on a stream socket and does, once
+    payloads get past a page or two; a partial read here would be silently
+    misparsed as the next tensor.
+    """
+    if n == 0:
+        return b""
+    chunks: list[bytes] = []
+    remaining = n
+    while remaining:
+        chunk = sock.recv(min(remaining, 1 << 20))
+        if not chunk:
+            raise ProtocolError(f"peer closed with {remaining} of {n} bytes outstanding")
+        chunks.append(chunk)
+        remaining -= len(chunk)
+    return b"".join(chunks) if len(chunks) > 1 else chunks[0]
+
+
+def send_message(
+    sock: socket.socket,
+    op: str,
+    body: dict[str, Any] | None = None,
+    tensors: dict[str, np.ndarray] | None = None,
+) -> None:
+    """Send one message. ``tensors`` are sent as raw buffers after the header."""
+    tensors = tensors or {}
+    described: list[dict[str, Any]] = []
+    buffers: list[memoryview] = []
+    for name, array in tensors.items():
+        array = np.asarray(array)
+        if (
+            not isinstance(name, str)
+            or not name
+            or array.dtype.kind not in "biuf"
+            or array.dtype.itemsize not in (1, 2, 4, 8)
+        ):
+            raise ProtocolError("unsupported tensor name or dtype")
+        if sum(b.nbytes for b in buffers) + array.nbytes > MAX_PAYLOAD_BYTES:
+            raise ProtocolError("outgoing payload exceeds limit")
+        # Record the shape *before* making it contiguous: ``ascontiguousarray``
+        # promotes a 0-d array to shape (1,), which would silently turn a
+        # scalar graph input into a length-1 vector on the far side.
+        shape = list(array.shape)
+        contiguous = np.ascontiguousarray(array)
+        described.append(
+            {
+                "name": name,
+                # ``str(dtype)`` loses endianness for multibyte types; ``.str``
+                # keeps the '<'/'>' so a big-endian peer would still decode.
+                "dtype": contiguous.dtype.str,
+                "shape": shape,
+                "nbytes": int(contiguous.nbytes),
+            }
+        )
+        # ``reshape(-1)`` rather than a direct cast: ``memoryview.cast`` refuses
+        # a 0-dimensional source.
+        buffers.append(memoryview(contiguous.reshape(-1)).cast("B"))
+
+    if op == OP_HELLO and os.environ.get("VLLM_OMNI_WORKER_IDENTITY"):
+        body = {**(body or {}), "process_identity": json.loads(os.environ["VLLM_OMNI_WORKER_IDENTITY"])}
+    header = json.dumps(
+        {"version": 1, "required_features": ["host-copy"], "op": op, "body": body or {}, "tensors": described}
+    ).encode("utf-8")
+    if len(header) > MAX_HEADER_BYTES:
+        raise ProtocolError(f"header is {len(header)} bytes, over the {MAX_HEADER_BYTES} limit")
+    sock.sendall(struct.pack(">I", len(header)))
+    sock.sendall(header)
+    for buffer in buffers:
+        sock.sendall(buffer)
+
+
+def recv_message(
+    sock: socket.socket, *, max_payload_bytes: int = MAX_PAYLOAD_BYTES
+) -> tuple[str, dict[str, Any], dict[str, np.ndarray]]:
+    """Receive one message as ``(op, body, tensors)``.
+
+    Raises :class:`WorkerError` on an ``err`` reply so callers do not each have
+    to remember to check; every other op is returned as-is.
+    """
+    (header_len,) = struct.unpack(">I", _recv_exactly(sock, 4))
+    if header_len > MAX_HEADER_BYTES:
+        raise ProtocolError(f"peer announced a {header_len}-byte header, over the limit")
+    try:
+        header = json.loads(_recv_exactly(sock, header_len).decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise ProtocolError(f"header is not JSON: {exc}") from exc
+
+    if not isinstance(header, dict):
+        raise ProtocolError("invalid message header")
+    try:
+        required = header.get("required_features", [])
+        if not isinstance(required, list) or any(not isinstance(f, str) for f in required):
+            raise ValueError("required_features must be a list of names")
+        # Missing version/features is the explicit legacy-v1 conversion.
+        negotiate(header.get("version", 1), required)
+    except ValueError as exc:
+        raise ProtocolError(str(exc)) from exc
+    if not isinstance(header.get("body", {}), dict) or not isinstance(header.get("tensors", []), list):
+        raise ProtocolError("invalid message schema")
+    total = 0
+    names = set()
+    try:
+        for desc in header.get("tensors", ()):
+            dtype = np.dtype(desc["dtype"])
+            shape, size, name = desc["shape"], desc["nbytes"], desc["name"]
+            if dtype.hasobject or dtype.kind not in "biuf" or dtype.itemsize not in (1, 2, 4, 8):
+                raise ValueError("unsupported tensor dtype")
+            if not isinstance(shape, list) or len(shape) > 32 or any(type(x) is not int or x < 0 for x in shape):
+                raise ValueError("invalid tensor dimensions")
+            if type(size) is not int or size != math.prod(shape) * dtype.itemsize:
+                raise ValueError("tensor byte length disagrees with dtype/shape")
+            if not isinstance(name, str) or not name or name in names:
+                raise ValueError("invalid or duplicate tensor name")
+            names.add(name)
+            total += size
+    except (KeyError, TypeError, ValueError) as exc:
+        raise ProtocolError(f"invalid tensor descriptor: {exc}") from exc
+    if total > min(max_payload_bytes, MAX_PAYLOAD_BYTES):
+        raise ProtocolError(f"peer announced a {total}-byte payload, over the limit")
+
+    tensors: dict[str, np.ndarray] = {}
+    for spec in header.get("tensors", ()):
+        raw = _recv_exactly(sock, int(spec["nbytes"]))
+        array = np.frombuffer(raw, dtype=np.dtype(spec["dtype"]))
+        # ``frombuffer`` is read-only and aliases ``raw``; a consumer that
+        # writes in place (ORT's IOBinding does) needs its own buffer.
+        tensors[spec["name"]] = array.reshape(spec["shape"]).copy()
+
+    op = str(header.get("op", ""))
+    body = header.get("body") or {}
+    if op == OP_ERR:
+        raise WorkerError(
+            str(body.get("message", "worker failed with no message")),
+            remote_traceback=str(body.get("traceback", "")),
+            code=str(body.get("code", "")),
+        )
+    return op, body, tensors

packages/omni-stage-contracts/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "omni-stage-contracts"
+version = "0.1.0"
+description = "Dependency-free stage contracts from the vLLM-Omni repository"
+requires-python = ">=3.10"
+readme = "README.md"
+
+[project.optional-dependencies]
+wire = ["numpy>=1.24"]
+
+[tool.setuptools]
+packages = ["omni_stage_contracts"]

pyproject.toml

@@ -168,9 +168,12 @@ vllm_omni_windows = "vllm_omni.windows:activate"
 vllm_omni_windows = "vllm_omni.windows:activate"
 
 
+[tool.setuptools.package-dir]
+omni_stage_contracts = "packages/omni-stage-contracts/omni_stage_contracts"
+
 [tool.setuptools.packages.find]
 where = ["."]
-include = ["vllm_omni*"]
+include = ["vllm_omni*", "omni_stage_contracts*"]
 
 [tool.setuptools.package-data]
 "vllm_omni" = ["_version.py", "py.typed", "transformers_utils/chat_templates/*.jinja"]

tests/edge/local/external/fake_worker.py

@@ -25,9 +25,7 @@ from pathlib import Path
 import numpy as np
 
 _HERE = Path(__file__).resolve()
-_PROTOCOL = (
-    _HERE.parents[4] / "vllm_omni" / "edge" / "local" / "external" / "protocol.py"
-)
+_PROTOCOL = _HERE.parents[4] / "vllm_omni" / "edge" / "local" / "external" / "protocol.py"
 
 spec = importlib.util.spec_from_file_location("_fake_protocol", _PROTOCOL)
 proto = importlib.util.module_from_spec(spec)
@@ -78,13 +76,15 @@ def main() -> int:
                 "session_create_s": 0.01,
                 "warmup_s": 0.001,
                 "rss_bytes": 111 * 2**20,
-                "inputs": [{"name": n, "type": "tensor(float)", "shape": list(t.shape)}
-                           for n, t in tensors.items()],
+                "inputs": [{"name": n, "type": "tensor(float)", "shape": list(t.shape)} for n, t in tensors.items()],
                 "outputs": [{"name": "y", "type": "tensor(float)", "shape": []}],
             }
             reply.update(config.get("load", {}))
             proto.send_message(sock, proto.OP_OK, reply)
         elif op == proto.OP_RUN:
+            import time
+
+            time.sleep(float(config.get("run", {}).get("sleep_s", 0)))
             scale = float(config.get("run", {}).get("scale", 2.0))
             proto.send_message(
                 sock,
@@ -94,16 +94,21 @@ def main() -> int:
             )
         elif op == proto.OP_STATS:
             proto.send_message(
-                sock, proto.OP_OK,
-                {"runs": 1, "total_run_s": 0.001, "rss_bytes": 111 * 2**20,
-                 "peak_rss_bytes": 122 * 2**20, "report": {}},
+                sock,
+                proto.OP_OK,
+                {
+                    "runs": 1,
+                    "total_run_s": 0.001,
+                    "rss_bytes": 111 * 2**20,
+                    "peak_rss_bytes": 122 * 2**20,
+                    "report": {},
+                },
             )
         elif op == proto.OP_CLOSE:
             proto.send_message(sock, proto.OP_OK, {})
             return 0
         else:
-            proto.send_message(sock, proto.OP_ERR,
-                               {"message": f"unknown op {op}", "code": "ValueError"})
+            proto.send_message(sock, proto.OP_ERR, {"message": f"unknown op {op}", "code": "ValueError"})
 
 
 if __name__ == "__main__":

tests/edge/local/external/test_protocol.py

@@ -16,9 +16,7 @@ from vllm_omni.edge.local.external import protocol
 def _round_trip(op, body=None, tensors=None):
     near, far = socket.socketpair()
     try:
-        threading.Thread(
-            target=lambda: protocol.send_message(near, op, body, tensors), daemon=True
-        ).start()
+        threading.Thread(target=lambda: protocol.send_message(near, op, body, tensors), daemon=True).start()
         return protocol.recv_message(far)
     finally:
         near.close()
@@ -66,8 +64,7 @@ def test_empty_message_and_no_tensors():
 
 def test_error_reply_raises_with_remote_detail():
     with pytest.raises(protocol.WorkerError) as excinfo:
-        _round_trip(protocol.OP_ERR, {"message": "boom", "code": "ValueError",
-                                      "traceback": "far side traceback"})
+        _round_trip(protocol.OP_ERR, {"message": "boom", "code": "ValueError", "traceback": "far side traceback"})
     assert excinfo.value.code == "ValueError"
     assert "far side traceback" in excinfo.value.remote_traceback
 
@@ -92,3 +89,30 @@ def test_oversized_header_is_refused_before_allocation():
     finally:
         near.close()
         far.close()
+
+
+@pytest.mark.parametrize(
+    "override",
+    [
+        {"version": 2},
+        {"version": True},
+        {"required_features": ["unknown-required-feature"]},
+        {"tensors": [{"name": "x", "dtype": "<f4", "shape": [2], "nbytes": 1}]},
+        {"tensors": [{"name": "x", "dtype": "|O", "shape": [1], "nbytes": 8}]},
+        {"tensors": [{"name": "x", "dtype": "<f4", "shape": [-1], "nbytes": -4}]},
+    ],
+)
+def test_invalid_frame_refused_without_reading_payload(override):
+    import json
+    import struct
+
+    near, far = socket.socketpair()
+    far.settimeout(0.2)
+    try:
+        header = json.dumps({"version": 1, "op": "run", "body": {}, "tensors": [], **override}).encode()
+        near.sendall(struct.pack(">I", len(header)) + header)
+        with pytest.raises(protocol.ProtocolError):
+            protocol.recv_message(far)
+    finally:
+        near.close()
+        far.close()

tests/engine/test_graph_backend.py

@@ -0,0 +1,526 @@
+"""The real StageRuntime/StagePool with a subprocess graph worker."""
+
+import asyncio
+import hashlib
+import json
+import subprocess
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from omni_stage_contracts import BufferRef, StateHandle, negotiate
+from vllm_omni.config.stage_config import (
+    DeployConfig,
+    PipelineConfig,
+    StageDeployConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+    merge_pipeline_deploy,
+)
+from vllm_omni.engine.resource_ledger import ResourceLedger, ResourceUnavailable
+from vllm_omni.engine.stage_runtime import StageRuntime
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def artifact(tmp_path):
+    (tmp_path / "graph.onnx").write_bytes(b"synthetic graph, not model evidence")
+    np.savez(tmp_path / "inputs.npz", x=np.ones((2, 3), dtype=np.float32))
+    files = {n: hashlib.sha256((tmp_path / n).read_bytes()).hexdigest() for n in ("graph.onnx", "inputs.npz")}
+    path = tmp_path / "manifest.json"
+    path.write_text(
+        json.dumps(
+            {
+                "schema_version": 1,
+                "component": "test",
+                "files": files,
+                "metadata": {"graph_file": "graph.onnx", "example_inputs_file": "inputs.npz"},
+            }
+        )
+    )
+    return path
+
+
+@pytest.fixture
+def graph_runtime(tmp_path, monkeypatch):
+    config_path = tmp_path / "worker.json"
+    config_path.write_text("{}")
+    monkeypatch.setenv("VLLM_OMNI_FAKE_WORKER_JSON", str(config_path))
+    monkeypatch.setenv("VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU", str(ROOT / "tests/edge/local/external/fake_worker.py"))
+    monkeypatch.setenv("VLLM_OMNI_EXTERNAL_PYTHON_ORT_CPU", sys.executable)
+    runtimes = []
+
+    def create(
+        *, stages=1, capacity=1024 << 20, worker_config=None, start=True, chain=False, real_ort=False, entrypoint=None
+    ):
+        config_path.write_text(json.dumps(worker_config or {}))
+        manifest = artifact(tmp_path)
+        if real_ort:
+            onnx = pytest.importorskip("onnx")
+            pytest.importorskip("onnxruntime")
+            from onnx import TensorProto, helper
+
+            graph = helper.make_graph(
+                [helper.make_node("Add", ["x", "x"], ["x_out"])],
+                "double",
+                [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])],
+                [helper.make_tensor_value_info("x_out", TensorProto.FLOAT, [2, 3])],
+            )
+            model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)], ir_version=9)
+            onnx.save(model, tmp_path / "graph.onnx")
+            data = json.loads(manifest.read_text())
+            data["files"]["graph.onnx"] = hashlib.sha256((tmp_path / "graph.onnx").read_bytes()).hexdigest()
+            manifest.write_text(json.dumps(data))
+            monkeypatch.delenv("VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU", raising=False)
+        pipeline = PipelineConfig(
+            model_type="test_graph",
+            stages=tuple(
+                StagePipelineConfig(
+                    stage_id=i,
+                    model_stage="graph",
+                    execution_type=StageExecutionType.GRAPH,
+                    final_output=(not chain or i == stages - 1),
+                    final_output_type="latent",
+                    input_sources=(i - 1,) if chain and i else (),
+                    custom_process_input_func=(
+                        "tests.engine.test_graph_backend.graph_handoff" if chain and i else None
+                    ),
+                )
+                for i in range(stages)
+            ),
+        )
+        deploy = DeployConfig(
+            async_chunk=False,
+            stages=[
+                StageDeployConfig(
+                    stage_id=i,
+                    backend={
+                        "name": "external.graph.v1",
+                        "route": "ort-cpu",
+                        "manifest": str(manifest),
+                        "start_timeout_s": 5,
+                    },
+                    resource_budget={"capacities": {"host_ram": capacity}, "demands": {"host_ram": 512 << 20}},
+                )
+                for i in range(stages)
+            ],
+        )
+        if entrypoint is not None:
+            import yaml
+
+            from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+
+            monkeypatch.setitem(OMNI_PIPELINES, pipeline.model_type, pipeline)
+            (tmp_path / "config.json").write_text('{"model_type":"bert"}')
+            deploy_path = tmp_path / "deploy.yaml"
+            deploy_path.write_text(
+                yaml.safe_dump(
+                    {
+                        "pipeline": pipeline.model_type,
+                        "async_chunk": False,
+                        "stages": [
+                            {"stage_id": s.stage_id, "backend": s.backend, "resource_budget": s.resource_budget}
+                            for s in deploy.stages
+                        ],
+                    }
+                )
+            )
+            engine = entrypoint(
+                model=str(tmp_path), deploy_config=str(deploy_path), stage_init_timeout=10, init_timeout=30
+            )
+            runtimes.append(engine)
+            return engine
+        configs = [s.to_omegaconf() for s in merge_pipeline_deploy(pipeline, deploy)]
+        runtime = StageRuntime(configs, "local-graph", "", stage_init_timeout=5, async_chunk=False)
+        runtimes.append(runtime)
+        if start:
+            runtime.initialize()
+        return runtime
+
+    yield create
+    for runtime in runtimes:
+        runtime.shutdown()
+
+
+def test_contract_import_does_not_import_engine():
+    code = "import omni_stage_contracts; import sys; assert not ({'torch','vllm','vllm_omni'} & set(sys.modules))"
+    subprocess.run([sys.executable, "-S", "-c", code], cwd=ROOT, check=True)
+
+
+def test_device_init_lock_is_host_local_and_preserves_exclusion(tmp_path, monkeypatch):
+    import os
+
+    from vllm_omni import host
+    from vllm_omni._filelock_compat import flock_exclusive_nb, funlock
+    from vllm_omni.engine.stage_init_utils import device_init_lock_path, open_device_lock_file
+    from vllm_omni.host import paths
+
+    assert host is not None
+    monkeypatch.setattr(paths, "device_lock_directory", lambda: str(tmp_path))
+    lock_path = device_init_lock_path(0)
+    assert Path(lock_path).parent == tmp_path
+    first, writable = open_device_lock_file(lock_path)
+    second = None
+    try:
+        assert writable
+        flock_exclusive_nb(first)
+        second, _ = open_device_lock_file(lock_path)
+        with pytest.raises(BlockingIOError):
+            flock_exclusive_nb(second)
+        funlock(first)
+        flock_exclusive_nb(second)
+        funlock(second)
+    finally:
+        os.close(first)
+        if second is not None:
+            os.close(second)
+
+
+def test_graph_metadata_matches_structured_and_legacy_config():
+    from vllm_omni.config import VllmOmniConfig, VllmOmniGraphStageConfig
+    from vllm_omni.engine.stage_init_utils import (
+        extract_legacy_stage_metadata,
+        extract_stage_metadata_from_omni_stage_config,
+    )
+
+    pipeline = PipelineConfig(
+        model_type="test_graph",
+        stages=(
+            StagePipelineConfig(
+                stage_id=0,
+                model_stage="graph",
+                execution_type=StageExecutionType.GRAPH,
+                final_output=True,
+                final_output_type="latent",
+            ),
+        ),
+    )
+    deploy = DeployConfig(
+        async_chunk=False,
+        stages=[
+            StageDeployConfig(
+                stage_id=0,
+                backend={"name": "external.graph.v1"},
+                resource_budget={"capacities": {"ram": 1024}, "demands": {"ram": 512}},
+            )
+        ],
+    )
+    config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
+    assert isinstance(config, VllmOmniGraphStageConfig)
+    structured = extract_stage_metadata_from_omni_stage_config(config)
+    legacy = extract_legacy_stage_metadata(merge_pipeline_deploy(pipeline, deploy)[0].to_omegaconf())
+    assert structured.stage_type == legacy.stage_type == "graph"
+    assert type(structured.default_sampling_params) is type(legacy.default_sampling_params)
+
+
+def test_language_independent_contract_fixtures():
+    from omni_stage_contracts import StageEvent, StageRequest
+
+    fixture = json.loads((ROOT / "packages/omni-stage-contracts/conformance/v1.json").read_text())
+    negotiate(fixture["protocol_version"], fixture["required_features"])
+    BufferRef(**fixture["buffer"])
+    StageRequest(**fixture["request"])
+    StageEvent(**fixture["event"])
+    state = StateHandle(**fixture["state"])
+    for override in fixture["reject_buffer_overrides"]:
+        with pytest.raises(ValueError):
+            BufferRef(**{**fixture["buffer"], **override})
+    for override in fixture["stale_state_overrides"]:
+        assert not state.accepts(StateHandle(**{**fixture["state"], **override}))
+
+
+def test_external_weight_payload_must_be_in_manifest(tmp_path):
+    onnx = pytest.importorskip("onnx")
+    from onnx import helper, numpy_helper
+
+    from vllm_omni.edge.local.external.worker_ort import _verify_artifact_members
+
+    weight = numpy_helper.from_array(np.ones((32, 32), dtype=np.float32), "w")
+    graph = helper.make_graph([helper.make_node("Add", ["x", "w"], ["y"])], "external", [], [], [weight])
+    path = tmp_path / "model.onnx"
+    onnx.save_model(
+        helper.make_model(graph),
+        path,
+        save_as_external_data=True,
+        all_tensors_to_one_file=True,
+        location="weights.bin",
+        size_threshold=0,
+    )
+    with pytest.raises(ValueError, match="absent from verified manifest"):
+        _verify_artifact_members(str(path), [str(path)])
+    _verify_artifact_members(str(path), [str(path), str(tmp_path / "weights.bin")])
+
+
+def test_contract_rejects_shape_and_generation_mismatch():
+    with pytest.raises(ValueError):
+        BufferRef("x", "a", "1", "float32", (2, 3), 23)
+    with pytest.raises(ValueError):
+        negotiate(2)
+    with pytest.raises(ValueError):
+        negotiate(1, ["zero-copy"])
+    a = StateHandle("s", "ort", "sha", worker_generation="old")
+    b = StateHandle("s", "ort", "sha", worker_generation="new")
+    assert not a.accepts(b)
+
+
+def test_atomic_shared_constraints_and_quarantine():
+    ledger = ResourceLedger({"ram": 100, "wsl": 60})
+    first = ledger.reserve("wsl", {"ram": 50, "wsl": 50})
+    with pytest.raises(ResourceUnavailable):
+        ledger.reserve("igpu+npu", {"ram": 60})
+    with pytest.raises(ResourceUnavailable):
+        ledger.reserve("second-wsl", {"ram": 10, "wsl": 20})
+    assert ledger.snapshot()["reserved"] == {"ram": 50, "wsl": 50}
+    assert not ledger.release(first, drained=False)
+    assert ledger.snapshot()["quarantined"] == ["wsl"]
+    assert ledger.release(first, drained=True)
+    second = ledger.reserve("wsl", {"ram": 5})
+    assert not ledger.release(first, drained=True)
+    assert ledger.snapshot()["reserved"]["ram"] == 5
+    ledger.release(second, drained=True)
+
+
+@pytest.mark.parametrize(
+    "kinds",
+    [
+        ("cpu",),
+        ("cpu", "gpu_integrated"),
+        ("cpu", "npu"),
+        ("cpu", "gpu_integrated", "npu"),
+        ("cpu", "gpu_integrated", "npu", "gpu_discrete"),
+    ],
+)
+def test_integrated_topology_has_one_physical_ram(kinds):
+    from vllm_omni.edge.topology import describe_topology
+
+    devices = [
+        SimpleNamespace(device_id=k, kind=k, extra={}, memory_pool="vram" if k == "gpu_discrete" else "host_ram")
+        for k in kinds
+    ]
+    descriptions = describe_topology(devices, machine_id="pc", domain_id="windows")
+    for descriptor in descriptions:
+        if descriptor.integration != "discrete":
+            assert descriptor.memory_pool_ids == ("pc:ram",)
+        assert descriptor.power_domain_ids == ()  # unknown is not invented
+
+
+async def wait_output(pool):
+    async def poll():
+        while (out := pool.poll_graph_output(0)) is None:
+            await asyncio.sleep(0.005)
+        return out
+
+    return await asyncio.wait_for(poll(), timeout=10)
+
+
+@pytest.mark.asyncio
+async def test_runtime_load_submit_lease_and_shutdown(graph_runtime):
+    runtime = graph_runtime()
+    pool = runtime.stage_pools[0]
+    state = SimpleNamespace(sampling_params_list=[None])
+    prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}
+    await pool.submit_initial("a", state, prompt)
+    output = await wait_output(pool)
+    np.testing.assert_array_equal(output.custom_output["tensors"]["x_out"], 2)
+    with pytest.raises(ResourceUnavailable):
+        await pool.submit_initial("b", state, prompt)
+    output.release_stage_buffers()
+    await asyncio.sleep(0)
+    await pool.submit_initial("b", state, prompt)
+    (await wait_output(pool)).release_stage_buffers()
+    runtime.shutdown()
+    assert runtime.resource_ledger.snapshot()["reserved"] == {"host_ram": 0}
+
+
+def test_joint_admission_before_any_worker(graph_runtime):
+    runtime = graph_runtime(stages=2, chain=True, capacity=600 << 20, start=False)
+    with pytest.raises(ResourceUnavailable):
+        runtime.initialize()
+    assert runtime.resource_ledger.snapshot()["owners"] == []
+    assert runtime.stage_pools == []
+
+
+@pytest.mark.parametrize("fanout", [False, True])
+def test_unimplemented_graph_topologies_refused_before_loading(graph_runtime, fanout):
+    runtime = graph_runtime(stages=3, chain=fanout, start=False)
+    if fanout:
+        runtime._stage_configs[2].engine_input_source = [0]
+    with pytest.raises(ValueError, match="linear chain with exactly one final output"):
+        runtime.initialize()
+    assert runtime.resource_ledger is None
+    assert runtime.stage_pools == []
+
+
+def test_placement_failure_releases_load(graph_runtime):
+    runtime = graph_runtime(worker_config={"load": {"fraction_on_target": 0.0}}, start=False)
+    with pytest.raises(RuntimeError, match="REFUSE_EP_PLACEMENT"):
+        runtime.initialize()
+    assert runtime.resource_ledger.snapshot()["owners"] == []
+
+
+@pytest.mark.asyncio
+async def test_worker_failure_produces_error_not_success(graph_runtime):
+    runtime = graph_runtime(worker_config={"die_on": "run"})
+    pool = runtime.stage_pools[0]
+    await pool.submit_initial(
+        "a", SimpleNamespace(sampling_params_list=[None]), {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}
+    )
+    output = await wait_output(pool)
+    assert output.finished and output.error
+
+
+@pytest.mark.asyncio
+async def test_cancel_fences_inflight_graph(graph_runtime):
+    runtime = graph_runtime(worker_config={"run": {"sleep_s": 30}})
+    pool = runtime.stage_pools[0]
+    await pool.submit_initial(
+        "a", SimpleNamespace(sampling_params_list=[None]), {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}
+    )
+    await asyncio.sleep(0.05)
+    await asyncio.wait_for(pool.abort_requests(["a"]), timeout=8)
+    assert pool.poll_graph_output(0) is None
+    assert runtime.resource_ledger.snapshot()["reserved"]["host_ram"] == 0
+
+
+def graph_handoff(outputs, prompt, requires_multimodal_data):
+    return {"tensors": {"x": outputs[0].custom_output["tensors"]["x_out"]}}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("event_driven", [False, True])
+async def test_two_graphs_through_existing_orchestrator(graph_runtime, event_driven):
+    from vllm_omni.engine.messages import StageSubmissionMessage
+    from vllm_omni.engine.orchestrator import Orchestrator
+
+    runtime = graph_runtime(stages=2, chain=True)
+    request_queue, output_queue, rpc_queue = asyncio.Queue(), asyncio.Queue(), asyncio.Queue()
+    orch = Orchestrator(request_queue, output_queue, rpc_queue, runtime.stage_pools)
+    loop = asyncio.create_task(orch._orchestration_loop_event_driven() if event_driven else orch._orchestration_loop())
+    prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}
+    try:
+        await orch._handle_add_request(
+            StageSubmissionMessage(
+                type="add_request",
+                request_id="pipeline",
+                prompt=prompt,
+                original_prompt=prompt,
+                output_prompt_text=None,
+                sampling_params_list=[None, None],
+                final_stage_id=1,
+                final_output_stage_ids=[1],
+                preprocess_ms=0,
+                request_timestamp=0,
+                enqueue_ts=0,
+            )
+        )
+        from vllm_omni.engine.messages import StageMetricsMessage
+
+        while True:
+            result = await asyncio.wait_for(output_queue.get(), timeout=10)
+            if not isinstance(result, StageMetricsMessage):
+                break
+        assert result.finished
+        np.testing.assert_array_equal(result.engine_outputs.custom_output["tensors"]["x_out"], 4)
+        result.engine_outputs.release_stage_buffers()
+        await asyncio.sleep(0)
+        assert not orch.request_states
+        assert all(pool.stage_client._active is None for pool in runtime.stage_pools)
+    finally:
+        orch._shutdown_event.set()
+        await asyncio.wait_for(loop, timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_stale_ack_cannot_release_reused_request_id(graph_runtime):
+    runtime = graph_runtime()
+    pool = runtime.stage_pools[0]
+    state = SimpleNamespace(sampling_params_list=[None])
+    prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}
+    await pool.submit_initial("same", state, prompt)
+    output = await wait_output(pool)
+    stale_ack = output._stage_release
+    output.release_stage_buffers()
+    await asyncio.sleep(0)
+    await pool.submit_initial("same", state, prompt)
+    next_output = await wait_output(pool)
+    stale_ack()
+    await asyncio.sleep(0)
+    assert pool.stage_client._active == "same"
+    next_output.release_stage_buffers()
+
+
+@pytest.mark.asyncio
+async def test_real_ort_cpu_graph_through_stage_runtime(graph_runtime):
+    runtime = graph_runtime(real_ort=True)
+    pool = runtime.stage_pools[0]
+    await pool.submit_initial(
+        "ort", SimpleNamespace(sampling_params_list=[None]), {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}
+    )
+    output = await wait_output(pool)
+    assert not output.error
+    np.testing.assert_array_equal(output.custom_output["tensors"]["x_out"], 2)
+    assert pool.stage_client.load_report.target_nodes > 0
+    output.release_stage_buffers()
+
+
+@pytest.mark.asyncio
+async def test_public_async_omni_graph_requests(graph_runtime):
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+    engine = graph_runtime(entrypoint=AsyncOmni)
+    for request_id in ("first", "second"):
+        outputs = [
+            out
+            async for out in engine.generate(
+                {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}, request_id=request_id
+            )
+        ]
+        assert len(outputs) == 1
+        assert outputs[0].request_id == request_id
+        assert outputs[0].custom_output["tensors"]["x_out"].shape == (2, 3)
+
+
+def test_public_sync_omni_graph_requests(graph_runtime):
+    from vllm_omni.entrypoints.omni import Omni
+
+    engine = graph_runtime(entrypoint=Omni)
+    for _ in range(2):
+        outputs = engine.generate({"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}, use_tqdm=False)
+        assert len(outputs) == 1
+        assert outputs[0].custom_output["tensors"]["x_out"].shape == (2, 3)
+
+
+@pytest.mark.asyncio
+async def test_graph_ingress_stays_bounded_until_consumer_ack(graph_runtime):
+    from vllm_omni.engine.messages import OutputMessage
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+    frontend = graph_runtime(entrypoint=AsyncOmni)
+    engine = frontend.engine
+    prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}
+    engine.add_request("slow", prompt)
+    for i in range(100):
+        with pytest.raises(ResourceUnavailable, match="capacity is one"):
+            engine.add_request(f"blocked-{i}", prompt)
+
+    async def take():
+        while True:
+            msg = await engine.try_get_output_async()
+            if isinstance(msg, OutputMessage):
+                return msg.engine_outputs
+            await asyncio.sleep(0.005)
+
+    output = await asyncio.wait_for(take(), 5)
+    with pytest.raises(ResourceUnavailable):
+        engine.add_request("still-blocked", prompt)
+    stale_release = output._stage_release
+    output.release_stage_buffers()
+    await asyncio.sleep(0.01)
+    engine.add_request("next", prompt)
+    stale_release()
+    with pytest.raises(ResourceUnavailable):
+        engine.add_request("late-ack-must-not-release-next", prompt)
+    (await asyncio.wait_for(take(), 5)).release_stage_buffers()

tests/windows/test_shared_buffers.py

@@ -0,0 +1,75 @@
+"""Ownership/framing tests also run on POSIX; native connector tests run on Windows."""
+
+import subprocess
+import sys
+import uuid
+
+import pytest
+
+from vllm_omni.host.shared_memory import SharedBufferOwner, read_buffer
+
+
+def test_producer_retains_until_a_different_process_reads():
+    owner = SharedBufferOwner()
+    meta = owner.write(b"payload", f"omni-test-{uuid.uuid4().hex}")
+    try:
+        assert owner.retained_bytes > 0
+        code = (
+            "from vllm_omni.host.shared_memory import read_buffer; assert read_buffer(" + repr(meta) + ") == b'payload'"
+        )
+        subprocess.run([sys.executable, "-c", code], check=True, timeout=60)
+        owner.collect()
+        assert owner.retained_bytes == 0
+    finally:
+        owner.release(meta["name"])
+
+
+def test_bounds_stale_generation_and_name_collision():
+    owner = SharedBufferOwner(max_bytes=128, max_buffers=1)
+    meta = owner.write(b"one", f"omni-test-{uuid.uuid4().hex}")
+    try:
+        with pytest.raises(FileExistsError):
+            SharedBufferOwner().write(b"collision", meta["name"])
+        with pytest.raises(ValueError, match="stale"):
+            read_buffer({**meta, "generation": "wrong"})
+        with pytest.raises(BufferError):
+            owner.write(b"two")
+        assert owner.retained_bytes > 0
+        assert read_buffer(meta) == b"one"
+        owner.collect()
+        assert owner.retained_bytes == 0
+        with pytest.raises(BufferError):
+            owner.write(b"x" * 129)
+    finally:
+        owner.release(meta["name"])
+
+
+@pytest.mark.skipif(sys.platform != "win32", reason="native Windows path and mapping lifetime")
+def test_native_connector_roundtrip_and_cleanup():
+    from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+    from vllm_omni.host.shared_memory import owner
+    from vllm_omni.windows.paths import shm_dir
+
+    assert not shm_dir().startswith("/dev/shm")
+    producer, consumer = SharedMemoryConnector({}), SharedMemoryConnector({})
+    key = f"omni-test-{uuid.uuid4().hex}"
+    try:
+        ok, size, _ = producer.put("0", "1", key, {"data": [1, 2, 3]})
+        assert ok and owner.retained_bytes > 0
+        assert consumer.get("0", "1", key) == ({"data": [1, 2, 3]}, size)
+        owner.collect()
+        assert owner.retained_bytes == 0
+        assert producer.put("0", "1", key, {"unconsumed": True})[0]
+        assert owner.retained_bytes > 0
+        producer.cleanup(key)
+        assert owner.retained_bytes == 0
+        assert consumer.get("0", "1", key) is None
+        for index in range(8):
+            next_key = f"{key}_{index}"
+            assert producer.put("0", "1", next_key, {"index": index})[0]
+            assert len(producer._pending_keys) <= 1
+            assert consumer.get("0", "1", next_key)[0] == {"index": index}
+            owner.collect()
+    finally:
+        producer.close()
+        consumer.close()

tests/windows/test_windows_compat.py

@@ -7,6 +7,7 @@ that decide whether a vLLM function is replaced or left alone."""
 from __future__ import annotations
 
 import asyncio
+import errno
 import os
 import sys
 import tempfile
@@ -34,7 +35,7 @@ class FakeLocking:
         self.calls.append((fd, mode))
         if mode == self.NBLCK:
             if fd in self.held:
-                raise OSError(36, "Resource deadlock avoided")
+                raise OSError(errno.EACCES, "Lock held")
             self.held.add(fd)
         elif mode == self.UNLCK:
             self.held.discard(fd)
@@ -86,6 +87,20 @@ def test_flock_blocking_waits_until_released(lock_fd):
     assert len(sleeps) == 2 and lock_fd in locking.held
 
 
+@pytest.mark.parametrize("operation", [shims.LOCK_EX, shims.LOCK_EX | shims.LOCK_NB])
+def test_flock_permanent_failure_does_not_retry(lock_fd, operation):
+    sleeps = []
+
+    def invalid_lock(*args):
+        raise OSError(errno.EINVAL, "Unsupported lock path")
+
+    shim = shims.FcntlShim(invalid_lock, 2, 0, sleep=sleeps.append)
+    with pytest.raises(OSError) as failure:
+        shim.flock(lock_fd, operation)
+    assert failure.value.errno == errno.EINVAL
+    assert sleeps == []
+
+
 def test_shared_lock_is_served_as_exclusive_and_counted(lock_fd):
     locking, sleeps = FakeLocking(), []
     s = _shim(locking, sleeps)
@@ -237,7 +252,7 @@ def test_shm_dir_exists_and_path_joins():
     d = shm_dir()
     assert os.path.isdir(d)
     assert shm_path("x.lock").startswith(d)
-    if os.path.isdir("/dev/shm"):
+    if sys.platform != "win32" and os.path.isdir("/dev/shm"):
         assert d == "/dev/shm"
 
 

tests/worker/test_forward_kwargs_filtering.py

@@ -43,9 +43,7 @@ def test_a_model_taking_kwargs_is_left_alone():
 
 def test_a_vllm_native_model_reports_exactly_its_parameters():
     accepted = _Probe(_VLLMNativeModel())._accepted_forward_kwargs()
-    assert accepted == frozenset(
-        {"input_ids", "positions", "intermediate_tensors", "inputs_embeds"}
-    )
+    assert accepted == frozenset({"input_ids", "positions", "intermediate_tensors", "inputs_embeds"})
     for omni_only in ("sampling_metadata", "logits_index", "sampler"):
         assert omni_only not in accepted
 
@@ -67,3 +65,15 @@ def test_an_uninspectable_forward_falls_back_to_passing_everything():
 
     accepted = _Probe(Weird())._accepted_forward_kwargs()
     assert accepted is None or "sampling_metadata" not in accepted
+
+
+def test_cuda_graph_wrapper_preserves_underlying_forward_signature():
+    from vllm.compilation.cuda_graph import CUDAGraphWrapper
+
+    # No CUDA context is needed for attribute delegation. Use the actual vLLM
+    # wrapper rather than a fake that could conceal API changes.
+    wrapper = object.__new__(CUDAGraphWrapper)
+    wrapper.runnable = _VLLMNativeModel()
+    assert _Probe(wrapper)._accepted_forward_kwargs() == _Probe(wrapper.runnable)._accepted_forward_kwargs()
+    wrapper.runnable = _OmniNativeModel()
+    assert _Probe(wrapper)._accepted_forward_kwargs() is None

vllm_omni/config/__init__.py

@@ -19,6 +19,7 @@ from vllm_omni.config.omni_config import (
     VllmOmniConfig,
     VllmOmniDiffusionStageConfig,
     VllmOmniGenerationStageConfig,
+    VllmOmniGraphStageConfig,
     VllmOmniOrchestratorConfig,
 )
 from vllm_omni.config.stage_config import (
@@ -60,6 +61,7 @@ __all__ = [
     "BaseVllmOmniStageConfig",
     "VllmOmniARStageConfig",
     "VllmOmniGenerationStageConfig",
+    "VllmOmniGraphStageConfig",
     "VllmOmniDiffusionStageConfig",
     "StageConfigType",
     "OmniConfigResolution",

vllm_omni/config/omni_config.py

@@ -56,6 +56,7 @@ _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str |
     StageExecutionType.LLM_AR: (StageType.LLM, "ar"),
     StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"),
     StageExecutionType.DIFFUSION: (StageType.DIFFUSION, None),
+    StageExecutionType.GRAPH: (StageType.GRAPH, None),
 }
 
 _PIPELINE_DEPLOY_CLI_FIELDS = PIPELINE_WIDE_ENGINE_FIELDS
@@ -197,6 +198,8 @@ class _SchedulerEngineOverrides(TypedDict, total=False):
 
 
 class _RuntimeEngineOverrides(TypedDict, total=False):
+    backend: dict[str, Any]
+    resource_budget: dict[str, Any]
     additional_config: dict[str, Any]
     distributed_executor_backend: Any
     worker_cls: str
@@ -528,6 +531,8 @@ class OmniStageRuntimeConfig:
 
     # LLM backend extensions; diffusion owns these in its config projection.
     additional_config: dict[str, Any] | None = None
+    backend: dict[str, Any] | None = None
+    resource_budget: dict[str, Any] | None = None
     distributed_executor_backend: Any = None
     worker_cls: str | None = None
     devices: str | None = None
@@ -1142,19 +1147,22 @@ _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
 # validation must use the execution-type-specific sets below.
 _STAGE_ENGINE_FIELDS = _LLM_STAGE_ENGINE_FIELDS | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS
 _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
-    StageExecutionType.LLM_AR: _LLM_STAGE_ENGINE_FIELDS,
-    StageExecutionType.LLM_GENERATION: _LLM_STAGE_ENGINE_FIELDS,
-    StageExecutionType.DIFFUSION: _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS,
+    StageExecutionType.LLM_AR: _LLM_STAGE_ENGINE_FIELDS - {"backend"},
+    StageExecutionType.LLM_GENERATION: _LLM_STAGE_ENGINE_FIELDS - {"backend"},
+    StageExecutionType.DIFFUSION: _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS - {"backend"},
+    StageExecutionType.GRAPH: frozenset({"model", "backend", "resource_budget", "log_level", "log_stats"}),
 }
 _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
     StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
     StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
     StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS,
+    StageExecutionType.GRAPH: frozenset(),
 }
 _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = {
     StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS,
     StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS,
     StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP),
+    StageExecutionType.GRAPH: frozenset(),
 }
 
 _LOAD_STAGE_ENGINE_FIELD_MAP = {
@@ -1254,6 +1262,10 @@ def _stage_engine_values(
     stage_cli_overrides: Mapping[str, Any] | None = None,
 ) -> _StageEngineValues:
     engine = _stage_engine_overrides(stage_deploy)
+    if topology.execution_type == StageExecutionType.GRAPH and engine.get("silence_ban_frames") == 0:
+        # StageDeployConfig carries this speech-specific default even when no
+        # speech option was supplied. It is not an external-graph argument.
+        engine.pop("silence_ban_frames")
     # Preserve legacy ordering: topology-owned KV roles override deploy
     # extras, while an explicit CLI override remains highest priority.
     if topology.omni_kv_config:
@@ -1467,7 +1479,14 @@ class VllmOmniDiffusionStageConfig(BaseVllmOmniStageConfig):
     diffusion_config: _DiffusionConfigProjection = field(default_factory=_DiffusionConfigProjection)
 
 
-StageConfigType: TypeAlias = VllmOmniARStageConfig | VllmOmniGenerationStageConfig | VllmOmniDiffusionStageConfig
+@config(config=ConfigDict(arbitrary_types_allowed=True))
+class VllmOmniGraphStageConfig(BaseVllmOmniStageConfig):
+    """A whole external graph with an explicit backend and resource budget."""
+
+
+StageConfigType: TypeAlias = (
+    VllmOmniARStageConfig | VllmOmniGenerationStageConfig | VllmOmniDiffusionStageConfig | VllmOmniGraphStageConfig
+)
 
 
 def _build_common_stage_config_kwargs(
@@ -1628,10 +1647,27 @@ def _build_diffusion_stage_config(
     )
 
 
+def _build_graph_stage_config(pipeline, deploy, topology, stage_deploy, engine, *, model):
+    if deploy.async_chunk or deploy.session_mode == "duplex":
+        raise ValueError("external graph v1 supports complete, stateless requests with async_chunk=False")
+    if stage_deploy is None or not engine.runtime.get("backend") or not engine.runtime.get("resource_budget"):
+        raise ValueError("graph stage requires backend and resource_budget")
+    common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs(
+        pipeline,
+        deploy,
+        topology,
+        stage_deploy,
+        engine,
+        model=model,
+    )
+    return _with_resolved_processors(VllmOmniGraphStageConfig(**common_kwargs), input_proc, next_stage_proc)
+
+
 _STAGE_CONFIG_BUILDERS = {
     StageExecutionType.LLM_AR: _build_ar_stage_config,
     StageExecutionType.LLM_GENERATION: _build_generation_stage_config,
     StageExecutionType.DIFFUSION: _build_diffusion_stage_config,
+    StageExecutionType.GRAPH: _build_graph_stage_config,
 }
 
 

vllm_omni/config/stage_config.py

@@ -189,6 +189,7 @@ class StageType(str, Enum):
     # TODO(@lishunyang12): remove once all models migrate to StageExecutionType
     LLM = "llm"
     DIFFUSION = "diffusion"
+    GRAPH = "graph"
 
 
 class StageExecutionType(str, Enum):
@@ -197,6 +198,7 @@ class StageExecutionType(str, Enum):
     LLM_AR = "llm_ar"
     LLM_GENERATION = "llm_generation"
     DIFFUSION = "diffusion"
+    GRAPH = "graph"
 
 
 def _resolve_scheduler(
@@ -386,6 +388,8 @@ class StageDeployConfig:
     devices: str | None = None
     num_replicas: int = 1
     env: dict[str, Any] | None = None
+    backend: dict[str, Any] | None = None
+    resource_budget: dict[str, Any] | None = None
 
     # Inter-stage connector wiring and request defaults.
     output_connectors: dict[str, str] | None = None
@@ -867,6 +871,7 @@ _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str |
     StageExecutionType.LLM_AR: (StageType.LLM, "ar"),
     StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"),
     StageExecutionType.DIFFUSION: (StageType.DIFFUSION, None),
+    StageExecutionType.GRAPH: (StageType.GRAPH, None),
 }
 
 

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -3,9 +3,8 @@
 
 # [edge-infer W1] This module is in the import closure of AsyncOmni, so a
 # module-scope `import fcntl` makes the whole engine unimportable on Windows
-# -- before anyone has chosen a connector. The transport itself is genuinely
-# POSIX (it hardcodes /dev/shm lock paths), so the honest Windows behaviour
-# is to stay importable and refuse at construction, not to fake a port.
+# -- before anyone has chosen a connector. The Windows plugin supplies locks;
+# host helpers keep mappings alive until the single consumer acknowledges them.
 try:
     import fcntl
 except ImportError:  # Windows
@@ -15,6 +14,7 @@ from multiprocessing import shared_memory as shm_pkg
 from typing import Any
 
 from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes
+from vllm_omni.windows.paths import shm_path
 
 from ..utils.logging import get_connector_logger
 from .base import OmniConnectorBase
@@ -25,7 +25,7 @@ logger = get_connector_logger(__name__)
 class SharedMemoryConnector(OmniConnectorBase):
     """Key-addressed local shared-memory connector.
 
-    SHM is a local-only transport: it reads/writes POSIX shared memory
+    SHM is a local-only transport: it reads/writes host shared memory
     segments identified purely by *key*.  It does **not** understand
     remote-transport metadata such as ``source_host`` / ``source_port``
     (that is the RDMA connector's job).  When such metadata is passed in,
@@ -35,9 +35,8 @@ class SharedMemoryConnector(OmniConnectorBase):
     def __init__(self, config: dict[str, Any]):
         if fcntl is None:
             raise RuntimeError(
-                "SharedMemoryConnector requires POSIX shared memory and fcntl "
-                "locking (it addresses lock files under /dev/shm); neither exists "
-                "on Windows. Select a different connector for this stage."
+                "SharedMemoryConnector requires host file-lock support. "
+                "On Windows, activate the vllm_omni.windows compatibility layer before importing connectors."
             )
         self.config = config
         self.stage_id = config.get("stage_id", -1)
@@ -59,7 +58,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             payload = self.serialize_obj(data)
             size = len(payload)
 
-            lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"
+            lock_file = shm_path(f"shm_{put_key}_lockfile.lock")
             with open(lock_file, "wb+") as lockf:
                 fcntl.flock(lockf, fcntl.LOCK_EX)
                 meta = shm_write_bytes(payload, name=put_key)
@@ -67,6 +66,12 @@ class SharedMemoryConnector(OmniConnectorBase):
 
             # meta contains {'name': ..., 'size': ...}
             metadata = {"shm": meta, "size": size}
+            if os.name == "nt":
+                from vllm_omni.host.shared_memory import owner
+
+                # ACKs arrive in another process. Do not retain an unbounded
+                # history of keys after their producer handles are released.
+                self._pending_keys.intersection_update(owner.retained_names)
             self._pending_keys.add(put_key)
 
             self._metrics["puts"] += 1
@@ -86,7 +91,7 @@ class SharedMemoryConnector(OmniConnectorBase):
                 data_bytes = shm_read_bytes(shm_handle)
                 fcntl.flock(lockf, fcntl.LOCK_UN)
             obj = self.deserialize_obj(data_bytes)
-            result = (obj, int(shm_handle.get("size", 0)))
+            result = (obj, len(data_bytes))
             deserialized = True
             return result
         except Exception as e:
@@ -106,7 +111,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             shm = shm_pkg.SharedMemory(name=get_key)
             if shm is None or shm.size == 0:
                 return None
-            lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock"
+            lock_file = shm_path(f"shm_{get_key}_lockfile.lock")
             shm_handle = {"name": get_key, "size": shm.size}
             result = self._get_data_with_lock(lock_file, shm_handle)
             if result is not None:
@@ -142,7 +147,7 @@ class SharedMemoryConnector(OmniConnectorBase):
 
             if isinstance(metadata, dict) and "shm" in metadata:
                 shm_handle = metadata["shm"]
-                lock_file = f"/dev/shm/shm_{shm_handle['name']}_lockfile.lock"
+                lock_file = shm_path(f"shm_{shm_handle['name']}_lockfile.lock")
                 result = self._get_data_with_lock(lock_file, shm_handle)
                 if result is not None:
                     self._pending_keys.discard(get_key)
@@ -171,6 +176,7 @@ class SharedMemoryConnector(OmniConnectorBase):
         ]
         for key in stale:
             self._pending_keys.discard(key)
+            self._release_owned_buffer(key)
             try:
                 seg = shm_pkg.SharedMemory(name=key)
                 seg.close()
@@ -180,7 +186,7 @@ class SharedMemoryConnector(OmniConnectorBase):
                 pass
             except Exception as e:
                 logger.debug("cleanup: failed to unlink SHM segment %s: %s", key, e)
-            lock_file = f"/dev/shm/shm_{key}_lockfile.lock"
+            lock_file = shm_path(f"shm_{key}_lockfile.lock")
             if os.path.exists(lock_file):
                 try:
                     os.remove(lock_file)
@@ -190,13 +196,14 @@ class SharedMemoryConnector(OmniConnectorBase):
     def close(self) -> None:
         """Unlink all remaining tracked SHM segments."""
         for key in list(self._pending_keys):
+            self._release_owned_buffer(key)
             try:
                 seg = shm_pkg.SharedMemory(name=key)
                 seg.close()
                 seg.unlink()
             except Exception:
                 pass
-            lock_file = f"/dev/shm/shm_{key}_lockfile.lock"
+            lock_file = shm_path(f"shm_{key}_lockfile.lock")
             if os.path.exists(lock_file):
                 try:
                     os.remove(lock_file)
@@ -204,5 +211,12 @@ class SharedMemoryConnector(OmniConnectorBase):
                     pass
         self._pending_keys.clear()
 
+    @staticmethod
+    def _release_owned_buffer(key: str) -> None:
+        if os.name == "nt":
+            from vllm_omni.host.shared_memory import owner
+
+            owner.release(key)
+
     def health(self) -> dict[str, Any]:
         return {"status": "healthy", **self._metrics}

vllm_omni/edge/local/external/client.py

@@ -28,10 +28,12 @@ ledger, not a rounding error. So the worker reports its own RSS and
 
 from __future__ import annotations
 
+import json
 import secrets
 import socket
 import subprocess
-import sys
+import tempfile
+import threading
 import time
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
@@ -42,6 +44,7 @@ import numpy as np
 from vllm_omni.edge.local.external import launch as _launch
 from vllm_omni.edge.local.external import protocol as _proto
 from vllm_omni.edge.local.external.protocol import ProtocolError, WorkerError
+from vllm_omni.host.process import spawn_worker, terminate_worker
 
 DEFAULT_START_TIMEOUT_S = 120.0
 """Generous on purpose: the first VitisAI session on a cold machine spends
@@ -212,12 +215,17 @@ class ExternalWorker:
         *,
         start_timeout_s: float = DEFAULT_START_TIMEOUT_S,
         call_timeout_s: float = DEFAULT_CALL_TIMEOUT_S,
+        max_payload_bytes: int = _proto.MAX_PAYLOAD_BYTES,
     ) -> None:
         if not route.available:
             raise WorkerStartError(f"route {route.name} is not available: {route.reason}")
         self.route = route
         self.start_timeout_s = start_timeout_s
         self.call_timeout_s = call_timeout_s
+        self.max_payload_bytes = max_payload_bytes
+        self._log = None
+        self._lifecycle_lock = threading.RLock()
+        self._drained = False
         self._process: subprocess.Popen[bytes] | None = None
         self._sock: socket.socket | None = None
         self._listener: socket.socket | None = None
@@ -228,6 +236,8 @@ class ExternalWorker:
 
     def start(self) -> dict[str, Any]:
         """Launch the worker and complete the handshake. Returns its ``hello``."""
+        if self._process is not None:
+            raise WorkerStartError("worker instances cannot be started twice")
         token = secrets.token_hex(16)
         host = _bind_address(is_windows=self.route.is_windows)
         listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -240,19 +250,29 @@ class ExternalWorker:
 
         argv = [
             str(self.route.interpreter),
+            _launch.to_worker_path(
+                Path(__file__).parents[3] / "host/worker_bootstrap.py", is_windows=self.route.is_windows
+            ),
             _launch.to_worker_path(self.route.worker, is_windows=self.route.is_windows),
-            "--host", host,
-            "--port", str(port),
-            "--token", token,
+            "--host",
+            host,
+            "--port",
+            str(port),
+            "--token",
+            token,
         ]
-        self._process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        self._log = tempfile.TemporaryFile()
+        try:
+            self._process = spawn_worker(argv, stdout=self._log, stderr=self._log)
+        except BaseException:
+            self.close()
+            raise
 
         try:
             conn, _ = listener.accept()
-        except socket.timeout as exc:
+        except TimeoutError as exc:
             raise WorkerStartError(
-                f"worker did not connect back within {self.start_timeout_s:.0f}s: "
-                f"{self._drain_process()}"
+                f"worker did not connect back within {self.start_timeout_s:.0f}s: {self._drain_process()}"
             ) from exc
         finally:
             listener.close()
@@ -262,7 +282,11 @@ class ExternalWorker:
         conn.settimeout(self.call_timeout_s)
         self._sock = conn
 
-        op, body, _ = _proto.recv_message(conn)
+        try:
+            op, body, _ = _proto.recv_message(conn, max_payload_bytes=0)
+        except BaseException:
+            self.terminate()
+            raise
         if op != _proto.OP_HELLO:
             self.close()
             raise WorkerStartError(f"first message was {op!r}, expected {_proto.OP_HELLO!r}")
@@ -275,20 +299,74 @@ class ExternalWorker:
         return body
 
     def _drain_process(self) -> str:
+        with self._lifecycle_lock:
+            return self._drain_process_locked()
+
+    def _drain_process_locked(self) -> str:
         process = self._process
         if process is None:
             return "worker was never started"
-        if process.poll() is None:
-            process.kill()
-        try:
-            out, err = process.communicate(timeout=5)
-        except subprocess.TimeoutExpired:  # pragma: no cover - kill just ran
+        if not terminate_worker(process):
             return f"worker pid {process.pid} did not exit"
-        text = (err or b"").decode("utf-8", "replace").strip()
-        if not text:
-            text = (out or b"").decode("utf-8", "replace").strip()
+        text = ""
+        if self._log is not None:
+            self._log.seek(0, 2)
+            self._log.seek(max(0, self._log.tell() - 2000))
+            text = self._log.read().decode("utf-8", "replace").strip()
         return f"exit={process.returncode} output={text[:2000]!r}" if text else f"exit={process.returncode}"
 
+    def terminate(self) -> bool:
+        """Interrupt a blocking native call without sending on its data socket.
+
+        The caller retains reservations if the process tree did not drain.
+        This method may run concurrently with the one in-flight `_call`.
+        """
+        with self._lifecycle_lock:
+            return self._terminate_locked()
+
+    def _terminate_locked(self) -> bool:
+        if self._drained:
+            return True
+        native_drained = True
+        if self._process is not None and self.route.is_windows and _launch.is_wsl():
+            identity = self.hello.get("process_identity") or {}
+            native_drained = False
+            if type(identity.get("pid")) is int and type(identity.get("created_filetime")) is int:
+                try:
+                    helper = Path(__file__).parents[3] / "host/windows_process.py"
+                    result = subprocess.run(
+                        [
+                            str(self.route.interpreter),
+                            _launch.to_worker_path(helper, is_windows=True),
+                            str(identity["pid"]),
+                            str(identity["created_filetime"]),
+                        ],
+                        capture_output=True,
+                        timeout=20,
+                        text=True,
+                    )
+                    native_drained = result.returncode == 0 and json.loads(result.stdout).get("drained") is True
+                except (OSError, subprocess.SubprocessError, ValueError):
+                    pass
+        # Keep the data connection open until native tree retirement completes:
+        # closing it first lets an idle worker exit before we can verify its tree.
+        sock, self._sock = self._sock, None
+        if sock is not None:
+            try:
+                sock.shutdown(socket.SHUT_RDWR)
+            except OSError:
+                pass
+            sock.close()
+        if self._listener is not None:
+            self._listener.close()
+            self._listener = None
+        drained = (self._process is None or terminate_worker(self._process)) and native_drained
+        self._drained = drained
+        if drained and self._log is not None:
+            self._log.close()
+            self._log = None
+        return drained
+
     def close(self, timeout_s: float = DEFAULT_CLOSE_TIMEOUT_S) -> None:
         """Ask the worker to exit; kill it if it will not. Safe to call twice."""
         sock, self._sock = self._sock, None
@@ -304,17 +382,20 @@ class ExternalWorker:
         if self._listener is not None:
             self._listener.close()
             self._listener = None
-        process, self._process = self._process, None
+        process = self._process
         if process is not None and process.poll() is None:
             try:
                 process.wait(timeout=timeout_s)
             except subprocess.TimeoutExpired:
-                process.kill()
-                process.wait(timeout=5)
+                if not self.terminate():
+                    raise WorkerStartError("worker tree did not drain during close")
         if process is not None:
             for stream in (process.stdout, process.stderr):
                 if stream is not None:
                     stream.close()
+        if self._log is not None:
+            self._log.close()
+            self._log = None
 
     def __enter__(self) -> ExternalWorker:
         self.start()
@@ -336,12 +417,13 @@ class ExternalWorker:
             raise WorkerStartError("worker is not running; call start() first")
         try:
             _proto.send_message(sock, op, body, tensors)
-            reply_op, reply_body, reply_tensors = _proto.recv_message(sock)
+            reply_op, reply_body, reply_tensors = _proto.recv_message(sock, max_payload_bytes=self.max_payload_bytes)
         except (OSError, ProtocolError) as exc:
             # The socket is the only liveness signal we have; a dead worker
             # shows up as a broken pipe or a short read, and the useful part of
             # the diagnosis is on its stderr.
             detail = self._drain_process()
+            sock.close()
             self._sock = None
             raise WorkerStartError(f"worker connection failed ({exc}); {detail}") from exc
         if reply_op != _proto.OP_OK:
@@ -357,6 +439,7 @@ class ExternalWorker:
         profile_prefix: str | Path | None = None,
         ep_dir: str | Path | None = None,
         intra_op_num_threads: int | None = None,
+        artifact_files: list[Path] | None = None,
     ) -> LoadReport:
         """Open a graph and measure where its nodes actually ran.
 
@@ -371,13 +454,15 @@ class ExternalWorker:
             "profile": bool(profile),
         }
         if profile_prefix is not None:
-            body["profile_prefix"] = _launch.to_worker_path(
-                profile_prefix, is_windows=self.route.is_windows
-            )
+            body["profile_prefix"] = _launch.to_worker_path(profile_prefix, is_windows=self.route.is_windows)
         if ep_dir is not None:
             body["ep_dir"] = str(ep_dir)
         if intra_op_num_threads is not None:
             body["intra_op_num_threads"] = int(intra_op_num_threads)
+        if artifact_files is not None:
+            body["artifact_files"] = [
+                _launch.to_worker_path(path, is_windows=self.route.is_windows) for path in artifact_files
+            ]
 
         reply, _ = self._call(_proto.OP_LOAD, body, example_inputs or {})
         self.report = LoadReport.from_body(reply, ep=self.route.ep, graph_path=str(graph_path))

vllm_omni/edge/local/external/launch.py

@@ -29,6 +29,7 @@ different work, and the planner has to be able to tell the caller which it hit.
 
 from __future__ import annotations
 
+import getpass
 import os
 import shutil
 import subprocess
@@ -63,20 +64,25 @@ covered on a machine with no AMD hardware and no ORT at all."""
 # Where the routes live on this machine. Defaults, not requirements: the env
 # override above is the supported way to point at a different install, and the
 # reason string names this list when nothing is found.
+_WINDOWS_HOME = Path.home() if os.name == "nt" else Path("/mnt/c/Users") / getpass.getuser()
 _DEFAULT_CANDIDATES: dict[str, tuple[str, ...]] = {
     ROUTE_VITISAI: (
-        r"/mnt/c/Users/zhout/npu-ep/Scripts/python.exe",
-        r"/mnt/c/Users/zhout/ryzenai/Scripts/python.exe",
+        str(_WINDOWS_HOME / "npu-ep/Scripts/python.exe"),
+        str(_WINDOWS_HOME / "ryzenai/Scripts/python.exe"),
     ),
     ROUTE_DML: (
-        r"/mnt/c/Users/zhout/dml-ep/Scripts/python.exe",
-        r"/mnt/c/Users/zhout/directml/Scripts/python.exe",
+        str(_WINDOWS_HOME / "dml-ep/Scripts/python.exe"),
+        str(_WINDOWS_HOME / "directml/Scripts/python.exe"),
+    ),
+    ROUTE_TORCH_DML: () if os.name == "nt" else (".venvs/dml/bin/python",),
+    ROUTE_CPU: (
+        (".venvs/omni-cpu/Scripts/python.exe",)
+        if os.name == "nt"
+        else (".venvs/omni-cuda/bin/python", ".venvs/omni-cpu/bin/python")
     ),
-    ROUTE_TORCH_DML: (".venvs/dml/bin/python",),
-    ROUTE_CPU: (".venvs/omni-cuda/bin/python", ".venvs/omni-cpu/bin/python"),
     ROUTE_WIN_CUDA: (
-        r"/mnt/c/Users/zhout/w0/w0venv/Scripts/python.exe",
-        r"/mnt/c/Users/zhout/vllm-win/Scripts/python.exe",
+        str(_WINDOWS_HOME / "w0/w0venv/Scripts/python.exe"),
+        str(_WINDOWS_HOME / "vllm-win/Scripts/python.exe"),
     ),
 }
 
@@ -175,9 +181,10 @@ def is_wsl() -> bool:
     """
     if os.name == "nt":
         return False
-    return Path("/proc/sys/fs/binfmt_misc/WSLInterop").exists() or Path(
-        "/proc/sys/fs/binfmt_misc/WSLInterop-late"
-    ).exists()
+    return (
+        Path("/proc/sys/fs/binfmt_misc/WSLInterop").exists()
+        or Path("/proc/sys/fs/binfmt_misc/WSLInterop-late").exists()
+    )
 
 
 def to_worker_path(path: str | Path, *, is_windows: bool) -> str:
@@ -194,9 +201,7 @@ def to_worker_path(path: str | Path, *, is_windows: bool) -> str:
     if not shutil.which("wslpath"):
         return text
     try:
-        out = subprocess.run(
-            ["wslpath", "-w", text], capture_output=True, text=True, timeout=20, check=True
-        )
+        out = subprocess.run(["wslpath", "-w", text], capture_output=True, text=True, timeout=20, check=True)
     except (subprocess.SubprocessError, OSError):
         return text
     return out.stdout.strip() or text
@@ -206,16 +211,22 @@ def _resolve_interpreter(route: str, root: Path) -> tuple[str | None, str]:
     env_key = f"{ENV_PREFIX}{route.upper().replace('-', '_')}"
     override = os.environ.get(env_key)
     if override:
-        if Path(override).exists():
-            return override, ""
+        try:
+            if Path(override).is_file():
+                return override, ""
+        except OSError as exc:
+            return None, f"cannot access configured interpreter: {exc}"
         return None, f"{env_key}={override!r} does not exist"
 
     tried: list[str] = []
     for candidate in _DEFAULT_CANDIDATES.get(route, ()):
         path = Path(candidate) if Path(candidate).is_absolute() else root / candidate
         tried.append(str(path))
-        if path.exists():
-            return str(path), ""
+        try:
+            if path.is_file():
+                return str(path), ""
+        except OSError:
+            continue
     return None, f"no interpreter found; tried {tried} and {env_key} is unset"
 
 
@@ -226,9 +237,7 @@ def resolve(route: str, *, root: Path | None = None) -> Route:
     root = root or repo_root()
     is_windows = route in (ROUTE_VITISAI, ROUTE_DML, ROUTE_WIN_CUDA)
     worker_key = f"{WORKER_ENV_PREFIX}{route.upper().replace('-', '_')}"
-    worker = os.environ.get(worker_key) or str(
-        Path(__file__).resolve().parent / _WORKER_FOR_ROUTE[route]
-    )
+    worker = os.environ.get(worker_key) or str(Path(__file__).resolve().parent / _WORKER_FOR_ROUTE[route])
 
     if is_windows and not is_wsl() and os.name != "nt":
         return Route(

vllm_omni/edge/local/external/protocol.py

@@ -1,185 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""The wire between the engine and a stage that runs in another process.
+"""Compatibility import for the extracted host-copy wire implementation.
 
-Deliberately tiny, and deliberately **stdlib + numpy only**. The worker on the
-other end of this socket may be a native-Windows interpreter that has no torch,
-no vllm and no vllm_omni on its path -- ``C:\\Users\\zhout\\npu-ep`` is exactly
-that: Python 3.12, onnxruntime, onnx, numpy, nothing else. Anything imported
-here has to be installable there, so this module imports neither.
-
-**Why a socket at all.** Two devices on this laptop cannot be reached from the
-engine's own interpreter, and for different reasons:
-
-* the **XDNA2 NPU** is an MCDM device. WSL2's GPU-PV forwards WDDM display
-  adapters only, so there is no ``/dev/accel`` inside WSL and no amount of
-  installing changes that. VitisAI is a Windows DLL.
-* the **Radeon 890M** is reachable from WSL through DirectML, but
-  ``torch-directml`` pins torch 2.4.1 and Omni runs 2.13.
-
-So out-of-process is forced, not preferred. One protocol covers both, plus the
-native-Windows DirectML route, which needs its own venv anyway because
-``onnxruntime-directml`` is pinned at 1.24.4 and cannot load the VitisAI EP
-(that needs ORT >= 1.25) -- the two AMD devices cannot share an interpreter.
-
-**Framing.** ``[4-byte big-endian header length][JSON header][raw tensor bytes]``.
-Tensors are described in the header (name, dtype, shape, nbytes) and their
-buffers are concatenated after it in the order listed, so a receive is two
-reads and one ``frombuffer`` per tensor -- no per-element work and no pickle.
-Measured on this machine, WSL <-> native Windows: 11 us half round trip,
-330 MB/s. That is ample for a stage boundary (a 2 s Code2Wav chunk is ~192 KB,
-about 0.6 ms) and nowhere near enough for per-layer traffic inside a decode
-loop, which is the arithmetic behind "coarse stages only".
-
-``pickle`` is not used anywhere in here on purpose: the header is JSON, the
-payload is raw buffers, so a hostile or merely broken peer cannot execute code
-by replying.
+Standalone legacy worker scripts load this file directly. The source-tree
+fallback loads only the portable package, never the vllm_omni initializer.
 """
 
-from __future__ import annotations
+import sys
+from pathlib import Path
 
-import json
-import socket
-import struct
-from typing import Any
-
-import numpy as np
-
-# Ops the client sends.
-OP_LOAD = "load"
-"""Open a session on a graph. Replies with the placement report."""
-OP_RUN = "run"
-"""One forward. Inputs in, outputs out."""
-OP_STATS = "stats"
-"""Counters and RSS. The worker's memory is invisible to the engine's sampler
-when the worker is a Windows process, so it has to report its own."""
-OP_CLOSE = "close"
-
-# Ops the worker sends.
-OP_HELLO = "hello"
-"""First message after connect-back: interpreter, versions, available EPs."""
-OP_OK = "ok"
-OP_ERR = "err"
-
-MAX_HEADER_BYTES = 1 << 20
-"""A header is metadata; a megabyte of it is a bug or an attack, not a graph."""
-
-MAX_PAYLOAD_BYTES = 2 << 30
-"""Refuse rather than allocate on a corrupt length. Two gibibytes is far above
-any legitimate stage payload here -- the largest is a vision tower's patch
-tensor, single-digit MB."""
-
-
-class ProtocolError(RuntimeError):
-    """The peer sent something this protocol cannot represent."""
-
-
-class WorkerError(RuntimeError):
-    """The peer ran and failed. Carries the remote traceback when there is one."""
-
-    def __init__(self, message: str, *, remote_traceback: str = "", code: str = "") -> None:
-        # The remote traceback goes in the message, not only in an attribute:
-        # a bare "AssertionError:" surfacing here with the actual stack one
-        # attribute away is a debugging loop nobody should have to discover.
-        detail = f"{message}\n--- worker traceback ---\n{remote_traceback}" if remote_traceback else message
-        super().__init__(detail)
-        self.message = message
-        self.remote_traceback = remote_traceback
-        self.code = code
-
-
-def _recv_exactly(sock: socket.socket, n: int) -> bytes:
-    """Read exactly ``n`` bytes or raise.
-
-    ``recv`` is allowed to return short reads on a stream socket and does, once
-    payloads get past a page or two; a partial read here would be silently
-    misparsed as the next tensor.
-    """
-    if n == 0:
-        return b""
-    chunks: list[bytes] = []
-    remaining = n
-    while remaining:
-        chunk = sock.recv(min(remaining, 1 << 20))
-        if not chunk:
-            raise ProtocolError(f"peer closed with {remaining} of {n} bytes outstanding")
-        chunks.append(chunk)
-        remaining -= len(chunk)
-    return b"".join(chunks) if len(chunks) > 1 else chunks[0]
-
-
-def send_message(
-    sock: socket.socket,
-    op: str,
-    body: dict[str, Any] | None = None,
-    tensors: dict[str, np.ndarray] | None = None,
-) -> None:
-    """Send one message. ``tensors`` are sent as raw buffers after the header."""
-    tensors = tensors or {}
-    described: list[dict[str, Any]] = []
-    buffers: list[memoryview] = []
-    for name, array in tensors.items():
-        array = np.asarray(array)
-        # Record the shape *before* making it contiguous: ``ascontiguousarray``
-        # promotes a 0-d array to shape (1,), which would silently turn a
-        # scalar graph input into a length-1 vector on the far side.
-        shape = list(array.shape)
-        contiguous = np.ascontiguousarray(array)
-        described.append(
-            {
-                "name": name,
-                # ``str(dtype)`` loses endianness for multibyte types; ``.str``
-                # keeps the '<'/'>' so a big-endian peer would still decode.
-                "dtype": contiguous.dtype.str,
-                "shape": shape,
-                "nbytes": int(contiguous.nbytes),
-            }
-        )
-        # ``reshape(-1)`` rather than a direct cast: ``memoryview.cast`` refuses
-        # a 0-dimensional source.
-        buffers.append(memoryview(contiguous.reshape(-1)).cast("B"))
-
-    header = json.dumps({"op": op, "body": body or {}, "tensors": described}).encode("utf-8")
-    if len(header) > MAX_HEADER_BYTES:
-        raise ProtocolError(f"header is {len(header)} bytes, over the {MAX_HEADER_BYTES} limit")
-    sock.sendall(struct.pack(">I", len(header)))
-    sock.sendall(header)
-    for buffer in buffers:
-        sock.sendall(buffer)
-
-
-def recv_message(sock: socket.socket) -> tuple[str, dict[str, Any], dict[str, np.ndarray]]:
-    """Receive one message as ``(op, body, tensors)``.
-
-    Raises :class:`WorkerError` on an ``err`` reply so callers do not each have
-    to remember to check; every other op is returned as-is.
-    """
-    (header_len,) = struct.unpack(">I", _recv_exactly(sock, 4))
-    if header_len > MAX_HEADER_BYTES:
-        raise ProtocolError(f"peer announced a {header_len}-byte header, over the limit")
-    try:
-        header = json.loads(_recv_exactly(sock, header_len).decode("utf-8"))
-    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
-        raise ProtocolError(f"header is not JSON: {exc}") from exc
-
-    total = sum(int(t["nbytes"]) for t in header.get("tensors", ()))
-    if total > MAX_PAYLOAD_BYTES:
-        raise ProtocolError(f"peer announced a {total}-byte payload, over the limit")
-
-    tensors: dict[str, np.ndarray] = {}
-    for spec in header.get("tensors", ()):
-        raw = _recv_exactly(sock, int(spec["nbytes"]))
-        array = np.frombuffer(raw, dtype=np.dtype(spec["dtype"]))
-        # ``frombuffer`` is read-only and aliases ``raw``; a consumer that
-        # writes in place (ORT's IOBinding does) needs its own buffer.
-        tensors[spec["name"]] = array.reshape(spec["shape"]).copy()
-
-    op = str(header.get("op", ""))
-    body = header.get("body") or {}
-    if op == OP_ERR:
-        raise WorkerError(
-            str(body.get("message", "worker failed with no message")),
-            remote_traceback=str(body.get("traceback", "")),
-            code=str(body.get("code", "")),
-        )
-    return op, body, tensors
+try:
+    from omni_stage_contracts.wire import *  # noqa: F403
+except ModuleNotFoundError as exc:
+    if exc.name != "omni_stage_contracts":
+        raise
+    sys.path.insert(0, str(Path(__file__).resolve().parents[4]))
+    from omni_stage_contracts.wire import *  # noqa: F403

vllm_omni/edge/local/external/worker_ort.py

@@ -164,12 +164,37 @@ def _provider_counts(profile_path: str) -> dict[str, int]:
     return counts
 
 
+def _verify_artifact_members(graph: str, files: list[str]) -> None:
+    """Reject ONNX external data omitted from the controller's verified bundle."""
+    import onnx
+
+    allowed = {Path(name).resolve() for name in files}
+    graph_path = Path(graph).resolve()
+    if graph_path not in allowed:
+        raise ValueError("graph is not a verified artifact member")
+    pending = [onnx.load_model(graph, load_external_data=False)]
+    while pending:
+        message = pending.pop()
+        if message.DESCRIPTOR.full_name == "onnx.TensorProto":
+            locations = [entry.value for entry in message.external_data if entry.key == "location"]
+            if message.data_location == onnx.TensorProto.EXTERNAL and not locations:
+                raise ValueError("external tensor has no payload location")
+            for location in locations:
+                if (graph_path.parent / location).resolve() not in allowed:
+                    raise ValueError(f"external tensor payload is absent from verified manifest: {location}")
+        for descriptor, value in message.ListFields():
+            if descriptor.type == descriptor.TYPE_MESSAGE:
+                pending.extend(value if descriptor.is_repeated else [value])
+
+
 def _make_session(body: dict[str, Any]) -> tuple[Any, dict[str, Any]]:
     """Build the ORT session for the requested execution provider."""
     import onnxruntime as ort
 
     ep = str(body.get("ep", "cpu"))
     graph = str(body["graph_path"])
+    if "artifact_files" in body:
+        _verify_artifact_members(graph, body["artifact_files"])
     info: dict[str, Any] = {"ep": ep, "onnxruntime": ort.__version__}
 
     options = ort.SessionOptions()
@@ -241,12 +266,8 @@ def _load(state: _Session, body: dict[str, Any], tensors: dict[str, np.ndarray])
     state.session = session
     state.output_names = [o.name for o in session.get_outputs()]
 
-    info["inputs"] = [
-        {"name": i.name, "type": i.type, "shape": list(i.shape)} for i in session.get_inputs()
-    ]
-    info["outputs"] = [
-        {"name": o.name, "type": o.type, "shape": list(o.shape)} for o in session.get_outputs()
-    ]
+    info["inputs"] = [{"name": i.name, "type": i.type, "shape": list(i.shape)} for i in session.get_inputs()]
+    info["outputs"] = [{"name": o.name, "type": o.type, "shape": list(o.shape)} for o in session.get_outputs()]
 
     # The placement evidence. One profiled run, because ORT only writes node
     # assignments once nodes have actually executed -- a session that has never

vllm_omni/edge/local/session.py

@@ -33,10 +33,7 @@ nothing that would let a caller reach inside it.
 from __future__ import annotations
 
 import asyncio
-import time
 import uuid
-from dataclasses import asdict, dataclass, field
-from typing import Any
 
 STATE_LAYOUT_VERSION = 1
 """Bumped when the meaning of a handle's state changes. A handle from a
@@ -45,79 +42,10 @@ to treat state as non-migratable until a specific backend pair has been
 measured, and a version that silently matches would hide that."""
 
 
-@dataclass(frozen=True)
-class StateHandle:
-    """A reference to session state the backend owns.
-
-    The engine never reads the state. It reads this, to decide whether a
-    request may continue a session or must start a new one.
-    """
-
-    session_id: str
-    backend: str
-    artifact_id: str
-    """Ties the state to the checkpoint that produced it. The same shapes from
-    a different quantization are not the same state."""
-    layout_version: int = STATE_LAYOUT_VERSION
-    epoch: int = 0
-    """Incremented on every cancel. Events from a lower epoch are stale."""
-    replayable: bool = True
-    """A text session can be rebuilt by re-submitting its committed prompt, so
-    recovery after a backend crash is possible at an output boundary."""
-    migratable: bool = False
-    """Never true in M0. Moving KV between backends needs a measured layout
-    conversion for the specific pair; assuming it is the failure the proposal
-    warns about."""
-    created_unix: float = field(default_factory=time.time)
-
-    def next_epoch(self) -> StateHandle:
-        from dataclasses import replace
-
-        return replace(self, epoch=self.epoch + 1)
-
-    def accepts(self, other: StateHandle) -> bool:
-        """Whether ``other`` may continue this session's state."""
-        return (
-            other.session_id == self.session_id
-            and other.artifact_id == self.artifact_id
-            and other.layout_version == self.layout_version
-            and other.epoch == self.epoch
-        )
-
-    def to_dict(self) -> dict[str, Any]:
-        return asdict(self)
+from omni_stage_contracts.legacy import ChunkEvent, StateHandle  # noqa: E402, F401
 
 
-@dataclass(frozen=True)
-class ChunkEvent:
-    """One unit of output, with everything needed to order and retire it."""
-
-    request_id: str
-    stage_id: int
-    seq: int
-    """Monotonic within (request, epoch). Gaps mean loss; repeats mean a bug."""
-    epoch: int
-    kind: str
-    """``token`` | ``done`` | ``error`` | ``cancelled``."""
-    payload: Any
-    """For ``token``: the incremental text and token ids. Never the whole
-    output so far -- a cumulative payload makes a dropped chunk invisible."""
-    started_unix: float
-    emitted_unix: float
-    final: bool = False
-    error: str | None = None
-    input_watermark: int = 0
-    """How much of the input the producer had consumed when this was emitted.
-    One number in M0 (prompt tokens); it is what a duplex turn will align on."""
-    release_token: str = ""
-    """Returned by the consumer to give credit back. Empty when the event
-    carries no resource the producer is holding on its behalf."""
-
-    def to_dict(self) -> dict[str, Any]:
-        return asdict(self)
-
-
-class StreamClosed(RuntimeError):
+class StreamClosed(RuntimeError):  # noqa: N818 - retained public API
     """Raised on a stream that has been closed while a consumer waited."""
 
 
@@ -157,7 +85,12 @@ class BoundedEventStream:
             return len(payload.encode("utf-8"))
         if isinstance(payload, dict):
             return sum(
-                len(str(k)) + (len(v.encode("utf-8")) if isinstance(v, str) else 8 * (len(v) if isinstance(v, (list, tuple)) else 1))
+                len(str(k))
+                + (
+                    len(v.encode("utf-8"))
+                    if isinstance(v, str)
+                    else 8 * (len(v) if isinstance(v, (list, tuple)) else 1)
+                )
                 for k, v in payload.items()
             )
         return 64
@@ -172,11 +105,13 @@ class BoundedEventStream:
         size = self._sizeof(event)
         async with self._credit:
             await self._credit.wait_for(
-                lambda: self._closed
-                or (self._chunks < self.max_chunks and self._bytes + size <= self.max_bytes)
-                # A single event larger than the whole byte bound must still go
-                # through once the queue is empty, or it deadlocks forever.
-                or (self._chunks == 0 and size > self.max_bytes)
+                lambda: (
+                    self._closed
+                    or (self._chunks < self.max_chunks and self._bytes + size <= self.max_bytes)
+                    # A single event larger than the whole byte bound must still go
+                    # through once the queue is empty, or it deadlocks forever.
+                    or (self._chunks == 0 and size > self.max_bytes)
+                )
             )
             if self._closed:
                 raise StreamClosed("stream closed while producing")

vllm_omni/edge/topology.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Map discovery facts into physical constraints, without guessing performance."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import asdict
+
+from omni_stage_contracts import DeviceDescriptor
+
+
+def describe_topology(
+    devices,
+    *,
+    machine_id: str,
+    domain_id: str,
+    physical_ids: Mapping[str, str] | None = None,
+    guest_pool: str | None = None,
+) -> list[DeviceDescriptor]:
+    """Normalize legacy discovery; caller reconciles cross-domain identities.
+
+    Physical IDs must be shared across routes when discovery proves they refer
+    to one GPU. Unknown bandwidth/power relationships stay empty. WSL quota is
+    an additional constraint on host allocations, never another RAM capacity.
+    """
+    if not machine_id or not domain_id:
+        raise ValueError("machine and domain IDs must be explicit")
+    physical_ids = physical_ids or {}
+    result = []
+    for device in devices:
+        extra = device.extra or {}
+        physical = physical_ids.get(device.device_id, device.device_id)
+        kind = {"gpu_integrated": "gpu", "gpu_discrete": "gpu"}.get(device.kind, device.kind)
+        integrated = device.kind == "gpu_integrated"
+        integration = "integrated" if integrated else "discrete" if device.kind == "gpu_discrete" else "unknown"
+        pool = f"{machine_id}:ram" if device.memory_pool == "host_ram" else f"{machine_id}:vram:{physical}"
+        pools = (pool,) + ((guest_pool,) if guest_pool and device.memory_pool == "host_ram" else ())
+        result.append(
+            DeviceDescriptor(
+                device_id=f"{machine_id}:{physical}",
+                kind=kind,
+                domain_id=domain_id,
+                memory_pool_ids=pools,
+                integration=extra.get("integration", integration),
+                package_id=extra.get("package_id"),
+                bandwidth_group_ids=tuple(extra.get("bandwidth_group_ids", ())),
+                power_domain_ids=tuple(extra.get("power_domain_ids", ())),
+            )
+        )
+    return result
+
+
+def topology_report(devices, **kwargs) -> list[dict]:
+    return [asdict(device) for device in describe_topology(devices, **kwargs)]

vllm_omni/engine/async_omni_engine.py

@@ -244,6 +244,11 @@ class AsyncOmniEngine:
             )
 
         self.num_stages = len(self.stage_configs)
+        from vllm_omni.engine.resource_ledger import GraphRequestGate
+
+        self._graph_request_gate = (
+            GraphRequestGate() if any(c.stage_type == "graph" for c in self.stage_configs) else None
+        )
         stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None
         self.async_chunk = bool(getattr(stage0_args, "async_chunk", False))
         self.stage_pools: list[StagePool] = []
@@ -454,6 +459,7 @@ class AsyncOmniEngine:
                 enable_duplex_control=self._duplex_control_enabled,
                 duplex_session_config=self.duplex_session_config,
             )
+            orchestrator.graph_request_gate = self._graph_request_gate
             if not startup_future.done():
                 startup_future.set_result(asyncio.get_running_loop())
             await orchestrator.run()
@@ -757,7 +763,7 @@ class AsyncOmniEngine:
         stage_type = self.stage_metadata[0].stage_type
         output_prompt_text: Any = None
         _preprocess_ms = 0.0
-        if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
+        if stage_type not in ("diffusion", "graph") and not isinstance(prompt, EngineCoreRequest):
             # Stage transforms and downstream stages must share the same
             # request identity, including when the transform replaces the
             # prompt object.
@@ -1145,6 +1151,8 @@ class AsyncOmniEngine:
         a queue + coroutine-switch round-trip.  The Orchestrator receives a
         ready-to-submit OmniEngineCoreRequest.
         """
+        graph_gate = getattr(self, "_graph_request_gate", None)
+        graph_ticket = graph_gate.acquire(request_id) if graph_gate is not None else None
         try:
             msg = self._build_add_request_message(
                 request_id=request_id,
@@ -1163,6 +1171,8 @@ class AsyncOmniEngine:
                 resumable=resumable,
             )
         except BaseException:
+            if graph_gate is not None:
+                graph_gate.release(graph_ticket)
             if isinstance(prompt, dict):
                 for artifact_dir in prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) or ():
                     if isinstance(artifact_dir, str):
@@ -1186,6 +1196,8 @@ class AsyncOmniEngine:
 
             self.request_queue.sync_q.put(msg)
         except BaseException:
+            if graph_gate is not None:
+                graph_gate.release(graph_ticket)
             for artifact_dir in msg.request_artifact_dirs or ():
                 shutil.rmtree(artifact_dir, ignore_errors=True)
             raise
@@ -1554,7 +1566,7 @@ class AsyncOmniEngine:
     def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None:
         """Read one output message from the Orchestrator output queue."""
         try:
-            return self.output_queue.sync_q.get(timeout=timeout)
+            return self._attach_graph_release(self.output_queue.sync_q.get(timeout=timeout))
         except queue.Empty:
             if not self.is_alive():
                 raise RuntimeError("Orchestrator died unexpectedly. See logs above.")
@@ -1563,7 +1575,7 @@ class AsyncOmniEngine:
     async def try_get_output_async(self) -> EngineQueueMessage | None:
         """Async read from the Orchestrator output queue."""
         try:
-            return self.output_queue.sync_q.get_nowait()
+            return self._attach_graph_release(self.output_queue.sync_q.get_nowait())
         except queue.Empty:
             if not self.is_alive():
                 raise RuntimeError("Orchestrator died unexpectedly. See logs above.")
@@ -1603,6 +1615,24 @@ class AsyncOmniEngine:
         msg = await loop.run_in_executor(executor, _drain_get)
         if msg is None and not self.is_alive():
             raise RuntimeError("Orchestrator died unexpectedly. See logs above.")
+        return self._attach_graph_release(msg)
+
+    def _attach_graph_release(self, msg):
+        from vllm_omni.outputs import OmniRequestOutput
+
+        gate = getattr(self, "_graph_request_gate", None)
+        if gate is not None and isinstance(msg, OutputMessage) and msg.finished:
+            ticket = gate.current(msg.request_id)
+            output = msg.engine_outputs
+            if isinstance(output, OmniRequestOutput):
+                previous = output._stage_release
+
+                def release():
+                    if previous is not None:
+                        previous()
+                    gate.release(ticket)
+
+                output._stage_release = release
         return msg
 
     def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
@@ -1673,6 +1703,10 @@ class AsyncOmniEngine:
             raise
         if not result_msg.success:
             raise RuntimeError(result_msg.error or "abort failed")
+        graph_gate = getattr(self, "_graph_request_gate", None)
+        if graph_gate is not None:
+            for request_id in request_ids:
+                graph_gate.release(graph_gate.current(request_id))
         return list(result_msg.abort_outputs or [])
 
     def submit_interaction(

vllm_omni/engine/backends/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Factories for non-vLLM stages owned by the existing StageRuntime."""
+
+
+def create_graph_client(metadata, config, ledger, reservation):
+    from .graph import GraphStageClient
+
+    if config.get("backend", {}).get("name") != "external.graph.v1":
+        raise ValueError("graph stage requires backend.name=external.graph.v1")
+    return GraphStageClient(metadata, config["backend"], ledger, reservation)

vllm_omni/engine/backends/graph.py

@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+"""One bounded, non-preemptible graph stage using the existing worker.
+
+The runtime owns startup and capacity. This adapter owns one in-flight call,
+including its output until acknowledged by the consumer. No token scheduler,
+provider fallback, or automatic worker replay is introduced here.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import math
+import uuid
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from omni_stage_contracts import ArtifactManifest, BufferRef, StageEvent, StageRequest
+from vllm_omni.edge.local.external.client import ExternalWorker
+from vllm_omni.engine.resource_ledger import ResourceUnavailable
+from vllm_omni.engine.stage_client import StageClientBase
+from vllm_omni.host.routes import resolve_route
+from vllm_omni.outputs import OmniRequestOutput
+
+
+class GraphStageClient(StageClientBase):
+    def __init__(self, metadata, config, ledger, reservation) -> None:
+        for name, value in vars(metadata).items():
+            setattr(self, name, value)
+        self.stage_type = "graph"
+        self._ledger, self._reservation = ledger, reservation
+        self._generation = uuid.uuid4().hex
+        self._config = dict(config)
+        self._worker = None
+        self._closed = False
+        self._active: str | None = None
+        self._task: asyncio.Task | None = None
+        self._output: OmniRequestOutput | None = None
+        self._epoch = 0
+        self._max_io_bytes = int(config.get("max_io_bytes", 8 << 20))
+        if self._max_io_bytes <= 0 or self._max_io_bytes > min(reservation.demands.values()):
+            raise ValueError("graph I/O limit must be positive and fit the stage reservation")
+        try:
+            manifest_path = Path(config["manifest"]).resolve()
+            manifest = ArtifactManifest.read(manifest_path)
+            graph_name = manifest.metadata["graph_file"]
+            examples_name = manifest.metadata["example_inputs_file"]
+            if graph_name not in manifest.files or examples_name not in manifest.files:
+                raise ValueError("graph and placement examples must be hashed artifact payloads")
+            if sum((manifest_path.parent / n).stat().st_size for n in manifest.files) > min(
+                reservation.demands.values()
+            ):
+                raise ResourceUnavailable("artifact files exceed stage memory reservation")
+            route = resolve_route(config["route"])
+            self._worker = ExternalWorker(
+                route, start_timeout_s=float(config.get("start_timeout_s", 120)), max_payload_bytes=self._max_io_bytes
+            )
+            self._worker.start()
+            with np.load(manifest_path.parent / examples_name, allow_pickle=False) as example:
+                inputs = {name: example[name] for name in example.files}
+            self._describe(inputs)
+            self.load_report = self._worker.load(
+                manifest_path.parent / graph_name,
+                example_inputs=inputs,
+                ep_dir=config.get("ep_dir"),
+                intra_op_num_threads=config.get("intra_op_num_threads"),
+                profile_prefix=config.get("profile_prefix"),
+                artifact_files=[manifest_path.parent / name for name in manifest.files],
+            )
+            report = self.load_report
+            minimum = float(config.get("min_fraction_on_target", 1.0))
+            if not 0 < minimum <= 1:
+                raise ValueError("placement threshold must be in (0, 1]")
+            if (
+                report.requested_provider_missing
+                or report.fraction_on_target is None
+                or not math.isfinite(report.fraction_on_target)
+                or not 0 < report.fraction_on_target <= 1
+                or report.target_nodes <= 0
+                or report.fraction_on_target < minimum
+            ):
+                raise RuntimeError(f"REFUSE_EP_PLACEMENT: {report.summary()}")
+            allowed = set(config.get("allowed_providers", report.session_providers if route.ep == "cpu" else ()))
+            if not allowed or set(report.node_counts) - allowed:
+                raise RuntimeError("REFUSE_EP_PLACEMENT: declare all permitted execution providers")
+            if report.rss_bytes > min(reservation.demands.values()):
+                raise ResourceUnavailable("loaded worker RSS exceeds stage memory reservation")
+            self._input_schema = {name: (value.dtype.str, value.shape) for name, value in inputs.items()}
+            from omni_stage_contracts.types import file_digest
+
+            self.execution_plan = {
+                "backend": "external.graph.v1",
+                "stage_id": self.stage_id,
+                "worker_generation": self._generation,
+                "artifact_sha256": file_digest(manifest_path),
+                "component": manifest.component,
+                "route": route.to_dict(),
+                "reserved_bytes": dict(reservation.demands),
+                "max_io_bytes": self._max_io_bytes,
+                "placement": report.to_dict(),
+                "worker": {k: v for k, v in self._worker.hello.items() if k != "token"},
+                "evidence": "B",
+                "stateful": False,
+            }
+        except BaseException:
+            self.shutdown()
+            raise
+
+    def _describe(self, tensors: dict[str, np.ndarray]) -> tuple[BufferRef, ...]:
+        refs = tuple(
+            BufferRef(
+                name, str(self.stage_id), self._generation, array.dtype.name, tuple(array.shape), int(array.nbytes)
+            )
+            for name, array in tensors.items()
+        )
+        if sum(r.nbytes for r in refs) > self._max_io_bytes:
+            raise ResourceUnavailable("graph payload exceeds admitted I/O bound")
+        return refs
+
+    async def add_request_async(self, request_id: str, prompt: Any, params: Any = None) -> None:
+        self.check_health()
+        if self._active is not None:
+            raise ResourceUnavailable("graph stage has an unacknowledged request; capacity is one")
+        if not isinstance(prompt, dict) or not isinstance(prompt.get("tensors"), dict):
+            raise ValueError("graph prompt must contain a tensors mapping supplied by its model adapter")
+        inputs = {name: np.asarray(value) for name, value in prompt["tensors"].items()}
+        refs = self._describe(inputs)
+        if {name: (value.dtype.str, value.shape) for name, value in inputs.items()} != self._input_schema:
+            raise ValueError("input names/dtypes/shapes differ from the placement-verified artifact bucket")
+        # Copy at submission: application mutation cannot change an admitted call.
+        inputs = {name: value.copy() for name, value in inputs.items()}
+        self._epoch += 1
+        epoch = self._epoch
+        request = StageRequest(request_id, self.stage_id, epoch, self._generation, inputs=refs)
+        self._active = request_id
+        self._task = asyncio.create_task(self._run(request, inputs), name=f"graph-{self.stage_id}-{request_id}")
+
+    async def _run(self, request: StageRequest, inputs: dict[str, np.ndarray]) -> None:
+        try:
+            outputs, timing = await asyncio.to_thread(self._worker.run, inputs)
+            refs = self._describe(outputs)
+            if sum(x.nbytes for x in outputs.values()) + sum(x.nbytes for x in inputs.values()) > self._max_io_bytes:
+                raise ResourceUnavailable("combined graph input/output exceeds admitted I/O bound")
+            event = StageEvent(
+                request.request_id, self.stage_id, request.epoch, 1, "tensor", self._generation, refs, terminal=True
+            )
+            output = OmniRequestOutput(
+                request_id=request.request_id,
+                stage_id=self.stage_id,
+                final_output_type=self.final_output_type or "latent",
+                _custom_output={"tensors": outputs, "stage_event": dataclasses.asdict(event)},
+                metrics={"graph_timing": timing.to_dict()},
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            output = OmniRequestOutput.from_error(request.request_id, str(exc))
+            output.stage_id = self.stage_id
+        if not self._closed and self._epoch == request.epoch:
+            loop = asyncio.get_running_loop()
+            output._stage_release = lambda: loop.call_soon_threadsafe(
+                self.acknowledge, request.request_id, request.epoch, request.worker_generation
+            )
+            self._output = output
+
+    def get_graph_output_nowait(self):
+        # Delivery is separate from release: the slot stays charged until ACK.
+        output, self._output = self._output, None
+        return output
+
+    def acknowledge(self, request_id: str, epoch: int, generation: str) -> None:
+        if (
+            self._active == request_id
+            and epoch == self._epoch
+            and generation == self._generation
+            and (self._task is None or self._task.done())
+        ):
+            self._active = None
+            self._task = None
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        if self._active not in request_ids:
+            return
+        self._epoch += 1
+        self._output = None
+        # A completed call needs no worker kill. Otherwise cancellation retires
+        # this single-session worker; the runtime will report it unavailable.
+        if self._task is not None and not self._task.done():
+            drained = await asyncio.to_thread(self._worker.terminate)
+            self._closed = True
+            done, _ = await asyncio.wait({self._task}, timeout=5)
+            self._ledger.release(self._reservation, drained=drained and bool(done) and not self._task.cancelled())
+        self._active = None
+
+    def check_health(self) -> None:
+        if self._closed:
+            from vllm.v1.engine.exceptions import EngineDeadError
+
+            raise EngineDeadError()
+
+    async def collective_rpc_async(self, method, timeout=None, args=(), kwargs=None):
+        raise NotImplementedError(f"graph backend does not implement collective RPC {method}")
+
+    def shutdown(self) -> None:
+        self._closed = True
+        self._output = None
+        drained = self._worker is None or self._worker.terminate()
+        task = self._task
+        if task is not None and not task.done():
+            # Process exit does not release Python input/transport references.
+            # Keep them charged until the call coroutine actually unwinds.
+            self._ledger.release(self._reservation, drained=False)
+            task.add_done_callback(
+                lambda done: self._ledger.release(self._reservation, drained=drained and not done.cancelled())
+            )
+        else:
+            self._ledger.release(self._reservation, drained=drained)

vllm_omni/engine/orchestrator.py

@@ -1191,8 +1191,12 @@ class Orchestrator:
                     # Shared catch so a dead replica on either poll path is
                     # evicted rather than tearing down every stage (#4285).
                     try:
-                        if pool.stage_type == "diffusion":
-                            diffusion_output = pool.poll_diffusion_output(replica_id)
+                        if pool.stage_type in ("diffusion", "graph"):
+                            diffusion_output = (
+                                pool.poll_graph_output(replica_id)
+                                if pool.stage_type == "graph"
+                                else pool.poll_diffusion_output(replica_id)
+                            )
                             if diffusion_output is None:
                                 continue
 
@@ -1282,7 +1286,11 @@ class Orchestrator:
                     got = False
                     for replica_id in pool.available_replica_ids():
                         try:
-                            output = pool.poll_diffusion_output(replica_id)
+                            output = (
+                                pool.poll_graph_output(replica_id)
+                                if pool.stage_type == "graph"
+                                else pool.poll_diffusion_output(replica_id)
+                            )
                         except EngineDeadError as e:
                             await ready_q.put(("error", stage_id, replica_id, e))
                             continue
@@ -1317,7 +1325,7 @@ class Orchestrator:
                     # No live replica yet; nothing to attach to. A later
                     # reconcile picks the stage up once one registers.
                     continue
-                if stage_type == "diffusion":
+                if stage_type in ("diffusion", "graph"):
                     live_pools.add(stage_id)
                     existing_poller = pollers.get(stage_id)
                     if existing_poller is None or existing_poller.done():
@@ -1500,6 +1508,8 @@ class Orchestrator:
         for output in outputs:
             req_state = self.request_states.get(output.request_id)
             if req_state is None:
+                if hasattr(output, "release_stage_buffers"):
+                    output.release_stage_buffers()
                 logger.warning(
                     "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)",
                     output.request_id,
@@ -1524,7 +1534,11 @@ class Orchestrator:
                 )
                 stage_metrics.pipeline_timings = dict(req_state.pipeline_timings)
 
-            await self._route_output(stage_id, replica_id, output, req_state, stage_metrics)
+            try:
+                await self._route_output(stage_id, replica_id, output, req_state, stage_metrics)
+            finally:
+                if pool.stage_type == "graph" and stage_id not in req_state.final_output_stage_ids:
+                    output.release_stage_buffers()
 
     async def _handle_stage_error(self, stage_id: int, output: Any) -> None:
         """Emit a frontend-visible error and clean up request state."""
@@ -1834,6 +1848,10 @@ class Orchestrator:
         try:
             if abort:
                 abort_outputs = await self._abort_request_ids(cleanup_ids)
+                graph_gate = getattr(self, "graph_request_gate", None)
+                if graph_gate is not None:
+                    for request_id in cleanup_ids:
+                        graph_gate.release(graph_gate.current(request_id))
             self._release_request_bindings(cleanup_ids)
             for request_id in cleanup_ids:
                 self._pd_kv_params.pop(request_id, None)
@@ -2513,6 +2531,17 @@ class Orchestrator:
         requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
         _t_submit_start = _time.perf_counter()
 
+        if next_pool.stage_type == "graph":
+            if is_streaming_session:
+                raise ValueError("external graph requires a complete upstream payload")
+            transform = next_client.custom_process_input_func
+            if transform is None:
+                raise ValueError("graph dependency requires a model-owned custom_process_input_func")
+            graph_prompt = transform(source_outputs, req_state.prompt, requires_multimodal_data)
+            await next_pool.submit_initial(req_id, req_state, graph_prompt)
+            req_state.stage_submit_ts[next_logical] = _time.time()
+            return
+
         if next_pool.stage_type == "diffusion":
             # Gate: never dispatch with an incomplete CFG bundle. Checked
             # non-destructively BEFORE popping — a pop-then-redefer would lose

vllm_omni/engine/resource_ledger.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Atomic reservations shared by native and external stage startup.
+
+Capacities are admission ceilings for this controller, captured before loads.
+They are not fresh `available` samples: resident allocations remain charged,
+so sampling available memory cannot subtract our own allocation a second time.
+Refresh ceilings only with an explicit, reconciled controller budget.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+
+
+class ResourceUnavailable(RuntimeError):  # noqa: N818 - public admission result, like PlanNotAdmitted
+    pass
+
+
+@dataclass(frozen=True)
+class Reservation:
+    owner: str
+    demands: Mapping[str, int]
+
+
+class ResourceLedger:
+    def __init__(self, capacities: Mapping[str, int]) -> None:
+        if not capacities or any(type(n) is not int or n < 0 for n in capacities.values()):
+            raise ValueError("explicit nonnegative memory ceilings are required")
+        self.capacities = MappingProxyType(dict(capacities))
+        self._active: dict[str, Reservation] = {}
+        self._quarantined: set[str] = set()
+        self._lock = threading.RLock()
+
+    def reserve(self, owner: str, demands: Mapping[str, int]) -> Reservation:
+        with self._lock:
+            if not owner or owner in self._active:
+                raise ValueError(f"reservation owner missing or already active: {owner}")
+            if not demands or any(type(n) is not int or n < 0 for n in demands.values()):
+                raise ValueError("nonnegative memory demands are required")
+            for pool, count in demands.items():
+                if pool not in self.capacities:
+                    raise ResourceUnavailable(f"unknown memory pool: {pool}")
+                used = sum(r.demands.get(pool, 0) for r in self._active.values())
+                if used + count > self.capacities[pool]:
+                    raise ResourceUnavailable(f"{pool}: {used} reserved + {count} requested > {self.capacities[pool]}")
+            reservation = Reservation(owner, MappingProxyType(dict(demands)))
+            self._active[owner] = reservation
+            return reservation
+
+    def release(self, reservation: Reservation, *, drained: bool) -> bool:
+        with self._lock:
+            if self._active.get(reservation.owner) is not reservation:
+                return False  # stale token must never release a newer allocation
+            if not drained:
+                self._quarantined.add(reservation.owner)
+                return False
+            del self._active[reservation.owner]
+            self._quarantined.discard(reservation.owner)
+            return True
+
+    def snapshot(self) -> dict:
+        with self._lock:
+            return {
+                "capacities": dict(self.capacities),
+                "reserved": {p: sum(r.demands.get(p, 0) for r in self._active.values()) for p in self.capacities},
+                "owners": list(self._active),
+                "quarantined": sorted(self._quarantined),
+            }
+
+
+class GraphRequestGate:
+    """One admitted pipeline request, held through the final consumer ACK.
+
+    Stage reservations already include its payload allowance. Rejecting excess
+    work before the ingress queue keeps slow consumers from accumulating inputs
+    or error outputs, and leaves the existing control queue available to abort.
+    """
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._active = None
+
+    def acquire(self, request_id):
+        with self._lock:
+            if self._active is not None:
+                raise ResourceUnavailable("graph pipeline capacity is one; consume or cancel the active request first")
+            ticket = (request_id, object())
+            self._active = ticket
+            return ticket
+
+    def current(self, request_id):
+        with self._lock:
+            return self._active if self._active is not None and self._active[0] == request_id else None
+
+    def release(self, ticket):
+        with self._lock:
+            if ticket is not None and self._active is ticket:
+                self._active = None

vllm_omni/engine/stage_client.py

@@ -71,6 +71,14 @@ class StagePoolClient(StageClient, Protocol):
     ) -> Any: ...
 
 
+class StagePoolGraphClient(StagePoolClient, Protocol):
+    async def add_request_async(self, request_id: str, prompt: Any, params: Any = None) -> None: ...
+
+    def get_graph_output_nowait(self) -> Any | None: ...
+
+    def acknowledge(self, request_id: str, epoch: int, generation: str) -> None: ...
+
+
 class StagePoolLLMClient(StagePoolClient, Protocol):
     """Pool-facing API for LLM-style stages."""
 

vllm_omni/engine/stage_init_utils.py

@@ -12,12 +12,6 @@ out of StageEngineCoreClient into reusable functions.
 from __future__ import annotations
 
 import copy
-# [edge-infer W1] fcntl does not exist on Windows; this module is in the
-# import closure of AsyncOmni, so the bare import makes the engine
-# unimportable before anything can be configured. The locking logic is
-# sound and torch-free -- only the primitive changes. The shim keeps the
-# BlockingIOError contract this file's stale-lock path depends on.
-from vllm_omni._filelock_compat import flock_exclusive_nb, funlock
 import importlib
 import json
 import multiprocessing as mp
@@ -41,6 +35,12 @@ from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine.input_processor import InputProcessor
 from vllm.v1.executor import Executor
 
+# [edge-infer W1] fcntl does not exist on Windows; this module is in the
+# import closure of AsyncOmni, so the bare import makes the engine
+# unimportable before anything can be configured. The locking logic is
+# sound and torch-free -- only the primitive changes. The shim keeps the
+# BlockingIOError contract this file's stale-lock path depends on.
+from vllm_omni._filelock_compat import flock_exclusive_nb, funlock
 from vllm_omni.config.omni_config import (
     _CACHE_STAGE_ENGINE_FIELD_MAP,
     _DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP,
@@ -549,7 +549,7 @@ class StageMetadata:
     """Lightweight stage attributes extracted from stage_config."""
 
     stage_id: int
-    stage_type: Literal["llm", "diffusion"]
+    stage_type: Literal["llm", "diffusion", "graph"]
     engine_output_type: str | None
     is_comprehension: bool
     requires_multimodal_data: bool
@@ -599,7 +599,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
     engine-argument and stage-init consumers together.
     """
     stage_id: int = stage_config.stage_id
-    stage_type: Literal["llm", "diffusion"] = _get_attr_or_item(stage_config, "stage_type", "llm")
+    stage_type: Literal["llm", "diffusion", "graph"] = _get_attr_or_item(stage_config, "stage_type", "llm")
     engine_args = stage_config.engine_args
 
     _apply_rocm_attention_backend(engine_args, stage_type)
@@ -711,8 +711,11 @@ def extract_stage_metadata_from_omni_stage_config(
     layout, engine-argument, remote-diffusion, and platform setup paths still
     require the legacy StageConfig/OmegaConf shape.
     """
-    stage_type: Literal["llm", "diffusion"] = "diffusion" if stage_config.stage_type == StageType.DIFFUSION else "llm"
-    sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
+    stage_type: Literal["llm", "diffusion", "graph"] = {
+        StageType.DIFFUSION: "diffusion",
+        StageType.GRAPH: "graph",
+    }.get(stage_config.stage_type, "llm")
+    sampling_params_cls = OmniDiffusionSamplingParams if stage_type == "diffusion" else SamplingParams
     sampling_params: OmniSamplingParams = sampling_params_cls(
         **(stage_config.model_config.default_sampling_params or {})
     )
@@ -738,7 +741,7 @@ def extract_stage_metadata_from_omni_stage_config(
 
     return StageMetadata(
         stage_id=stage_config.stage_id,
-        stage_type="llm",
+        stage_type=stage_type,
         engine_output_type=stage_config.engine_output_type,
         is_comprehension=stage_config.is_comprehension,
         requires_multimodal_data=stage_config.requires_multimodal_data,
@@ -1574,7 +1577,7 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
     return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer)
 
 
-def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str:
+def device_init_lock_path(device_id: int, lock_dir: str | None = None) -> str:
     """Return the per-physical-device initialization lock file path.
 
     Shared by the orchestrator-side ``acquire_device_locks`` (legacy full-init
@@ -1585,7 +1588,9 @@ def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str:
     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")
+    from vllm_omni.host.paths import device_lock_directory
+
+    return os.path.join(lock_dir or device_lock_directory(), f"vllm_omni_device_{device_id}_init.lock")
 
 
 def _open_existing_lock_file(lock_file: str) -> tuple[int, bool] | None:
@@ -1637,6 +1642,12 @@ def open_device_lock_file(lock_file: str) -> tuple[int, bool]:
     very failure this avoids. ``link`` also fails cleanly if another process wins
     the race, which keeps the "first creator wins" inode stable.
     """
+    if os.name == "nt":
+        # Windows byte-range locks require a writable handle. O_CREAT without
+        # O_TRUNC atomically opens/creates the stable file; there is no fchmod
+        # or POSIX cross-user mode publication on this host.
+        return os.open(lock_file, os.O_RDWR | os.O_CREAT, 0o600), True
+
     existing = _open_existing_lock_file(lock_file)
     if existing is not None:
         return existing

vllm_omni/engine/stage_phase_lock.py

@@ -76,6 +76,7 @@ def _unlock(fd: int) -> None:
 
     funlock(fd)
 
+
 logger = init_logger(__name__)
 
 # Default upper bound on how long a single phase will wait for its device locks
@@ -144,7 +145,7 @@ class DevicePhaseLock:
         device_ids: list[int],
         *,
         timeout_s: float = _DEFAULT_LOCK_TIMEOUT_S,
-        lock_dir: str = "/tmp",
+        lock_dir: str | None = None,
     ) -> None:
         self._device_ids = sorted(device_ids)
         self._timeout_s = timeout_s

vllm_omni/engine/stage_pool.py

@@ -990,6 +990,14 @@ class StagePool:
         if self.stage_type == "diffusion":
             params = OmniDiffusionSamplingParams.from_params(params)
         submit_kwargs = dict(submit_kwargs or {})
+        if self.stage_type == "graph":
+            replica_id = await self._pick_or_select(request_id, affinity_request_id=affinity_request_id)
+            try:
+                await self.clients[replica_id].add_request_async(request_id, request, params)
+            except BaseException:
+                self.release_binding(request_id)
+                raise
+            return replica_id
         if self.stage_type == "diffusion":
             if isinstance(request, list):
                 raise ValueError(
@@ -1050,6 +1058,8 @@ class StagePool:
         prompt_text: Any = None,
     ) -> int:
         """Submit a streaming update to an already admitted request."""
+        if self.stage_type == "graph":
+            raise ValueError("external.graph.v1 accepts complete graphs, not streaming updates")
         params = req_state.sampling_params_list[self.stage_id]
         if self.stage_type == "diffusion":
             params = OmniDiffusionSamplingParams.from_params(params)
@@ -1207,6 +1217,12 @@ class StagePool:
             return None
         return cast(StagePoolDiffusionClient, raw_client).get_diffusion_output_nowait()
 
+    def poll_graph_output(self, replica_id: int) -> Any | None:
+        if not self.is_replica_available(replica_id):
+            return None
+        client = self.clients[replica_id]
+        return None if client is None else client.get_graph_output_nowait()
+
     # ---- Stage-local control plane ----
 
     async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]:

vllm_omni/engine/stage_runtime.py

@@ -155,6 +155,9 @@ class StageRuntime:
         # device groups can initialize concurrently.
         self._replica_launch_lock = threading.Lock()
         self._init_visible_devices_baseline: str | None = None
+        self.resource_ledger = None
+        self._resource_reservations = {}
+        self._resource_started = set()
 
     @staticmethod
     def _client_addresses_from_zmq(addresses: Any) -> dict[str, str]:
@@ -221,18 +224,21 @@ class StageRuntime:
         return collected
 
     @staticmethod
-    def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None:
+    def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> set[int]:
         """Best-effort shutdown for attached clients after init failure."""
+        drained = set()
         for client in reversed(list(clients)):
             if client is None:
                 continue
             try:
                 client.shutdown()
+                drained.add(id(client))
             except Exception as cleanup_error:
                 logger.warning(
                     "[StageRuntime] Failed to shutdown initialized client after init failure: %s",
                     cleanup_error,
                 )
+        return drained
 
     def initialize(self) -> None:
         """Run the full stage initialization sequence."""
@@ -241,6 +247,7 @@ class StageRuntime:
             plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans
         }
         try:
+            self._reserve_stage_resources(stage_plans)
             self._before_initialize_stage_replicas(stage_plans)
             initialized_clients = self._initialize_stage_replicas(stage_plans, self._stage_init_timeout)
             initialized_clients_by_stage = initialized_clients
@@ -259,16 +266,80 @@ class StageRuntime:
                 "[StageRuntime] Stage initialization failed; shutting down %s initialized client(s)",
                 len(cleanup_clients),
             )
-            self._shutdown_initialized_clients(cleanup_clients)
+            drained = self._shutdown_initialized_clients(cleanup_clients) or set()
+            if self.resource_ledger is not None:
+                for client in cleanup_clients:
+                    if client.stage_type != "graph":
+                        reservation = self._resource_reservations.get((client.stage_id, client.replica_id))
+                        if reservation is not None:
+                            self.resource_ledger.release(reservation, drained=id(client) in drained)
             self._cleanup_after_initialize_failure()
+            self._release_unstarted_resources(initialized_clients_by_stage)
             raise exc
 
+    def _reserve_stage_resources(self, stage_plans) -> None:
+        """Explicit controller ceilings are mandatory for mixed graph plans.
+
+        Native stage budgets join the same transaction; absent estimates fail
+        closed rather than treating native allocations as free memory.
+        """
+        from vllm_omni.engine.resource_ledger import ResourceLedger
+
+        replicas = [r for plan in stage_plans for r in plan.replicas]
+        if not any(r.metadata.stage_type == "graph" for r in replicas):
+            return
+        # One final consumer owns the ingress ticket. Fan-out or multiple final
+        # outputs would need reference-counted leases before any early ACK can
+        # release that ticket. Refuse those topologies in the v1 graph backend.
+        previous_stage = None
+        for index, plan in enumerate(stage_plans):
+            metadata = plan.replicas[0].metadata
+            expected_sources = [] if previous_stage is None else [previous_stage]
+            if (
+                list(metadata.engine_input_source or []) != expected_sources
+                or metadata.final_output != (index == len(stage_plans) - 1)
+            ):
+                raise ValueError("graph v1 requires a linear chain with exactly one final output")
+            previous_stage = metadata.stage_id
+        budgets = [r.stage_cfg.engine_args.get("resource_budget") for r in replicas]
+        if any(not b for b in budgets):
+            raise ValueError("every stage in a graph pipeline must declare resource_budget capacities and demands")
+        capacities = dict(budgets[0]["capacities"])
+        if any(dict(b["capacities"]) != capacities for b in budgets):
+            raise ValueError("all stages must agree on physical memory ceilings")
+        self.resource_ledger = ResourceLedger(capacities)
+        try:
+            for replica, budget in zip(replicas, budgets):
+                key = (replica.metadata.stage_id, replica.replica_id)
+                self._resource_reservations[key] = self.resource_ledger.reserve(str(key), dict(budget["demands"]))
+        except BaseException:
+            for reservation in self._resource_reservations.values():
+                self.resource_ledger.release(reservation, drained=True)
+            self._resource_reservations.clear()
+            raise
+
+    def _release_unstarted_resources(self, initialized) -> None:
+        if self.resource_ledger is None:
+            return
+        for (stage_id, replica_id), reservation in self._resource_reservations.items():
+            clients = initialized.get(stage_id, [])
+            if replica_id >= len(clients) or clients[replica_id] is None:
+                # Graph factories release or quarantine their own failed load.
+                if reservation.owner not in self.resource_ledger.snapshot()["quarantined"]:
+                    self.resource_ledger.release(
+                        reservation, drained=(stage_id, replica_id) not in self._resource_started
+                    )
+
     def shutdown(self) -> None:
         for pool in self.stage_pools:
             for client in pool.clients:
                 if client is not None and hasattr(client, "shutdown"):
                     try:
                         client.shutdown()
+                        if self.resource_ledger is not None and client.stage_type != "graph":
+                            reservation = self._resource_reservations.get((client.stage_id, client.replica_id))
+                            if reservation is not None:
+                                self.resource_ledger.release(reservation, drained=True)
                     except Exception:
                         logger.warning("[StageRuntime] client shutdown failed", exc_info=True)
         if self._stage_init_executor is not None:
@@ -374,7 +445,7 @@ class StageRuntime:
             stage_vllm_config = None
             executor_class = None
             engine_args_dict = None
-            if base_metadata.stage_type != "diffusion":
+            if base_metadata.stage_type not in ("diffusion", "graph"):
                 # The stable adapter entry point still receives the same
                 # legacy stage object as replica planning. Its implementation
                 # switches only at the coordinated RFC #4021 cutover.
@@ -384,6 +455,7 @@ class StageRuntime:
                     stage_connector_spec=stage_connector_spec,
                     cli_tokenizer=self._tokenizer,
                 )
+                engine_args_dict.pop("resource_budget", None)
                 inject_omni_kv_connector_config(
                     engine_args_dict,
                     omni_kv_connector,
@@ -597,8 +669,7 @@ class StageRuntime:
         if local and all(self._cpu_placed(r) for r in local):
             import psutil
 
-            self._run_host_admission(stage_plans, check_host_admission,
-                                     lambda: int(psutil.virtual_memory().available))
+            self._run_host_admission(stage_plans, check_host_admission, lambda: int(psutil.virtual_memory().available))
             return
 
         def _resolve(replica: ReplicaInitPlan) -> list[int] | AdmissionExempt | None:
@@ -718,6 +789,8 @@ class StageRuntime:
 
     def _init_group_key_override(self, replica: ReplicaInitPlan) -> str | None:
         """Key that bypasses device-overlap grouping, or ``None`` to use it."""
+        if replica.metadata.stage_type == "graph":
+            return "inline:graph"
         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.
@@ -736,6 +809,18 @@ class StageRuntime:
         plan: ReplicaInitPlan,
         stage_init_timeout: int,
     ) -> StagePoolClient:
+        if plan.metadata.stage_type == "graph":
+            if type(self) is not StageRuntime or plan.launch_mode != "local":
+                raise ValueError("graph backends support one local device only")
+            from vllm_omni.engine.backends import create_graph_client
+
+            return create_graph_client(
+                plan.metadata,
+                plan.stage_cfg.engine_args,
+                self.resource_ledger,
+                self._resource_reservations[(plan.metadata.stage_id, plan.replica_id)],
+            )
+        self._resource_started.add((plan.metadata.stage_id, plan.replica_id))
         if plan.launch_mode == "remote":
             return self._initialize_remote_replica(plan, stage_init_timeout)
         if plan.metadata.stage_type == "diffusion":
@@ -985,7 +1070,7 @@ class StageRuntime:
             clients: list[StagePoolClient] = [client for client in replica_clients if client is not None]
             stage_vllm_config = None
             output_processor = None
-            if plan.replicas[0].metadata.stage_type != "diffusion":
+            if plan.replicas[0].metadata.stage_type not in ("diffusion", "graph"):
                 stage_vllm_config = plan.replicas[0].stage_vllm_config
                 if stage_vllm_config is None:
                     raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config")

vllm_omni/entrypoints/async_omni.py

@@ -918,7 +918,10 @@ class AsyncOmni(EngineClient, OmniBase):
                     stage_id,
                     getattr(output_to_yield, "final_output_type", None),
                 )
-                yield output_to_yield
+                try:
+                    yield output_to_yield
+                finally:
+                    output_to_yield.release_stage_buffers()
 
             # The Orchestrator sets "finished" when the final stage is done
             if result.finished:

vllm_omni/entrypoints/omni.py

@@ -40,7 +40,7 @@ class Omni(OmniBase):
             sp = copy.deepcopy(params)
             stage_meta = self.engine.get_stage_metadata(stage_id)
             if (
-                stage_meta.stage_type != "diffusion"
+                stage_meta.stage_type not in ("diffusion", "graph")
                 and hasattr(sp, "output_kind")
                 and sp.output_kind != RequestOutputKind.DELTA
             ):
@@ -191,7 +191,10 @@ class Omni(OmniBase):
                     final_stage_id_for_e2e=req_final_stage_ids[req_id],
                 )
                 if output_to_yield is not None:
-                    yield output_to_yield
+                    try:
+                        yield output_to_yield
+                    finally:
+                        output_to_yield.release_stage_buffers()
 
                 if isinstance(msg, OutputMessage) and msg.finished:
                     active_reqs.discard(req_id)

vllm_omni/entrypoints/stage_utils.py

@@ -191,6 +191,10 @@ def shm_write_bytes(payload: bytes, name: str | None = None) -> dict[str, Any]:
 
     Caller should close the segment; the receiver should unlink.
     """
+    if os.name == "nt":
+        from vllm_omni.host.shared_memory import owner
+
+        return owner.write(payload, name)
     try:
         shm = _shm.SharedMemory(create=True, size=len(payload), name=name)
     except FileExistsError:
@@ -218,6 +222,10 @@ def shm_write_bytes(payload: bytes, name: str | None = None) -> dict[str, Any]:
 
 def shm_read_bytes(meta: dict[str, Any]) -> bytes:
     """Read bytes from SharedMemory by meta {name,size} and cleanup."""
+    if os.name == "nt":
+        from vllm_omni.host.shared_memory import read_buffer
+
+        return read_buffer(meta)
     shm = _shm.SharedMemory(name=meta["name"])  # type: ignore[index]
     mv = memoryview(shm.buf)
     data = bytes(mv[: meta["size"]])

vllm_omni/host/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Host mechanics, independent of accelerator and model selection."""

vllm_omni/host/paths.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Host-local paths, independent of the checkout and model location."""
+
+import os
+import tempfile
+
+
+def device_lock_directory() -> str:
+    # A rooted POSIX path resolves against the current UNC share on Windows.
+    # Opening /tmp from a WSL checkout can block inside the filesystem driver.
+    # All same-user controllers/workers use the same native temporary directory.
+    return tempfile.gettempdir() if os.name == "nt" else "/tmp"

vllm_omni/host/process.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Portable worker ownership; termination reports whether resources drained."""
+
+from __future__ import annotations
+
+import os
+import signal
+import subprocess
+
+
+def spawn_worker(argv: list[str], **kwargs) -> subprocess.Popen:
+    if os.name == "nt":
+        kwargs.setdefault("creationflags", subprocess.CREATE_NO_WINDOW | subprocess.CREATE_NEW_PROCESS_GROUP)
+    else:
+        kwargs.setdefault("start_new_session", True)
+    return subprocess.Popen(argv, **kwargs)
+
+
+def terminate_worker(process: subprocess.Popen, timeout: float = 5.0) -> bool:
+    """Kill descendants and parent, retaining the caller's resource charge on failure."""
+    import psutil
+
+    try:
+        parent = psutil.Process(process.pid)
+        children = parent.children(recursive=True)
+    except psutil.NoSuchProcess:
+        process.wait(timeout=timeout)
+        if os.name == "nt":
+            return False  # absent parent cannot prove that descendants exited
+        try:
+            os.killpg(process.pid, signal.SIGKILL)
+        except ProcessLookupError:
+            return True
+        except OSError:
+            return False
+        return False  # an orphaned group was found; retain its charge until verified
+    except psutil.Error:
+        return False
+    for child in reversed(children):
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            pass
+        except psutil.Error:
+            return False
+    try:
+        if os.name == "nt":
+            process.kill()
+        else:
+            try:
+                os.killpg(process.pid, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+        process.wait(timeout=timeout)
+        _, alive = psutil.wait_procs(children, timeout=timeout)
+        return not alive
+    except (OSError, subprocess.TimeoutExpired):
+        return False

vllm_omni/host/routes.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Deployment-owned worker routes, independent of model and physical device."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from pathlib import Path
+
+from vllm_omni.edge.local.external.launch import Route, is_wsl, resolve
+
+
+def resolve_route(spec: str | Mapping) -> Route:
+    """Names retain compatibility; explicit profiles have no machine-local defaults."""
+    if isinstance(spec, str):
+        return resolve(spec)
+    required = {"name", "interpreter", "worker", "ep", "os_domain"}
+    if not required <= set(spec) or set(spec) - required:
+        raise ValueError(f"explicit route requires exactly {sorted(required)}")
+    if spec["os_domain"] not in {"windows", "posix"}:
+        raise ValueError("route os_domain must be windows or posix")
+    windows = spec["os_domain"] == "windows"
+    interpreter = Path(os.path.expandvars(str(spec["interpreter"]))).expanduser()
+    worker = Path(os.path.expandvars(str(spec["worker"]))).expanduser()
+    reason = ""
+    if windows and os.name != "nt" and not is_wsl():
+        reason = "Windows route requires Windows or WSL interop on this machine"
+    elif not windows and os.name == "nt":
+        reason = "POSIX route needs a POSIX controller; no implicit remote/WSL launch"
+    elif not interpreter.is_file() or not worker.is_file():
+        reason = "configured interpreter or worker script is missing"
+    return Route(str(spec["name"]), str(interpreter), str(worker), str(spec["ep"]), windows, not reason, reason)

vllm_omni/host/shared_memory.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Bounded, single-consumer shared buffers with explicit Windows ownership.
+
+Windows destroys a named mapping when its last handle closes. Keep a producer
+handle until a generation-specific acknowledgement arrives from the reader.
+The header also records payload length: opening a Windows mapping can report a
+page-rounded allocation size. POSIX's existing unlink-based path is unchanged.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import struct
+import threading
+import uuid
+from multiprocessing import shared_memory
+from pathlib import Path
+
+_HEADER = struct.Struct("<8s16sQ")
+_MAGIC = b"OMNISHM1"
+
+
+def _ack_path(name: str, generation: bytes) -> Path:
+    from vllm_omni.windows.paths import shm_dir
+
+    identity = hashlib.sha256(name.encode() + generation).hexdigest()
+    return Path(shm_dir()) / f"ack_{identity}"
+
+
+class SharedBufferOwner:
+    def __init__(self, max_bytes: int = 64 << 20, max_buffers: int = 256):
+        self.max_bytes, self.max_buffers = max_bytes, max_buffers
+        self._owned = {}
+        self._lock = threading.RLock()
+        self._collector = None
+
+    @property
+    def retained_bytes(self) -> int:
+        with self._lock:
+            return sum(size for _, _, size in self._owned.values())
+
+    @property
+    def retained_names(self) -> frozenset[str]:
+        with self._lock:
+            return frozenset(self._owned)
+
+    def collect(self) -> None:
+        with self._lock:
+            for name, (_, generation, _) in list(self._owned.items()):
+                if _ack_path(name, generation).exists():
+                    self.release(name)
+
+    def release(self, name: str) -> None:
+        with self._lock:
+            item = self._owned.get(name)
+            if item is not None:
+                segment, generation, _ = item
+                segment.close()
+                segment.unlink()  # no-op on Windows; useful for contract tests
+                self._owned.pop(name)
+                _ack_path(name, generation).unlink(missing_ok=True)
+
+    def _collect_loop(self) -> None:
+        event = threading.Event()
+        while not event.wait(0.05):
+            try:
+                self.collect()
+            except OSError:
+                # Keep ownership charged if acknowledgement cleanup fails.
+                continue
+
+    def write(self, payload: bytes, name: str | None = None) -> dict:
+        with self._lock:
+            self.collect()
+            size = _HEADER.size + len(payload)
+            if len(self._owned) >= self.max_buffers or self.retained_bytes + size > self.max_bytes:
+                raise BufferError("Windows shared-buffer ownership limit exceeded")
+            generation = uuid.uuid4().bytes
+            segment = shared_memory.SharedMemory(create=True, size=size, name=name)
+            try:
+                segment.buf[: _HEADER.size] = _HEADER.pack(_MAGIC, generation, len(payload))
+                segment.buf[_HEADER.size : size] = payload
+            except BaseException:
+                segment.close()
+                segment.unlink()
+                raise
+            self._owned[segment.name] = (segment, generation, size)
+            if self._collector is None:
+                self._collector = threading.Thread(target=self._collect_loop, daemon=True, name="omni-shm-owner")
+                self._collector.start()
+            return {"name": segment.name, "size": len(payload), "generation": generation.hex()}
+
+
+def read_buffer(meta: dict) -> bytes:
+    segment = shared_memory.SharedMemory(name=meta["name"])
+    try:
+        magic, generation, size = _HEADER.unpack(segment.buf[: _HEADER.size])
+        if magic != _MAGIC or size > segment.size - _HEADER.size:
+            raise ValueError("invalid Windows shared-buffer header")
+        if meta.get("generation", generation.hex()) != generation.hex():
+            raise ValueError("stale Windows shared-buffer generation")
+        payload = bytes(segment.buf[_HEADER.size : _HEADER.size + size])
+    finally:
+        segment.close()
+    _ack_path(meta["name"], generation).touch()
+    return payload
+
+
+owner = SharedBufferOwner()

vllm_omni/host/windows_process.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Standalone native-Windows process identity and bounded tree retirement.
+
+This file is executed by the worker interpreter from WSL; it imports no engine
+or third-party packages. A launcher PID is never treated as a Windows PID.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import json
+import os
+import subprocess
+import sys
+from ctypes import wintypes
+
+
+def _api():
+    kernel = ctypes.WinDLL("kernel32", use_last_error=True)
+    kernel.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
+    kernel.OpenProcess.restype = wintypes.HANDLE
+    kernel.CloseHandle.argtypes = [wintypes.HANDLE]
+    kernel.GetProcessTimes.argtypes = [wintypes.HANDLE] + [ctypes.POINTER(wintypes.FILETIME)] * 4
+    kernel.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD]
+    return kernel
+
+
+def _created(kernel, handle) -> int:
+    times = [wintypes.FILETIME() for _ in range(4)]
+    if not kernel.GetProcessTimes(handle, *(ctypes.byref(t) for t in times)):
+        raise ctypes.WinError(ctypes.get_last_error())
+    return (times[0].dwHighDateTime << 32) | times[0].dwLowDateTime
+
+
+def identity() -> dict:
+    kernel = _api()
+    handle = kernel.OpenProcess(0x1000, False, os.getpid())  # QUERY_LIMITED_INFORMATION
+    if not handle:
+        raise ctypes.WinError(ctypes.get_last_error())
+    try:
+        return {"pid": os.getpid(), "created_filetime": _created(kernel, handle)}
+    finally:
+        kernel.CloseHandle(handle)
+
+
+def retire(pid: int, created_filetime: int) -> bool:
+    kernel = _api()
+    handle = kernel.OpenProcess(0x1000 | 0x100000, False, pid)  # query + synchronize
+    if not handle:
+        return False  # parent absence alone says nothing about its descendants
+    try:
+        if _created(kernel, handle) != created_filetime:
+            return False  # stale identity must not kill a reused PID
+        result = subprocess.run(
+            ["taskkill.exe", "/PID", str(pid), "/T", "/F"],
+            capture_output=True,
+            timeout=10,
+            creationflags=subprocess.CREATE_NO_WINDOW,
+        )
+        return result.returncode == 0 and kernel.WaitForSingleObject(handle, 5000) == 0
+    finally:
+        kernel.CloseHandle(handle)
+
+
+if __name__ == "__main__":
+    try:
+        print(json.dumps({"drained": retire(int(sys.argv[1]), int(sys.argv[2]))}))
+    except Exception as exc:
+        print(json.dumps({"drained": False, "error": str(exc)}))
+        sys.exit(1)

vllm_omni/host/worker_bootstrap.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Standalone worker bootstrap: record OS identity without importing Omni."""
+
+import json
+import os
+import runpy
+import sys
+from pathlib import Path
+
+if os.name == "nt":
+    from windows_process import identity
+
+    os.environ["VLLM_OMNI_WORKER_IDENTITY"] = json.dumps(identity())
+else:
+    os.environ.pop("VLLM_OMNI_WORKER_IDENTITY", None)
+
+worker = sys.argv.pop(1)
+sys.argv[0] = worker
+sys.path.insert(0, str(Path(worker).resolve().parent))
+runpy.run_path(worker, run_name="__main__")

vllm_omni/outputs/__init__.py

@@ -96,6 +96,7 @@ _REQUEST_OUTPUT_CONTENT_ATTRS = (
 # Omni-specific content copied when the wrapped stage output is itself an
 # OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline).
 _OMNI_CONTENT_ATTRS = (
+    "_stage_release",
     "images",
     "latents",
     "trajectory_latents",
@@ -170,6 +171,13 @@ class OmniRequestOutput(RequestOutput):
     trajectory_decoded: list | None = None
     _multimodal_output: dict[str, Any] = field(default_factory=dict)
     _custom_output: dict[str, Any] = field(default_factory=dict)
+    _stage_release: Any = field(default=None, repr=False, compare=False)
+
+    def release_stage_buffers(self) -> None:
+        """Acknowledge consumption of an externally owned stage payload."""
+        release, self._stage_release = self._stage_release, None
+        if release is not None:
+            release()
 
     # profiling data
     stage_durations: dict[str, float] = field(default_factory=dict)

vllm_omni/windows/paths.py

@@ -20,7 +20,7 @@ def shm_dir() -> str:
     use. Callers that used to write ``f"/dev/shm/..."`` go through this so the
     path is the same on every platform that has one and a real one elsewhere.
     """
-    if os.path.isdir(_POSIX_SHM):
+    if sys.platform != "win32" and os.path.isdir(_POSIX_SHM):
         return _POSIX_SHM
     base = os.environ.get("LOCALAPPDATA") or tempfile.gettempdir()
     path = Path(base) / "vllm_omni" / "shm"

vllm_omni/windows/shims.py

@@ -18,6 +18,7 @@ built on injectable primitives so the logic is testable on Linux.
 from __future__ import annotations
 
 import asyncio
+import errno
 import os
 import sys
 import time
@@ -88,8 +89,10 @@ class FcntlShim:
                 self._at_lock_offset(fd, lambda: self._locking(fd, self._nblck, 1))
                 return
             except OSError as exc:
+                if exc.errno not in (errno.EACCES, errno.EAGAIN, errno.EDEADLK):
+                    raise
                 if non_blocking:
-                    raise BlockingIOError(str(exc)) from exc
+                    raise BlockingIOError(exc.errno, str(exc)) from exc
                 self._sleep(self._poll_interval)
 
     def lockf(self, fd: int, cmd: int, len: int = 0, start: int = 0, whence: int = 0) -> None:  # noqa: A002

vllm_omni/worker/gpu_model_runner.py

@@ -2035,7 +2035,11 @@ class OmniGPUModelRunner(GPUModelRunner):
 
         accepted: frozenset[str] | None
         try:
-            parameters = inspect.signature(type(self.model).forward).parameters
+            # CUDA graph wrappers delegate attributes to their wrapped model;
+            # their class does not define ``forward``. Inspect the resolved
+            # callable so eager and captured native models use the same filter.
+            forward = getattr(self.model, "forward", self.model)
+            parameters = inspect.signature(forward).parameters
         except (TypeError, ValueError):
             accepted = None  # not introspectable; assume it copes, as before
         else: