[edge-infer] WIP snapshot: uncommitted M0/W1/W4/W5 edge work as of 2026-09-18

Verbatim copy of the edge/omni-edge working tree (tracked diffs + untracked
edge/local, tests/edge/local, docs/edge, pipeline files). Taken so that
vLLM-Omni v0.29.0rc1 can be merged on a branch without touching the
user's working tree. logs_edge/ and benchmark .npy outputs excluded.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

tests/config/test_vllm_native_pipelines.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Pipelines for text models vLLM implements and Omni does not own.
+
+A model being in vLLM's registry does not make it reachable through
+``Omni(model=...)``: the orchestrator resolves a pipeline by ``model_type``
+before any weight is read, so a ``model_type`` with no entry fails early and
+opaquely. Spark-X2.5 hit exactly that; these tests pin the general case.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.stage_config import StageExecutionType
+from vllm_omni.config.vllm_native_pipelines import LLAMA_PIPELINE
+
+
+def test_llama_resolves_from_the_registry():
+    """MiniCPM5-2B declares model_type "llama"; without this entry the
+    orchestrator cannot resolve it at all."""
+    assert OMNI_PIPELINES["llama"] is LLAMA_PIPELINE
+
+
+def test_it_is_one_autoregressive_text_stage():
+    assert LLAMA_PIPELINE.get_validation_errors() == []
+    assert len(LLAMA_PIPELINE.stages) == 1
+    stage = LLAMA_PIPELINE.stages[0]
+    assert stage.execution_type is StageExecutionType.LLM_AR
+    assert stage.input_sources == ()
+    assert stage.final_output and stage.final_output_type == "text"
+    assert stage.owns_tokenizer
+    assert stage.model_arch == "LlamaForCausalLM"
+
+
+def test_the_stage_declares_no_cross_chunk_state():
+    """A causal LM's state is its KV cache, which the backend owns. Claiming
+    otherwise would make the scheduler park requests to hold capacity."""
+    assert not LLAMA_PIPELINE.stages[0].retains_state_across_chunks
+
+
+def test_registering_llama_is_keyed_on_architecture_not_checkpoint():
+    """Worth pinning because it is the surprising part: this entry covers every
+    checkpoint whose config says "llama", not only the one it was added for.
+    That is the right granularity — pipeline shape is a property of the
+    architecture — but it is a claim about shape, not a verification of any
+    particular weights."""
+    assert LLAMA_PIPELINE.model_type == "llama"
+    assert LLAMA_PIPELINE.hf_architectures == ("LlamaForCausalLM",)

tests/edge/local/conftest.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fixtures for the M0 local text mode.
+
+Everything here is synthetic on purpose: a hardware profile built from
+literals, and a checkpoint directory with a real ``config.json`` and
+zero-filled weight files. The planner reads sizes and geometry, never tensors,
+so this exercises the real code path without a GPU or 7.7 GB of weights. The
+one test module that does need real weights says so and skips.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.edge.hardware_probe import (
+    ACCEL_GPU_DISCRETE,
+    ACCEL_GPU_INTEGRATED,
+    ACCEL_NPU,
+    HardwareProfile,
+)
+
+# A 28-layer Spark-shaped decoder: 7 full-attention and 21 sliding@512, which
+# is the 1.7B layout. Small hidden/vocab so nothing here is slow.
+SPARK_SHAPED_CONFIG = {
+    "architectures": ["Spark2_5ForCausalLM"],
+    "model_type": "spark2_5",
+    "dtype": "bfloat16",
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 2,
+    "head_dim": 256,
+    "hidden_act": "gelu",
+    "vocab_size": 131072,
+    "sliding_window": 512,
+    "max_position_embeddings": 1048576,
+    "tie_word_embeddings": True,
+    "layer_types": [
+        "full_attention" if (i + 1) % 4 == 0 else "sliding_attention"
+        for i in range(28)
+    ],
+    "rope_parameters": {
+        "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25},
+        "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0},
+    },
+}
+
+INT8_QUANT_CONFIG = {
+    "quant_method": "compressed-tensors",
+    "format": "int-quantized",
+    "quantization_status": "compressed",
+    "ignore": ["lm_head"],
+    "config_groups": {
+        "group_0": {
+            "targets": ["Linear"],
+            "weights": {"num_bits": 8, "type": "int", "symmetric": True, "strategy": "channel"},
+        }
+    },
+}
+
+FP8_QUANT_CONFIG = {
+    "quant_method": "compressed-tensors",
+    "format": "float-quantized",
+    "config_groups": {},
+}
+
+
+def write_checkpoint(
+    root: Path,
+    *,
+    quantization: dict | None = None,
+    weight_bytes: int = 4 * 1024 * 1024,
+    shards: int = 1,
+) -> Path:
+    """A checkpoint directory the manifest and planner can read end to end."""
+    root.mkdir(parents=True, exist_ok=True)
+    config = dict(SPARK_SHAPED_CONFIG)
+    if quantization is not None:
+        config["quantization_config"] = quantization
+    (root / "config.json").write_text(json.dumps(config), encoding="utf-8")
+    (root / "tokenizer_config.json").write_text('{"model_max_length": 4096}', encoding="utf-8")
+    (root / "chat_template.jinja").write_text("{{ messages }}", encoding="utf-8")
+    per_shard = max(weight_bytes // shards, 1)
+    for i in range(shards):
+        name = "model.safetensors" if shards == 1 else f"model-{i + 1:05d}-of-{shards:05d}.safetensors"
+        with (root / name).open("wb") as fh:
+            fh.truncate(per_shard)
+    return root
+
+
+@pytest.fixture
+def dense_checkpoint(tmp_path: Path) -> Path:
+    return write_checkpoint(tmp_path / "dense", weight_bytes=8 * 1024 * 1024, shards=2)
+
+
+@pytest.fixture
+def int8_checkpoint(tmp_path: Path) -> Path:
+    return write_checkpoint(tmp_path / "int8", quantization=INT8_QUANT_CONFIG)
+
+
+def blackwell_laptop(ram_available_bytes: int = 28 * 2**30) -> HardwareProfile:
+    """The reference machine: sm_120 discrete GPU, an iGPU and an NPU beside it.
+
+    Modelled on the probe output recorded in the architecture proposal, so the
+    tests exercise the same shape the real host produces.
+    """
+    return HardwareProfile(
+        arch="x86_64",
+        cpu_model="AMD Ryzen AI 9 HX 370 w/ Radeon 890M",
+        cpu_flags=["avx2", "avx512", "avx512_bf16", "avx512_vnni", "fma"],
+        n_logical=24,
+        n_physical=12,
+        clusters=[{"max_khz": 5000000, "cpus": list(range(24))}],
+        ram_total_bytes=int(30.91 * 2**30),
+        ram_available_bytes=ram_available_bytes,
+        accelerator="cuda_discrete",
+        gpu_name="NVIDIA GeForce RTX 5090 Laptop GPU",
+        gpu_mem_bytes=24463 * 2**20,
+        gpu_sm_count=110,
+        gpu_capability=[12, 0],
+        supports_bf16=True,
+        host_os="wsl2",
+        accelerators=[
+            {"kind": ACCEL_GPU_DISCRETE, "vendor": "nvidia",
+             "name": "NVIDIA GeForce RTX 5090 Laptop GPU", "usable": True,
+             "source": "torch.cuda", "mem_bytes": 24463 * 2**20},
+            {"kind": ACCEL_NPU, "vendor": "amd", "name": "NPU Compute Accelerator Device",
+             "usable": False, "source": "wsl-host", "reason": "no device node under WSL2"},
+            {"kind": ACCEL_GPU_INTEGRATED, "vendor": "amd", "name": "AMD Radeon(TM) 890M Graphics",
+             "usable": False, "source": "wsl-host", "reason": "no /dev/dri under WSL2"},
+        ],
+        source="test-fixture",
+    )
+
+
+def cpu_only_box() -> HardwareProfile:
+    """The no-NVIDIA deployment class, as a machine rather than as a mask."""
+    profile = blackwell_laptop()
+    profile.accelerator = "none"
+    profile.gpu_name = ""
+    profile.gpu_mem_bytes = 0
+    profile.gpu_capability = []
+    profile.accelerators = [a for a in profile.accelerators if a["vendor"] != "nvidia"]
+    return profile
+
+
+@pytest.fixture
+def laptop_profile() -> HardwareProfile:
+    return blackwell_laptop()

tests/edge/local/external/conftest.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fixtures that stand in for the two AMD devices.
+
+Nothing here needs a Radeon, an NPU, onnxruntime or Windows: the worker is
+replaced by ``fake_worker.py`` and the devices by hand-built capabilities. The
+paths that matter -- framing, the placement gate, the refusal codes -- are
+logic, and logic that only runs on one laptop is logic nobody can change
+safely.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.edge.local.capabilities import (
+    FORMAT_ONNX_A16W8,
+    FORMAT_ONNX_FP16,
+    FORMAT_ONNX_FP32,
+    DeviceCapability,
+)
+from vllm_omni.edge.local.external import launch
+
+FAKE_WORKER = str(Path(__file__).resolve().parent / "fake_worker.py")
+
+
+@pytest.fixture
+def fake_route(monkeypatch, tmp_path):
+    """Point the ``ort-cpu`` route at the fake worker and this interpreter.
+
+    Returns a callable that writes the worker's reply script and hands back the
+    resolved route.
+    """
+
+    def configure(**config) -> launch.Route:
+        config_path = tmp_path / "fake_worker.json"
+        config_path.write_text(json.dumps(config))
+        monkeypatch.setenv("VLLM_OMNI_FAKE_WORKER_JSON", str(config_path))
+        monkeypatch.setenv("VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU", FAKE_WORKER)
+        monkeypatch.setenv("VLLM_OMNI_EXTERNAL_PYTHON_ORT_CPU", sys.executable)
+        return launch.resolve(launch.ROUTE_CPU)
+
+    return configure
+
+
+def _device(device_id: str, kind: str, formats, route: str, *, runnable: bool = True,
+            reason: str = "") -> DeviceCapability:
+    return DeviceCapability(
+        device_id=device_id,
+        kind=kind,
+        vendor="amd",
+        name=f"test {kind}",
+        backend="ort:cpu" if runnable else None,
+        runnable=runnable,
+        memory_pool="host_ram",
+        memory_bytes=32 * 2**30,
+        compute_capability=None,
+        weight_formats=frozenset(formats),
+        evidence="B" if runnable else "D",
+        reason=reason,
+        extra={"worker_route": route, "routes": [{"name": route, "available": runnable,
+                                                  "reason": reason}]},
+    )
+
+
+@pytest.fixture
+def npu_device() -> DeviceCapability:
+    """A16W8 only -- the real constraint, not a simplification."""
+    return _device("npu:amd", "npu", {FORMAT_ONNX_A16W8}, launch.ROUTE_CPU)
+
+
+@pytest.fixture
+def igpu_device() -> DeviceCapability:
+    return _device("igpu:amd", "gpu_integrated", {FORMAT_ONNX_FP16, FORMAT_ONNX_FP32},
+                   launch.ROUTE_CPU)
+
+
+@pytest.fixture
+def unreachable_npu() -> DeviceCapability:
+    return _device(
+        "npu:amd", "npu", {FORMAT_ONNX_A16W8}, launch.ROUTE_VITISAI,
+        runnable=False,
+        reason="an MCDM device; WSL2 forwards display adapters only, so there is no device node here",
+    )
+
+
+@pytest.fixture
+def a16w8_graph(tmp_path) -> Path:
+    """A file that stands in for an exported graph. Never loaded by ORT here."""
+    path = tmp_path / "stage.onnx"
+    path.write_bytes(b"not really onnx, and nothing in these tests parses it" * 64)
+    return path

tests/edge/local/external/fake_worker.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""A stage worker that speaks the protocol and never touches a device.
+
+Selected with ``VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU``, so the client, the
+framing, the placement gate and every refusal path are exercised on a machine
+with no AMD hardware, no onnxruntime and no Windows. What it replies is read
+from the JSON file named by ``VLLM_OMNI_FAKE_WORKER_JSON``:
+
+``load``      body merged into the load reply -- set ``node_counts`` and
+              ``fraction_on_target`` to stage a declined graph.
+``run``       ``{"scale": float}``; outputs are the inputs times it.
+``die_on``    op name after which the worker exits without replying, for the
+              "worker vanished mid-request" path.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+_HERE = Path(__file__).resolve()
+_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)
+sys.modules["_fake_protocol"] = proto
+spec.loader.exec_module(proto)
+
+
+def main() -> int:
+    host, port, token = sys.argv[2], int(sys.argv[4]), sys.argv[6]
+    config = {}
+    config_path = os.environ.get("VLLM_OMNI_FAKE_WORKER_JSON")
+    if config_path and Path(config_path).is_file():
+        config = json.loads(Path(config_path).read_text())
+
+    import socket
+
+    sock = socket.create_connection((host, port), timeout=30)
+    proto.send_message(
+        sock,
+        proto.OP_HELLO,
+        {
+            "token": config.get("token", token),
+            "executable": sys.executable,
+            "python": sys.version,
+            "platform": sys.platform,
+            "pid": os.getpid(),
+            "onnxruntime": "fake",
+            "numpy": np.__version__,
+            "available_providers": ["CPUExecutionProvider"],
+            "rss_bytes": 1024,
+        },
+    )
+
+    while True:
+        op, body, tensors = proto.recv_message(sock)
+        if config.get("die_on") == op:
+            sock.close()
+            return 0
+        if op == proto.OP_LOAD:
+            reply = {
+                "ep": body.get("ep", "cpu"),
+                "session_providers": ["CPUExecutionProvider"],
+                "available_providers": ["CPUExecutionProvider"],
+                "node_counts": {"CPUExecutionProvider": 4},
+                "total_nodes": 4,
+                "target_nodes": 4,
+                "fraction_on_target": 1.0,
+                "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()],
+                "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:
+            scale = float(config.get("run", {}).get("scale", 2.0))
+            proto.send_message(
+                sock,
+                proto.OP_OK,
+                {"run_s": 0.001},
+                {f"{k}_out": (v * scale) for k, v in tensors.items()},
+            )
+        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": {}},
+            )
+        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"})
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/edge/local/external/test_client.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Starting, driving and losing a worker process."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from vllm_omni.edge.local.external import launch
+from vllm_omni.edge.local.external.client import ExternalWorker, WorkerStartError
+
+
+def test_handshake_and_round_trip(fake_route):
+    with ExternalWorker(fake_route()) as worker:
+        assert worker.hello["onnxruntime"] == "fake"
+        report = worker.load("/nonexistent/stage.onnx", example_inputs={"x": np.ones(4, np.float32)})
+        assert report.fraction_on_target == 1.0
+        assert report.rss_bytes == 111 * 2**20
+
+        outputs, timing = worker.run({"x": np.ones(4, np.float32)})
+        assert np.array_equal(outputs["x_out"], np.full(4, 2.0, np.float32))
+        assert timing.round_trip_s >= timing.worker_s
+        assert timing.transport_s == pytest.approx(timing.round_trip_s - timing.worker_s)
+
+        assert worker.stats()["peak_rss_bytes"] == 122 * 2**20
+
+
+def test_wrong_token_is_refused(fake_route):
+    """The listener is on a host-visible address; a stray peer must not be served."""
+    worker = ExternalWorker(fake_route(token="not-the-token"))
+    with pytest.raises(WorkerStartError, match="wrong token"):
+        worker.start()
+    worker.close()
+
+
+def test_worker_that_vanishes_mid_request_reports_rather_than_hangs(fake_route):
+    worker = ExternalWorker(fake_route(die_on="run"))
+    worker.start()
+    worker.load("/nonexistent/stage.onnx", example_inputs={"x": np.ones(2, np.float32)})
+    with pytest.raises(WorkerStartError, match="worker connection failed"):
+        worker.run({"x": np.ones(2, np.float32)})
+    worker.close()
+
+
+def test_unavailable_route_refuses_at_construction():
+    route = launch.Route(
+        name="ort-vitisai", interpreter=None, worker="w.py", ep="vitisai",
+        is_windows=True, available=False, reason="no interpreter installed",
+    )
+    with pytest.raises(WorkerStartError, match="no interpreter installed"):
+        ExternalWorker(route)
+
+
+def test_close_is_idempotent(fake_route):
+    worker = ExternalWorker(fake_route())
+    worker.start()
+    worker.close()
+    worker.close()
+
+
+def test_bind_address_is_loopback_for_a_same_os_worker():
+    from vllm_omni.edge.local.external import client
+
+    assert client._bind_address(is_windows=False) == "127.0.0.1"
+
+
+def test_bind_address_for_a_windows_worker_is_not_loopback_under_wsl():
+    """The regression this exists for: probing the WSL nameserver instead of the
+    gateway yields 10.255.255.254, an address only this distro can reach, and
+    the Windows worker then dies with ECONNREFUSED after a silent start."""
+    from vllm_omni.edge.local.external import client
+
+    if not launch.is_wsl():
+        pytest.skip("not running under WSL interop")
+    assert not client._bind_address(is_windows=True).startswith(("127.", "10.255."))
+
+
+def test_external_data_counts_toward_the_artifact_size(tmp_path):
+    """An ONNX graph with external weights is a few MB of topology pointing at
+    gigabytes beside it. Charging only the .onnx told the admission ledger that
+    a 1.8 GB vision tower costs 2 MiB."""
+    import onnx
+    from onnx import helper, numpy_helper
+
+    from vllm_omni.edge.local.manifest import build_graph_artifact
+
+    weight = numpy_helper.from_array(np.zeros((256, 256), np.float32), name="w")
+    node = helper.make_node("MatMul", ["x", "w"], ["y"])
+    graph = helper.make_graph(
+        [node], "g",
+        [helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 256])],
+        [helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1, 256])],
+        [weight],
+    )
+    path = tmp_path / "m.onnx"
+    onnx.save(helper.make_model(graph), str(path), save_as_external_data=True,
+              all_tensors_to_one_file=True, location="m.onnx.data")
+
+    artifact = build_graph_artifact(path, fmt="onnx:fp32", opset=21, source_model="t",
+                                    component="c", exporter="e")
+    assert artifact.bytes > path.stat().st_size
+    assert artifact.bytes >= 256 * 256 * 4

tests/edge/local/external/test_protocol.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The wire format, which has to survive a peer that is not us."""
+
+from __future__ import annotations
+
+import socket
+import threading
+
+import numpy as np
+import pytest
+
+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()
+        return protocol.recv_message(far)
+    finally:
+        near.close()
+        far.close()
+
+
+def test_tensors_survive_dtype_and_shape():
+    sent = {
+        "f32": np.arange(12, dtype=np.float32).reshape(3, 4),
+        "i64": np.array([[1, 2], [3, 4]], dtype=np.int64),
+        "u16": np.array([7, 8, 9], dtype=np.uint16),
+        "scalar": np.array(3.5, dtype=np.float64),
+    }
+    op, body, got = _round_trip(protocol.OP_RUN, {"n": 1}, sent)
+    assert op == protocol.OP_RUN and body == {"n": 1}
+    for name, array in sent.items():
+        assert got[name].dtype == array.dtype
+        assert got[name].shape == array.shape
+        assert np.array_equal(got[name], array)
+
+
+def test_received_tensors_are_writable():
+    """``frombuffer`` hands back a read-only view aliasing the socket buffer.
+
+    ORT's IOBinding writes into the arrays it is given, so a read-only result
+    would fail at the far end of a working connection rather than here.
+    """
+    _, _, got = _round_trip(protocol.OP_RUN, {}, {"x": np.ones(4, dtype=np.float32)})
+    got["x"][0] = 5.0  # must not raise
+    assert got["x"][0] == 5.0
+
+
+def test_payload_larger_than_one_recv():
+    """Stream sockets short-read; a naive recv would misparse the next tensor."""
+    big = np.arange(4 * 1024 * 1024, dtype=np.float32)
+    _, _, got = _round_trip(protocol.OP_RUN, {}, {"big": big, "after": np.array([1.0])})
+    assert np.array_equal(got["big"], big)
+    assert np.array_equal(got["after"], np.array([1.0]))
+
+
+def test_empty_message_and_no_tensors():
+    op, body, got = _round_trip(protocol.OP_STATS)
+    assert (op, body, got) == (protocol.OP_STATS, {}, {})
+
+
+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"})
+    assert excinfo.value.code == "ValueError"
+    assert "far side traceback" in excinfo.value.remote_traceback
+
+
+def test_closed_peer_is_a_protocol_error_not_a_hang():
+    near, far = socket.socketpair()
+    near.close()
+    with pytest.raises(protocol.ProtocolError):
+        protocol.recv_message(far)
+    far.close()
+
+
+def test_oversized_header_is_refused_before_allocation():
+    """A corrupt length must not become a multi-gigabyte allocation."""
+    near, far = socket.socketpair()
+    try:
+        import struct
+
+        near.sendall(struct.pack(">I", protocol.MAX_HEADER_BYTES + 1))
+        with pytest.raises(protocol.ProtocolError, match="over the limit"):
+            protocol.recv_message(far)
+    finally:
+        near.close()
+        far.close()

tests/edge/local/external/test_stage.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Placing a stage on an external device, and refusing to.
+
+The gate tested here is the one that keeps a CPU fallback from being reported
+as NPU execution. It cannot be replaced by a cheaper check: a VitisAI session
+lists the EP in ``get_providers()`` whether or not it claimed a node, and a
+graph it declined runs on the CPU and returns bit-identical numbers. So these
+tests assert on node assignments, never on output values.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from vllm_omni.edge.local.capabilities import FORMAT_ONNX_A16W8, FORMAT_ONNX_FP16
+from vllm_omni.edge.local.external.stage import (
+    ExternalStage,
+    PlacementRefused,
+    plan_external_stage,
+    verify_placement,
+)
+from vllm_omni.edge.local.manifest import build_graph_artifact
+from vllm_omni.edge.local.plan import (
+    REFUSE_EP_PLACEMENT,
+    REFUSE_NO_ARTIFACT,
+    REFUSE_NO_DEVICE,
+    REFUSE_ROUTE,
+)
+
+
+def _artifact(path, fmt=FORMAT_ONNX_A16W8, component="stage"):
+    return build_graph_artifact(
+        path, fmt=fmt, opset=21, source_model="test", component=component,
+        exporter="tests.edge.local.external",
+    )
+
+
+def test_a16w8_graph_is_placed_on_the_npu(a16w8_graph, npu_device, igpu_device):
+    plan = plan_external_stage(_artifact(a16w8_graph), [igpu_device, npu_device])
+    assert plan.admitted
+    assert plan.device is not None and plan.device.device_id == "npu:amd"
+    # The iGPU was considered and rejected on format; the record says so rather
+    # than only naming the winner.
+    assert [r.code for r in plan.refusals] == [REFUSE_NO_ARTIFACT]
+
+
+def test_fp16_graph_is_refused_by_the_npu_not_quietly_downgraded(a16w8_graph, npu_device):
+    """A16W8 is the only door on XDNA2, and the failure mode for anything else
+    is silent: no error, every node on the CPU. So it is refused at plan time."""
+    plan = plan_external_stage(_artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device])
+    assert not plan.admitted
+    assert plan.refusals[0].code == REFUSE_NO_ARTIFACT
+    assert "a16w8" in plan.refusals[0].remedy.lower()
+
+
+def test_require_refuses_rather_than_substituting(a16w8_graph, npu_device, igpu_device):
+    """"Does the NPU run this" must not be answered with the iGPU."""
+    plan = plan_external_stage(
+        _artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device, igpu_device],
+        require="npu:amd",
+    )
+    assert not plan.admitted
+    assert {r.device_id for r in plan.refusals} == {"npu:amd"}
+
+
+def test_prefer_falls_through(a16w8_graph, npu_device, igpu_device):
+    plan = plan_external_stage(
+        _artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device, igpu_device],
+        prefer="npu:amd",
+    )
+    assert plan.admitted and plan.device.device_id == "igpu:amd"
+
+
+def test_unreachable_device_refuses_with_the_route_reason(a16w8_graph, unreachable_npu):
+    plan = plan_external_stage(_artifact(a16w8_graph), [unreachable_npu])
+    assert not plan.admitted
+    assert plan.refusals[0].code == REFUSE_ROUTE
+    assert "MCDM" in plan.refusals[0].message
+    assert "VLLM_OMNI_EXTERNAL_PYTHON_ORT_VITISAI" in plan.refusals[0].remedy
+
+
+def test_no_external_device_at_all(a16w8_graph):
+    plan = plan_external_stage(_artifact(a16w8_graph), [])
+    assert not plan.admitted and plan.refusals[0].code == REFUSE_NO_DEVICE
+
+
+def test_required_device_absent_from_enumeration(a16w8_graph, igpu_device):
+    plan = plan_external_stage(_artifact(a16w8_graph), [igpu_device], require="npu:amd")
+    assert not plan.admitted and plan.refusals[0].code == REFUSE_NO_DEVICE
+
+
+def test_budget_is_on_the_shared_pool(a16w8_graph, npu_device):
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device])
+    assert {r.pool for r in plan.reservations} == {"host_ram"}
+    assert plan.budget_bytes > 0
+
+
+def test_open_verifies_placement_and_closes_the_ledger(a16w8_graph, npu_device, fake_route, tmp_path):
+    fake_route()
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device])
+    with ExternalStage(plan) as stage:
+        report = stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path)
+        assert report.fraction_on_target == 1.0
+    runtime = next(r for r in plan.reservations if r.purpose == "worker runtime")
+    assert runtime.actual_upper_bound_bytes == 111 * 2**20
+    assert runtime.bytes >= runtime.actual_upper_bound_bytes, "budget must bound the measurement"
+
+
+def test_declined_graph_is_refused_even_though_it_ran(a16w8_graph, npu_device, fake_route, tmp_path):
+    """The whole point: it produced correct output, on the CPU."""
+    fake_route(load={"node_counts": {"CPUExecutionProvider": 9, "vitisai": 1},
+                     "total_nodes": 10, "target_nodes": 1, "fraction_on_target": 0.1})
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device])
+    stage = ExternalStage(plan)
+    with pytest.raises(PlacementRefused) as excinfo:
+        stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path)
+    assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT
+    assert "1 of 10" in excinfo.value.refusal.message
+    assert "Light" in excinfo.value.refusal.remedy
+
+
+def test_unverified_placement_is_refused_not_assumed(a16w8_graph, npu_device, fake_route, tmp_path):
+    """``None`` means nobody looked. It must not be rounded to "fine"."""
+    fake_route(load={"fraction_on_target": None, "node_counts": {},
+                     "placement_unknown_reason": "no profiled run"})
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device])
+    stage = ExternalStage(plan)
+    with pytest.raises(PlacementRefused) as excinfo:
+        stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path)
+    assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT
+    assert "not measured" in excinfo.value.refusal.message
+
+
+def test_open_without_example_inputs_is_a_usage_error(a16w8_graph, npu_device, fake_route):
+    fake_route()
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device])
+    stage = ExternalStage(plan)
+    with pytest.raises(ValueError, match="example_inputs are required"):
+        stage.open({})
+    stage.close()
+
+
+def test_verify_placement_accepts_a_real_split(a16w8_graph, npu_device):
+    from vllm_omni.edge.local.external.client import LoadReport
+
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device])
+    plan.min_fraction_on_target = 0.3
+    report = LoadReport(ep="vitisai", graph_path="g.onnx", total_nodes=3, target_nodes=1,
+                        fraction_on_target=1 / 3, node_counts={"vitisai": 1, "CPUExecutionProvider": 2})
+    assert verify_placement(plan, report) is None
+
+
+def test_refusing_a_plan_cannot_be_opened(a16w8_graph, npu_device):
+    plan = plan_external_stage(_artifact(a16w8_graph, fmt=FORMAT_ONNX_FP16), [npu_device])
+    with pytest.raises(PlacementRefused):
+        ExternalStage(plan)
+
+
+def test_a_masked_device_is_told_to_drop_the_mask(a16w8_graph, npu_device):
+    """The remedy has to name the cause the caller actually has.
+
+    A masked device and a device WSL cannot forward are both "unreachable", and
+    explaining MCDM forwarding to somebody who simply passed ``--mask npu``
+    sends them after a problem they do not have.
+    """
+    import dataclasses
+
+    masked = dataclasses.replace(npu_device, runnable=False, masked=True,
+                                 backend=None, reason="masked out by --mask (npu)")
+    plan = plan_external_stage(_artifact(a16w8_graph), [masked])
+    assert not plan.admitted
+    assert plan.refusals[0].code == REFUSE_ROUTE
+    assert "--mask" in plan.refusals[0].remedy
+    assert "MCDM" not in plan.refusals[0].remedy
+
+
+def test_a_stage_larger_than_the_shared_pool_is_refused(a16w8_graph, npu_device):
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device], pool_capacity_bytes=1)
+    assert not plan.admitted
+    from vllm_omni.edge.local.plan import REFUSE_CAPACITY
+
+    assert plan.refusals[0].code == REFUSE_CAPACITY
+    assert "same pool the CPU stage is budgeted against" in plan.refusals[0].message
+
+
+def test_a_silently_swapped_provider_is_refused_by_name(a16w8_graph, npu_device, fake_route, tmp_path):
+    """onnxruntime does not raise when a provider fails to initialise.
+
+    It prints "Falling back to ['CPUExecutionProvider'] and retrying" and hands
+    back a working CPU session — which is how a 27B vision tower ran end to end
+    on the CPU here while the caller believed it was on the iGPU. The report
+    carries the swap, and it outranks the node count because "the provider
+    never initialised" says what to fix and "0 of 1171 nodes" does not.
+    """
+    fake_route(load={
+        "session_providers": ["CPUExecutionProvider"],
+        "requested_provider_missing": "DmlExecutionProvider",
+        "fallback_note": "DmlExecutionProvider failed to initialise and onnxruntime silently fell back",
+        "node_counts": {"CPUExecutionProvider": 113},
+        "total_nodes": 113, "target_nodes": 0, "fraction_on_target": 0.0,
+    })
+    plan = plan_external_stage(_artifact(a16w8_graph), [npu_device])
+    stage = ExternalStage(plan)
+    with pytest.raises(PlacementRefused) as excinfo:
+        stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path)
+    assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT
+    assert "silently fell back" in excinfo.value.refusal.message

tests/edge/local/external/test_vision_tower.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shape arithmetic for the exported vision tower.
+
+The tower itself needs a 29 GB checkpoint, so what is covered here is the part
+that decides whether an export is even well-formed — and that is where a silent
+mistake would land: a grid that does not divide by the merge size produces a
+graph that exports, runs, and returns the wrong number of tokens to the LLM.
+"""
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from vllm_omni.edge.local.artifacts.vision_tower import tower_shape
+
+VISION_CONFIG = {
+    "depth": 27, "hidden_size": 1152, "in_channels": 3, "intermediate_size": 4304,
+    "num_heads": 16, "num_position_embeddings": 2304, "out_hidden_size": 5120,
+    "patch_size": 16, "spatial_merge_size": 2, "temporal_patch_size": 2,
+}
+
+
+@pytest.fixture
+def checkpoint(tmp_path):
+    (tmp_path / "config.json").write_text(json.dumps({"vision_config": VISION_CONFIG}))
+    return tmp_path
+
+
+def test_448_gives_784_patches_and_196_tokens(checkpoint):
+    shape = tower_shape(checkpoint, 448)
+    assert (shape.grid_h, shape.grid_w) == (28, 28)
+    assert shape.patches == 784
+    # 3 channels x 2 temporal x 16 x 16
+    assert shape.patch_dim == 1536
+    # 784 patches merged 2x2
+    assert shape.merged_tokens == 196
+    assert shape.out_hidden_size == 5120
+
+
+def test_image_size_must_divide_by_the_patch_size(checkpoint):
+    with pytest.raises(ValueError, match="multiple of patch_size"):
+        tower_shape(checkpoint, 450)
+
+
+def test_grid_must_divide_by_the_spatial_merge(checkpoint):
+    """A 3x3 grid of patches cannot be merged 2x2. Exporting it anyway would
+    hand the language model a silently wrong number of image tokens."""
+    with pytest.raises(ValueError, match="spatial_merge_size"):
+        tower_shape(checkpoint, 48)
+
+
+def test_a_checkpoint_without_a_tower_is_refused(tmp_path):
+    (tmp_path / "config.json").write_text(json.dumps({"model_type": "spark2_5"}))
+    with pytest.raises(ValueError, match="no vision_config"):
+        tower_shape(tmp_path, 448)
+
+
+@pytest.mark.parametrize("size,patches,tokens", [(224, 196, 49), (448, 784, 196), (896, 3136, 784)])
+def test_token_count_scales_with_area(checkpoint, size, patches, tokens):
+    shape = tower_shape(checkpoint, size)
+    assert (shape.patches, shape.merged_tokens) == (patches, tokens)

tests/edge/local/external/test_win_cuda_route.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""[W4] The native-Windows vLLM route, and why it is refused.
+
+``win-cuda`` is the first route whose far side is torch rather than
+onnxruntime, and the first that exists to leave WSL rather than to reach a
+device WSL cannot see. Declaring it makes a Windows engine a row in the route
+table instead of a parallel universe.
+
+It is also the first route that **cannot pass the placement gate**, and these
+tests pin that down so nobody later mistakes the refusal for a bug. M3 made a
+measured device attribution a hard admission condition after a VitisAI session
+listed its EP, claimed zero nodes, and returned output bit-identical to the
+CPU's. vLLM emits no per-operation device map, so a ``win-cuda`` worker has
+nothing of that kind to report -- and the rule is about evidence, not about
+whose code it is, so our own engine is refused on it exactly as the NPU is.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from vllm_omni.edge.local.capabilities import FORMAT_ONNX_A16W8
+from vllm_omni.edge.local.external import launch
+from vllm_omni.edge.local.external.stage import (
+    ExternalStage,
+    PlacementRefused,
+    plan_external_stage,
+)
+from vllm_omni.edge.local.manifest import build_graph_artifact
+from vllm_omni.edge.local.plan import REFUSE_EP_PLACEMENT, REFUSE_NO_DEVICE
+
+from .conftest import _device
+
+
+def _artifact(path):
+    return build_graph_artifact(
+        path, fmt=FORMAT_ONNX_A16W8, opset=21, source_model="test",
+        component="stage", exporter="tests.edge.local.external",
+    )
+
+
+def test_win_cuda_is_a_known_route():
+    """It resolves through the same machinery as every other route."""
+    route = launch.resolve(launch.ROUTE_WIN_CUDA)
+    assert route.name == "win-cuda"
+    assert route.ep == "cuda"
+    # It needs a native-Windows interpreter, like the two AMD Windows routes.
+    assert route.is_windows is True
+    # Available or not is environment-dependent; a reason is not.
+    assert route.available or route.reason
+
+
+def test_win_cuda_reports_no_placement_evidence():
+    """The constant is the contract: ``None`` means nobody can look yet.
+
+    Asserted rather than assumed, because the tempting edit -- setting it to
+    1.0 because the worker "obviously" runs on the GPU -- is exactly the
+    reasoning M3 disproved.
+    """
+    assert launch.WIN_CUDA_PLACEMENT_EVIDENCE is None
+
+
+def test_unverified_placement_is_refused_for_our_own_engine(a16w8_graph, fake_route, tmp_path):
+    """A worker that cannot say where it ran is refused, not trusted.
+
+    The fake worker stands in for a ``win-cuda`` one: it runs, it returns
+    correct-looking numbers, and it reports ``fraction_on_target=None`` because
+    vLLM has no per-op device map to derive one from. That is precisely the
+    shape of a silent CPU fallback, so it must not be admitted.
+    """
+    fake_route(load={
+        "fraction_on_target": None,
+        "node_counts": {},
+        "placement_unknown_reason": "vllm emits no per-operation device map",
+    })
+    device = _device("npu:amd", "npu", [FORMAT_ONNX_A16W8], launch.ROUTE_CPU)
+    plan = plan_external_stage(_artifact(a16w8_graph), [device])
+    stage = ExternalStage(plan)
+    with pytest.raises(PlacementRefused) as excinfo:
+        stage.open({"x": np.ones(4, np.float32)}, profile_dir=tmp_path)
+    assert excinfo.value.refusal.code == REFUSE_EP_PLACEMENT
+    assert "not measured" in excinfo.value.refusal.message
+
+
+def test_external_planner_does_not_yet_admit_a_discrete_gpu(a16w8_graph, fake_route):
+    """A documented gap, pinned so the next person does not rediscover it.
+
+    ``plan_external_stage`` filters candidates to ``("gpu_integrated", "npu")``
+    (``stage.py``). Every external route so far reaches a device WSL cannot see
+    *and* that is integrated or an NPU, so the filter was exactly right. A
+    ``win-cuda`` route breaks that coupling: its device is a **discrete** GPU
+    reached out-of-process for an OS reason, not a visibility one.
+
+    This test asserts the current refusal rather than a desired behaviour. When
+    the planner learns about externally-routed discrete GPUs, this test should
+    fail and be rewritten -- that failure is the signal, not a regression.
+    """
+    fake_route()
+    device = _device("cuda:0", "gpu_discrete", [FORMAT_ONNX_A16W8], launch.ROUTE_CPU)
+    plan = plan_external_stage(_artifact(a16w8_graph), [device])
+    assert plan.device is None
+    assert plan.refusals
+    assert plan.refusals[0].code == REFUSE_NO_DEVICE
+    assert "integrated GPU or NPU" in plan.refusals[0].message

tests/edge/local/test_capabilities.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""What this process can run on, as opposed to what is plugged in."""
+
+import pytest
+
+from vllm_omni.edge.hardware_probe import ACCEL_GPU_DISCRETE
+from vllm_omni.edge.local.capabilities import (
+    FORMAT_DENSE,
+    FORMAT_FP8,
+    FORMAT_GPTQ,
+    FORMAT_INT8,
+    FORMAT_NVFP4,
+    MASK_ENV,
+    cuda_weight_formats,
+    enumerate_devices,
+    parse_mask,
+    runs_exported_graphs,
+)
+
+from .conftest import blackwell_laptop, cpu_only_box
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _by_id(devices):
+    return {d.device_id: d for d in devices}
+
+
+# ---------------------------------------------------------- format gating
+def test_blackwell_has_no_int8_path():
+    """sm_120 is inside the SM >= 100 range where cutlass c3x has no int8
+    scaled_mm. This is the single fact that turns a 40-second crash into a
+    30-millisecond refusal, so it gets its own test."""
+    formats = cuda_weight_formats((12, 0))
+    assert FORMAT_INT8 not in formats
+    assert FORMAT_FP8 in formats
+    assert FORMAT_NVFP4 in formats
+
+
+def test_ampere_keeps_int8_and_has_no_nvfp4():
+    formats = cuda_weight_formats((8, 6))
+    assert FORMAT_INT8 in formats
+    assert FORMAT_NVFP4 not in formats
+    assert FORMAT_FP8 not in formats
+
+
+def test_dense_runs_on_every_cuda_generation():
+    for capability in ((7, 5), (8, 0), (9, 0), (12, 0)):
+        assert FORMAT_DENSE in cuda_weight_formats(capability)
+
+
+# ------------------------------------------------------------ runnability
+def test_a_cuda_wheel_cannot_drive_the_cpu_and_says_so():
+    devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda"))
+    assert devices["cuda:0"].runnable
+    assert not devices["cpu"].runnable
+    assert "omni-cpu" in devices["cpu"].reason
+
+
+def test_a_cpu_wheel_cannot_drive_the_card_and_says_so():
+    devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cpu"))
+    assert devices["cpu"].runnable
+    assert not devices["cuda:0"].runnable
+    assert "omni-cuda" in devices["cuda:0"].reason
+
+
+def test_the_igpu_and_npu_are_reported_not_dropped():
+    """'The machine has an NPU' and 'a stage can run on it' are different
+    claims; the plan has to be able to show which one it relied on.
+
+    Both are now reachable (M3), through a worker in another interpreter, so
+    the invariant is no longer "never runnable" -- it is that the record always
+    says which of the two claims it is making, and never leaves a caller to
+    guess.
+    """
+    devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda"))
+    assert "npu:amd" in devices
+    assert "gpu_integrated:amd" in devices
+    for key in ("npu:amd", "gpu_integrated:amd"):
+        device = devices[key]
+        assert device.evidence in ("B", "D")
+        if device.runnable:
+            assert device.backend, "a runnable device must name the backend that drives it"
+            assert device.extra["worker_route"], "and the route it was resolved through"
+        else:
+            assert device.reason, "and an unrunnable one must say why"
+            assert device.backend is None
+        # Either way these execute exported graphs, never a checkpoint.
+        assert runs_exported_graphs(device)
+
+
+def test_the_amd_devices_refuse_with_the_missing_interpreter_named(monkeypatch):
+    """A device that works and a venv that is not installed are different
+    problems; the reason has to distinguish them well enough to act on."""
+    for route in ("ORT_VITISAI", "ORT_DML", "TORCH_DML"):
+        monkeypatch.setenv(f"VLLM_OMNI_EXTERNAL_PYTHON_{route}", "/nonexistent/python")
+    devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda"))
+    for key in ("npu:amd", "gpu_integrated:amd"):
+        device = devices[key]
+        assert not device.runnable and device.backend is None
+        assert "/nonexistent/python" in device.reason or "native-Windows" in device.reason
+
+
+def test_the_amd_devices_get_no_memory_budget_of_their_own():
+    """Section 6: the iGPU and the NPU allocate out of the CPU's RAM.
+
+    Carrying the pool size on their rows would let a caller add three copies of
+    the same memory together, so the number lives on the host row only and the
+    stage planner reads it from there.
+    """
+    devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cpu"))
+    for key in ("npu:amd", "gpu_integrated:amd"):
+        assert devices[key].memory_bytes == 0
+        assert devices[key].memory_pool == "host_ram"
+        assert devices[key].extra["host_pool_bytes"] > 0
+
+
+def test_a_cpu_only_box_reports_no_cuda_row_at_all():
+    """A masked GPU and an absent GPU are different situations and must not
+    produce the same record."""
+    devices = _by_id(enumerate_devices(cpu_only_box(), platform_device_type="cpu"))
+    assert "cuda:0" not in devices
+    assert devices["cpu"].runnable
+
+
+# ------------------------------------------------------------------ masking
+def test_masking_disables_without_hiding():
+    devices = _by_id(enumerate_devices(
+        blackwell_laptop(), mask=frozenset({ACCEL_GPU_DISCRETE}), platform_device_type="cuda"
+    ))
+    gpu = devices["cuda:0"]
+    assert gpu.masked and not gpu.runnable
+    # Still enumerated, with its real memory: a mask is a policy, not an
+    # amnesia, and the record has to show the card was there.
+    assert gpu.memory_bytes > 0
+    assert "mask" in gpu.reason
+
+
+def test_mask_is_read_from_the_environment(monkeypatch):
+    monkeypatch.setenv(MASK_ENV, "gpu_discrete, cpu")
+    assert parse_mask() == frozenset({"gpu_discrete", "cpu"})
+    monkeypatch.delenv(MASK_ENV)
+    assert parse_mask() == frozenset()
+
+
+def test_an_explicit_mask_beats_the_environment(monkeypatch):
+    monkeypatch.setenv(MASK_ENV, "gpu_discrete")
+    assert parse_mask("") == frozenset()
+
+
+# ---------------------------------------------------------------- pooling
+def test_cpu_and_accelerators_share_one_host_pool():
+    """Section 6 forbids giving the iGPU and the NPU each their own RAM."""
+    devices = enumerate_devices(blackwell_laptop(), platform_device_type="cpu")
+    host_pool = [d for d in devices if d.memory_pool == "host_ram"]
+    assert {d.device_id for d in host_pool} >= {"cpu", "npu:amd", "gpu_integrated:amd"}
+    # Only the CPU row carries the bytes; the others are 0 so nothing can add
+    # three copies of the same RAM together.
+    assert sum(d.memory_bytes for d in host_pool) == next(
+        d.memory_bytes for d in host_pool if d.device_id == "cpu"
+    )
+
+
+def test_the_discrete_gpu_has_its_own_pool():
+    devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cuda"))
+    assert devices["cuda:0"].memory_pool == "vram"
+    assert devices["cpu"].memory_pool == "host_ram"
+
+
+def test_cpu_weight_formats_include_int8_on_x86():
+    devices = _by_id(enumerate_devices(blackwell_laptop(), platform_device_type="cpu"))
+    assert FORMAT_INT8 in devices["cpu"].weight_formats
+    assert FORMAT_GPTQ in devices["cpu"].weight_formats

tests/edge/local/test_e2e_local_text.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""End to end: the M0 acceptance criteria, executed against a real checkpoint.
+
+This is the test that says the local text mode works, so it tests the thing the
+roadmap actually asks for rather than a proxy for it:
+
+* both startups -- device-visible and capability-masked -- reach a decision,
+  and a refusal is a pass, because "选择兼容方案或明确拒绝" is an either/or;
+* the same fixed set of twelve prompts, each generating at least 128 tokens;
+* TTFT and decode rate recorded per request;
+* the budgeted peak is an upper bound on the measured one, and nothing OOMs;
+* cancellation releases, and no event from the retired epoch is delivered.
+
+It needs real weights and a real backend, so it skips rather than fails when
+there is no checkpoint. Which one it runs is whatever this interpreter can
+drive: the CUDA venv gets the bf16 build, the CPU venv gets the int8 one --
+which is itself the point, since the int8 build is the one the CUDA plan
+refuses. Override with ``VLLM_OMNI_LOCAL_TEST_MODEL``.
+
+Roughly a minute on the reference laptop's GPU and three on its CPU. Marked
+``local_model`` because it is not a unit test and should not run in a PR gate.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+import pytest_asyncio
+
+from vllm_omni.edge.local import (
+    LocalTextEngine,
+    PlanNotAdmitted,
+    acceptance_prompts,
+    plan_text_session,
+)
+from vllm_omni.edge.local.capabilities import FORMAT_INT8
+
+# One event loop for the whole module. The engine fixture is module-scoped --
+# loading 7.7 GB per test would be absurd -- and the asyncio primitives inside
+# AsyncOmni belong to whichever loop created them. Leaving the tests on
+# pytest-asyncio's default per-function loop means a cancel awaits a task whose
+# backend cleanup lives on a different, no-longer-running loop, and it hangs
+# there rather than failing.
+pytestmark = [pytest.mark.local_model, pytest.mark.asyncio(loop_scope="module")]
+
+MIN_TOKENS = 128
+MAX_MODEL_LEN = 4096
+
+# Checkpoints under the repo's models/ directory, by the platform that can run
+# them. The int8 build is deliberately the CPU entry: Blackwell refuses it.
+_CANDIDATES = {
+    "cuda": ("Spark-X2.5-4B", "Spark-X2.5-1.7B-int8"),
+    "cpu": ("Spark-X2.5-1.7B-int8", "Spark-X2.5-4B"),
+}
+
+
+def _models_root() -> Path:
+    # tests/edge/local/ -> vllm-omni/ -> the workspace that holds models/
+    return Path(__file__).resolve().parents[3].parent / "models"
+
+
+def _platform() -> str:
+    from vllm.platforms import current_platform
+
+    return str(getattr(current_platform, "device_type", "") or "")
+
+
+def _find_model() -> str:
+    override = os.environ.get("VLLM_OMNI_LOCAL_TEST_MODEL")
+    if override:
+        return override
+    root = _models_root()
+    for name in _CANDIDATES.get(_platform(), ()):
+        if (root / name / "config.json").is_file():
+            return str(root / name)
+    pytest.skip(
+        f"no Spark checkpoint for platform {_platform()!r} under {root}; "
+        "set VLLM_OMNI_LOCAL_TEST_MODEL to run this"
+    )
+
+
+@pytest.fixture(scope="module")
+def model_dir() -> str:
+    return _find_model()
+
+
+@pytest.fixture(scope="module")
+def plan(model_dir: str):
+    p = plan_text_session(model_dir, max_model_len=MAX_MODEL_LEN, max_num_seqs=1)
+    if not p.admitted:
+        pytest.skip(f"this host cannot run {model_dir}:\n{p.summary()}")
+    return p
+
+
+@pytest_asyncio.fixture(scope="module", loop_scope="module")
+async def engine(plan):
+    eng = LocalTextEngine(plan)
+    await eng.start()
+    try:
+        yield eng
+    finally:
+        await eng.close()
+
+
+# ------------------------------------------------------------------ planning
+def test_the_masked_startup_reaches_a_decision(model_dir):
+    """The no-NVIDIA class must produce a plan or an explicit refusal -- never
+    a crash and never a silent substitution."""
+    masked = plan_text_session(
+        model_dir, max_model_len=MAX_MODEL_LEN, mask=frozenset({"gpu_discrete"})
+    )
+    if masked.admitted:
+        assert masked.selected.memory_pool == "host_ram"
+        assert masked.selected.kind != "gpu_discrete"
+    else:
+        assert masked.refusals
+        assert all(r.message and r.remedy for r in masked.refusals)
+        assert any("mask" in r.message for r in masked.refusals) or _platform() != "cuda"
+
+
+def test_a_refused_plan_cannot_be_started(model_dir):
+    refused = plan_text_session(
+        model_dir, max_model_len=MAX_MODEL_LEN, mask=frozenset({"gpu_discrete", "cpu"})
+    )
+    assert not refused.admitted
+    with pytest.raises(PlanNotAdmitted):
+        LocalTextEngine(refused)
+
+
+@pytest.mark.skipif(
+    not (_models_root() / "Spark-X2.5-1.7B-int8" / "config.json").is_file(),
+    reason="the int8 checkpoint is not present",
+)
+def test_int8_is_refused_on_cuda_and_admitted_on_cpu():
+    """The refusal and its remedy, checked against the real checkpoint on
+    whichever backend this interpreter has."""
+    int8 = str(_models_root() / "Spark-X2.5-1.7B-int8")
+    p = plan_text_session(int8, max_model_len=MAX_MODEL_LEN)
+    assert p.manifest.weight_format == FORMAT_INT8
+    if _platform() == "cuda":
+        assert not p.admitted
+        formats = [r for r in p.refusals if r.code == "weight_format_unsupported"]
+        assert formats and "omni-cpu" in formats[0].remedy
+    else:
+        assert p.admitted
+        assert p.selected.device_id == "cpu"
+
+
+# ------------------------------------------------------------------- loading
+async def test_the_engine_loads_and_reports_where_it_landed(engine):
+    assert engine.load_seconds and engine.load_seconds > 0
+    placement = engine.report_placement()
+    configured = placement["configured"]
+    assert "error" not in configured
+    # Read back from the stage's own config, not echoed from the plan.
+    assert configured["device_type"] == _platform()
+    assert configured["max_model_len"] == MAX_MODEL_LEN
+    assert configured["enforce_eager"] is True
+
+
+async def test_the_kv_pool_is_the_size_the_plan_asked_for(engine, plan):
+    """An absolute budget, honoured -- not whatever was left inside a
+    utilisation fraction."""
+    configured = engine.report_placement()["configured"]
+    assert configured["kv_cache_memory_bytes"] == plan.kv_budget.bytes_total
+
+
+# ---------------------------------------------------------- the twelve prompts
+async def test_twelve_prompts_each_reach_128_tokens(engine):
+    prompts = acceptance_prompts()
+    assert len(prompts) == 12
+
+    records = []
+    for name, text in prompts:
+        session = engine.open_session()
+        request_id, stream = await engine.submit(
+            session, text, max_tokens=MIN_TOKENS, temperature=0.0, ignore_eos=True
+        )
+        seqs = []
+        async for event in stream:
+            if event.kind == "token":
+                seqs.append(event.seq)
+            assert event.epoch == session.epoch
+            assert event.error is None, f"{name}: {event.error}"
+        record = engine.records[request_id]
+        engine.close_session(session.session_id)
+
+        assert record.finished, f"{name} did not finish"
+        assert record.output_tokens >= MIN_TOKENS, f"{name} produced {record.output_tokens}"
+        assert record.ttft_s is not None and record.ttft_s > 0
+        assert record.decode_tok_per_s is not None and record.decode_tok_per_s > 0
+        # Incremental, in order, no repeats: seq is what detects a lost chunk.
+        assert seqs == sorted(seqs)
+        assert len(seqs) == len(set(seqs))
+        records.append(record)
+
+    # The long prompt crosses the 512 sliding window many times; if window
+    # rollover were broken it would be broken here and not on the short ones.
+    long_record = next(r for r in records if r.prompt_tokens and r.prompt_tokens > 2000)
+    assert long_record.output_tokens >= MIN_TOKENS
+
+
+async def test_the_budget_bounded_the_measured_peak(engine, plan):
+    """Admission's whole claim: the plan's number is an upper bound, proven
+    before anything allocated."""
+    measured = engine.measured_peak()
+    assert measured["budgeted_peak_bytes"] == plan.peak_bytes
+    attributable = measured["attributable_peak_bytes"]
+    assert attributable is not None and attributable > 0
+    assert attributable <= measured["budgeted_peak_bytes"], (
+        f"measured {attributable / 2**30:.2f} GiB exceeded the budgeted "
+        f"{measured['budgeted_peak_bytes'] / 2**30:.2f} GiB"
+    )
+
+
+async def test_usage_separates_what_was_measured_from_what_was_estimated(engine):
+    usage = engine.report_usage()
+    by_purpose = {r["purpose"]: r for r in usage["reservations"]}
+    assert by_purpose["weights"]["evidence"] == "D"
+    assert by_purpose["kv_cache"]["evidence"] == "E"
+    # Nothing gets an actual it cannot support: only the weights line is
+    # bounded by a counter on this host.
+    bounded = [p for p, r in by_purpose.items() if r["actual_upper_bound_bytes"]]
+    assert bounded == ["weights"]
+
+
+# ------------------------------------------------------------- cancellation
+async def test_cancel_releases_and_delivers_nothing_from_the_retired_epoch(engine):
+    session = engine.open_session()
+    request_id, stream = await engine.submit(
+        session, acceptance_prompts()[0][1], max_tokens=4096, ignore_eos=True
+    )
+    seen = 0
+    async for event in stream:
+        if event.kind == "token":
+            seen += 1
+        if seen >= 8:
+            break
+    assert seen >= 8
+
+    report = await engine.cancel(request_id)
+    assert report["aborted_in_backend"]
+    assert report["producer_stopped"], "the producer task did not unwind"
+    assert report["inflight_after"] == 0
+    # The session's state is retired, not reused.
+    assert report["session_epoch_now"] == session.epoch + 1
+
+    leaked = []
+    while True:
+        event = await stream.get()
+        if event is None:
+            break
+        leaked.append(event)
+    assert leaked == [], f"{len(leaked)} events escaped the retired epoch"
+
+
+async def test_a_retired_handle_cannot_be_resubmitted(engine):
+    session = engine.open_session()
+    request_id, stream = await engine.submit(session, "Hello", max_tokens=32)
+    async for _ in stream:
+        break
+    await engine.cancel(request_id)
+    with pytest.raises(RuntimeError, match="stale"):
+        await engine.submit(session, "Hello again", max_tokens=8)
+
+
+async def test_the_engine_still_generates_after_a_cancel(engine):
+    """A cancelled request must not poison the engine for the next one."""
+    session = engine.open_session()
+    request_id, stream = await engine.submit(
+        session, "The capital of Australia is", max_tokens=16, ignore_eos=True
+    )
+    async for _ in stream:
+        pass
+    record = engine.records[request_id]
+    assert record.finished and record.output_tokens >= 16

tests/edge/local/test_gpu_attribution.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The placement report must measure per-process GPU attribution, not assume it.
+
+[edge-infer W4] The W1 acceptance on native Windows carried a hardcoded "under
+WSL2" explanation for an empty NVML process list. These tests pin the repaired
+behaviour: the reason names the host it actually ran on, and a successful
+attribution (NVML lists one of our PIDs) is reported as data, not a reason.
+"""
+
+from __future__ import annotations
+
+import builtins
+import sys
+import types
+
+import pytest
+
+from vllm_omni.edge.local import engine
+
+
+def test_host_kind_is_one_of_the_three_hosts():
+    assert engine._host_kind() in {"windows", "wsl2", "linux"}
+
+
+def test_reason_names_the_host_when_nvml_is_unavailable(monkeypatch):
+    real_import = builtins.__import__
+
+    def no_pynvml(name, *args, **kwargs):
+        if name == "pynvml":
+            raise ImportError("pynvml is not installed in this venv")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", no_pynvml)
+    out = engine.gpu_attribution_for_tree()
+    assert out["per_process_gpu_attribution"] is None
+    assert engine._host_kind() in out["per_process_gpu_attribution_reason"]
+    assert "WSL2" not in out["per_process_gpu_attribution_reason"] or engine._host_kind() == "wsl2"
+
+
+def _fake_pynvml(pids):
+    m = types.ModuleType("pynvml")
+    m.nvmlInit = lambda: None
+    m.nvmlShutdown = lambda: None
+    m.nvmlDeviceGetHandleByIndex = lambda i: object()
+    m.nvmlDeviceGetComputeRunningProcesses = lambda h: [
+        types.SimpleNamespace(pid=p, usedGpuMemory=1234) for p in pids
+    ]
+    return m
+
+
+def test_attribution_is_data_when_nvml_lists_our_pid(monkeypatch):
+    import os
+
+    monkeypatch.setitem(sys.modules, "pynvml", _fake_pynvml([os.getpid(), 1]))
+    out = engine.gpu_attribution_for_tree()
+    assert out["per_process_gpu_attribution"] == [
+        {"pid": os.getpid(), "used_gpu_memory": 1234}
+    ]
+    assert out["per_process_gpu_attribution_reason"] is None
+
+
+def test_reason_states_the_measured_count_when_none_are_ours(monkeypatch):
+    monkeypatch.setitem(sys.modules, "pynvml", _fake_pynvml([1, 2]))
+    out = engine.gpu_attribution_for_tree()
+    assert out["per_process_gpu_attribution"] is None
+    reason = out["per_process_gpu_attribution_reason"]
+    assert reason.startswith("nvmlDeviceGetComputeRunningProcesses returned 2 process(es)")
+    host = engine._host_kind()
+    expected = {"wsl2": "WSL2", "windows": "WDDM", "linux": "no compute process"}[host]
+    assert expected in reason

tests/edge/local/test_manifest.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Identifying a checkpoint well enough to tell it from its siblings."""
+
+import json
+
+import pytest
+
+from vllm_omni.edge.local.capabilities import FORMAT_DENSE, FORMAT_FP8, FORMAT_INT8
+from vllm_omni.edge.local.manifest import build_manifest, runtime_versions
+
+from .conftest import FP8_QUANT_CONFIG, INT8_QUANT_CONFIG, write_checkpoint
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_quantization_format_is_read_from_the_config(tmp_path):
+    assert build_manifest(write_checkpoint(tmp_path / "d")).weight_format == FORMAT_DENSE
+    assert build_manifest(
+        write_checkpoint(tmp_path / "i", quantization=INT8_QUANT_CONFIG)
+    ).weight_format == FORMAT_INT8
+    assert build_manifest(
+        write_checkpoint(tmp_path / "f", quantization=FP8_QUANT_CONFIG)
+    ).weight_format == FORMAT_FP8
+
+
+def test_compressed_tensors_int8_and_fp8_are_distinguished_by_format_not_method(tmp_path):
+    """Both say ``compressed-tensors``; only ``format`` separates the one
+    Blackwell can execute from the one it cannot."""
+    int8 = build_manifest(write_checkpoint(tmp_path / "i", quantization=INT8_QUANT_CONFIG))
+    fp8 = build_manifest(write_checkpoint(tmp_path / "f", quantization=FP8_QUANT_CONFIG))
+    assert int8.quantization["quant_method"] == fp8.quantization["quant_method"]
+    assert int8.weight_format != fp8.weight_format
+
+
+def test_two_quantizations_of_one_model_get_different_ids(tmp_path):
+    dense = build_manifest(write_checkpoint(tmp_path / "d"))
+    int8 = build_manifest(write_checkpoint(tmp_path / "i", quantization=INT8_QUANT_CONFIG))
+    assert dense.artifact_id != int8.artifact_id
+
+
+def test_the_id_is_stable_across_reads(tmp_path):
+    path = write_checkpoint(tmp_path / "d")
+    assert build_manifest(path).artifact_id == build_manifest(path).artifact_id
+
+
+def test_a_changed_tokenizer_changes_the_id(tmp_path):
+    """The tokenizer decides what the model is asked, and moves independently
+    of the weights."""
+    path = write_checkpoint(tmp_path / "d")
+    before = build_manifest(path).artifact_id
+    (path / "chat_template.jinja").write_text("{{ different }}", encoding="utf-8")
+    assert build_manifest(path).artifact_id != before
+
+
+def test_weights_are_not_hashed_unless_asked(tmp_path):
+    path = write_checkpoint(tmp_path / "d")
+    assert build_manifest(path).weight_sha256 is None
+    assert any("digest_weights" in n for n in build_manifest(path).notes)
+    assert build_manifest(path, digest_weights=True).weight_sha256 is not None
+
+
+def test_weight_bytes_and_shards_are_recorded(tmp_path):
+    path = write_checkpoint(tmp_path / "d", weight_bytes=8192, shards=4)
+    manifest = build_manifest(path)
+    assert len(manifest.weight_files) == 4
+    assert manifest.weight_bytes == 8192
+
+
+def test_a_directory_without_a_config_is_an_error(tmp_path):
+    (tmp_path / "empty").mkdir()
+    with pytest.raises(FileNotFoundError):
+        build_manifest(tmp_path / "empty")
+
+
+def test_an_hf_snapshot_path_yields_its_revision(tmp_path):
+    root = tmp_path / "models--org--name" / "snapshots" / "deadbeef1234"
+    write_checkpoint(root)
+    assert build_manifest(root).revision == "deadbeef1234"
+
+
+def test_a_local_checkpoint_reports_no_revision_rather_than_guessing(tmp_path):
+    assert build_manifest(write_checkpoint(tmp_path / "d")).revision is None
+
+
+def test_the_runtime_records_both_the_wheel_and_the_checkout():
+    versions = runtime_versions()
+    assert versions.python
+    # The import path, not just the version string: an editable install and a
+    # site-packages copy report the same version and are different code.
+    assert versions.vllm_omni_file.endswith("__init__.py")
+    assert isinstance(versions.checkouts, dict)
+
+
+def test_the_manifest_is_json_serialisable(tmp_path):
+    manifest = build_manifest(write_checkpoint(tmp_path / "d"))
+    back = json.loads(json.dumps(manifest.to_dict(), default=str))
+    assert back["artifact_id"] == manifest.artifact_id
+    assert back["architectures"] == ["Spark2_5ForCausalLM"]

tests/edge/local/test_parity.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference consistency: does the Omni stage path change the model's answer?
+
+The comparison logic is unit-tested here; the run that actually loads two
+engines is at the bottom and skips without a checkpoint. It is separated from
+``test_e2e_local_text`` because it cannot share that module's engine fixture --
+it needs the reference engine loaded and released *before* the Omni one, and
+two copies of a 7.7 GB checkpoint do not fit the budget this milestone is about.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.edge.local.parity import compare_ids, summarize
+from vllm_omni.edge.local.plan import plan_text_session
+
+MAX_MODEL_LEN = 4096
+
+
+# ------------------------------------------------------------------- unit
+@pytest.mark.core_model
+@pytest.mark.cpu
+class TestCompare:
+    def test_identical_sequences_are_exact(self):
+        assert compare_ids([1, 2, 3], [1, 2, 3]) == (True, None, 3)
+
+    def test_a_divergence_reports_where(self):
+        assert compare_ids([1, 2, 3, 4], [1, 2, 9, 4]) == (False, 2, 2)
+
+    def test_a_short_engine_sequence_is_not_exact(self):
+        """A truncation is a divergence at the end, not agreement."""
+        assert compare_ids([1, 2, 3], [1, 2]) == (False, 2, 2)
+
+    def test_a_long_engine_sequence_is_not_exact(self):
+        assert compare_ids([1, 2], [1, 2, 3]) == (False, 2, 2)
+
+    def test_divergence_at_the_first_token(self):
+        assert compare_ids([5, 6], [7, 6]) == (False, 0, 0)
+
+    def test_empty_sequences_agree_vacuously(self):
+        assert compare_ids([], []) == (True, None, 0)
+
+    def test_summary_counts_and_reports_the_worst_prefix(self):
+        from vllm_omni.edge.local.parity import PromptParity
+
+        rows = [
+            PromptParity("a", 4, 8, 8, True, None, 8),
+            PromptParity("b", 4, 8, 8, False, 5, 5),
+            PromptParity("c", 4, 8, 8, False, 1, 1),
+        ]
+        stats = summarize(rows)
+        assert stats["exact"] == 1
+        assert stats["prompts"] == 3
+        assert stats["min_agreement_prefix"] == 1
+        assert stats["mean_agreement_prefix"] == 3.0
+
+    def test_an_all_exact_summary_has_no_divergence_stats(self):
+        from vllm_omni.edge.local.parity import PromptParity
+
+        stats = summarize([PromptParity("a", 4, 8, 8, True, None, 8)])
+        assert stats["exact_fraction"] == 1.0
+        assert stats["min_agreement_prefix"] is None
+
+
+# -------------------------------------------------------------------- e2e
+def _models_root() -> Path:
+    return Path(__file__).resolve().parents[3].parent / "models"
+
+
+def _platform() -> str:
+    from vllm.platforms import current_platform
+
+    return str(getattr(current_platform, "device_type", "") or "")
+
+
+def _find_model() -> str:
+    override = os.environ.get("VLLM_OMNI_LOCAL_TEST_MODEL")
+    if override:
+        return override
+    candidates = ("Spark-X2.5-4B",) if _platform() == "cuda" else ("Spark-X2.5-1.7B-int8",)
+    for name in candidates:
+        if (_models_root() / name / "config.json").is_file():
+            return str(_models_root() / name)
+    pytest.skip(f"no Spark checkpoint for platform {_platform()!r}")
+
+
+@pytest.mark.local_model
+@pytest.mark.asyncio(loop_scope="module")
+async def test_the_omni_path_decodes_what_plain_vllm_decodes():
+    """The hard acceptance condition: same checkpoint, same limits, same
+    greedy decode -- the only difference is the path through the Omni stage.
+
+    Four prompts rather than the full twelve: this loads two engines and the
+    point is to detect a systematic difference in the stage path, which four
+    prompts of different shapes (short chat, tool-formatted, long, plain) show
+    as well as twelve. The full set runs from the CLI's ``parity`` command.
+    """
+    from vllm_omni.edge.local.cli import _engine_token_ids
+    from vllm_omni.edge.local.parity import reference_token_ids
+    from vllm_omni.edge.local.prompts import acceptance_prompts
+
+    model_dir = _find_model()
+    plan = plan_text_session(model_dir, max_model_len=MAX_MODEL_LEN, max_num_seqs=1)
+    if not plan.admitted:
+        pytest.skip(f"this host cannot run {model_dir}")
+
+    all_prompts = dict(acceptance_prompts())
+    prompts = [(n, all_prompts[n]) for n in ("chat_short", "chat_tools", "chat_long", "plain_0")]
+    max_tokens = 64
+
+    reference = reference_token_ids(
+        model_dir, prompts, engine_kwargs=plan.engine_kwargs, max_tokens=max_tokens
+    )
+    engine = await _engine_token_ids(plan, prompts, max_tokens)
+
+    mismatches = []
+    for name, _ in prompts:
+        ref = reference[name]["token_ids"]
+        eng = engine[name]["token_ids"]
+        assert len(ref) == max_tokens, f"{name}: reference produced {len(ref)}"
+        exact, first, prefix = compare_ids(ref, eng)
+        if not exact:
+            mismatches.append(f"{name}: diverged at {first}/{len(ref)} (prefix {prefix})")
+    assert not mismatches, "the Omni stage path changed the decode:\n" + "\n".join(mismatches)

tests/edge/local/test_plan.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Admission: say yes with a budget, or no with a reason and a remedy."""
+
+import pytest
+
+from vllm_omni.edge.local.plan import (
+    GiB,
+    REFUSE_CAPACITY,
+    REFUSE_FORMAT,
+    REFUSE_NO_DEVICE,
+    plan_text_session,
+)
+
+from .conftest import FP8_QUANT_CONFIG, blackwell_laptop, cpu_only_box, write_checkpoint
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _plan(model_dir, profile=None, platform="cuda", **kw):
+    return plan_text_session(
+        str(model_dir),
+        profile=profile if profile is not None else blackwell_laptop(),
+        platform_device_type=platform,
+        **kw,
+    )
+
+
+def _codes(plan):
+    return {r.code for r in plan.refusals}
+
+
+# ------------------------------------------------------------------ the yes
+def test_a_dense_checkpoint_is_admitted_to_the_gpu(dense_checkpoint):
+    plan = _plan(dense_checkpoint)
+    assert plan.admitted
+    assert plan.selected.device_id == "cuda:0"
+    assert plan.backend == "vllm:cuda"
+
+
+def test_an_admitted_plan_prices_every_line_of_section_6_2(dense_checkpoint):
+    """The budget is not 'the weights'. It is the sum the design specifies,
+    and each term has to be visible so a reviewer can argue with one."""
+    plan = _plan(dense_checkpoint)
+    purposes = {r.purpose for r in plan.reservations}
+    assert purposes == {
+        "weights", "kv_cache", "activations", "backend_workspace",
+        "load_transient", "external_reserve", "safety_margin",
+    }
+    assert all(r.pool == "vram" for r in plan.reservations)
+    assert all(r.evidence in {"D", "E"} for r in plan.reservations)
+    # The weights line is read off the files, not estimated.
+    weights = next(r for r in plan.reservations if r.purpose == "weights")
+    assert weights.evidence == "D"
+    assert weights.bytes == plan.manifest.weight_bytes
+
+
+def test_graph_capture_is_charged_only_when_graphs_are_on(dense_checkpoint):
+    eager = _plan(dense_checkpoint, enforce_eager=True)
+    graphs = _plan(dense_checkpoint, enforce_eager=False)
+    assert "graph_capture" not in {r.purpose for r in eager.reservations}
+    assert "graph_capture" in {r.purpose for r in graphs.reservations}
+    assert graphs.peak_bytes > eager.peak_bytes
+
+
+def test_the_peak_is_a_max_not_a_sum(dense_checkpoint):
+    """The loader's transient is gone before the KV pool is sized, so adding
+    both would refuse plans that fit."""
+    plan = _plan(dense_checkpoint)
+    everything = sum(r.bytes for r in plan.reservations)
+    assert plan.peak_bytes <= everything
+
+
+def test_kv_reservation_tracks_the_context(dense_checkpoint):
+    short = _plan(dense_checkpoint, max_model_len=1024)
+    long = _plan(dense_checkpoint, max_model_len=8192)
+    kv_short = next(r for r in short.reservations if r.purpose == "kv_cache").bytes
+    kv_long = next(r for r in long.reservations if r.purpose == "kv_cache").bytes
+    assert kv_long > kv_short
+
+
+def test_sliding_layers_are_not_charged_full_length(dense_checkpoint):
+    """21 of 28 layers hold a 512 window, not the whole context. Charging them
+    all at full length is the arithmetic the kv_budget module exists to fix."""
+    plan = _plan(dense_checkpoint, max_model_len=8192)
+    assert plan.kv_budget.sliding_layers == 21
+    assert plan.kv_budget.full_layers == 7
+    assert plan.kv_budget.bytes_total_hybrid < plan.kv_budget.bytes_total
+
+
+# ------------------------------------------------------------------- the no
+def test_int8_is_refused_on_blackwell_although_it_fits(int8_checkpoint):
+    """4 MiB of weights against 24 GiB of VRAM: every capacity check passes and
+    the answer is still no."""
+    plan = _plan(int8_checkpoint)
+    assert not plan.admitted
+    assert REFUSE_FORMAT in _codes(plan)
+    refusal = next(r for r in plan.refusals if r.code == REFUSE_FORMAT)
+    assert refusal.device_id == "cuda:0"
+    assert "int8" in refusal.message
+    assert "omni-cpu" in refusal.remedy
+
+
+def test_the_same_int8_checkpoint_is_admitted_on_the_cpu(int8_checkpoint):
+    """The remedy the refusal names has to actually be a plan."""
+    plan = _plan(int8_checkpoint, profile=cpu_only_box(), platform="cpu")
+    assert plan.admitted
+    assert plan.selected.device_id == "cpu"
+    assert plan.backend == "vllm:cpu"
+
+
+def test_fp8_is_admitted_where_int8_is_not(tmp_path):
+    fp8 = write_checkpoint(tmp_path / "fp8", quantization=FP8_QUANT_CONFIG)
+    assert _plan(fp8).admitted
+
+
+def test_a_checkpoint_too_large_for_the_pool_is_refused_with_a_number(tmp_path):
+    huge = write_checkpoint(tmp_path / "huge", weight_bytes=40 * GiB, shards=8)
+    plan = _plan(huge)
+    assert not plan.admitted
+    assert REFUSE_CAPACITY in _codes(plan)
+    refusal = next(r for r in plan.refusals if r.code == REFUSE_CAPACITY)
+    assert "GiB" in refusal.message
+
+
+def test_a_capacity_refusal_never_silently_shortens_the_context(tmp_path):
+    """Trimming max_model_len is offered as a remedy the operator applies, not
+    something the planner does on its behalf."""
+    # Sized so the overrun is smaller than the KV cache it asked for: that is
+    # the case where a shorter context would in fact fit, which is exactly when
+    # an engine that "helpfully" truncates would do damage.
+    big = write_checkpoint(tmp_path / "big", weight_bytes=16 * GiB, shards=8)
+    plan = _plan(big, max_model_len=131072)
+    assert not plan.admitted
+    refusal = next(r for r in plan.refusals if r.code == REFUSE_CAPACITY)
+    assert "max-model-len" in refusal.remedy
+    assert "forbids" in refusal.remedy
+    # And the plan still reports the context that was asked for, not a
+    # quietly reduced one.
+    assert plan.request_limits["max_model_len"] == 131072
+
+
+def test_masking_the_gpu_leaves_nothing_runnable_in_a_cuda_process(dense_checkpoint):
+    """The capability-masked startup of the acceptance run: a refusal here is
+    the correct outcome, not a failure."""
+    plan = _plan(dense_checkpoint, mask=frozenset({"gpu_discrete"}))
+    assert not plan.admitted
+    assert REFUSE_NO_DEVICE in _codes(plan)
+    assert any("mask" in r.message for r in plan.refusals)
+
+
+def test_every_unusable_device_is_explained(dense_checkpoint):
+    plan = _plan(dense_checkpoint)
+    explained = {r.device_id for r in plan.refusals}
+    assert {"cpu", "npu:amd", "gpu_integrated:amd"} <= explained
+    assert all(r.message and r.remedy for r in plan.refusals)
+
+
+# ------------------------------------------------------------ engine kwargs
+def test_the_memory_fraction_is_derived_not_defaulted(dense_checkpoint):
+    """vLLM reads this field as host RAM on the CPU backend, so leaving it at
+    its default there asks for 92% of the machine. It is always derived."""
+    for platform, profile in (("cuda", blackwell_laptop()), ("cpu", cpu_only_box())):
+        plan = _plan(dense_checkpoint, profile=profile, platform=platform)
+        assert plan.admitted
+        fraction = plan.engine_kwargs["gpu_memory_utilization"]
+        assert 0 < fraction <= 0.95
+        assert fraction >= plan.peak_bytes / plan.selected.memory_bytes * 0.5
+
+
+def test_the_kv_pool_is_absolute_not_leftover(dense_checkpoint):
+    plan = _plan(dense_checkpoint)
+    assert plan.engine_kwargs["kv_cache_memory_bytes"] == plan.kv_budget.bytes_total
+
+
+def test_request_limits_reach_the_backend(dense_checkpoint):
+    plan = _plan(dense_checkpoint, max_model_len=2048, max_num_seqs=2)
+    assert plan.engine_kwargs["max_model_len"] == 2048
+    assert plan.engine_kwargs["max_num_seqs"] == 2
+    assert plan.engine_kwargs["enforce_eager"] is True
+
+
+# ---------------------------------------------------------------- the stage
+def test_the_plan_declares_exactly_one_text_stage(dense_checkpoint):
+    plan = _plan(dense_checkpoint)
+    assert plan.stage["stage_id"] == 0
+    assert plan.stage["execution_type"] == "LLM_AR"
+    assert plan.stage["final_output_type"] == "text"
+    assert plan.stage["input_sources"] == []
+
+
+def test_the_record_round_trips_to_json(dense_checkpoint):
+    import json
+
+    plan = _plan(dense_checkpoint)
+    text = json.dumps(plan.to_dict(), default=str)
+    back = json.loads(text)
+    assert back["admitted"]
+    assert back["selected_device"]["device_id"] == "cuda:0"
+    assert back["manifest"]["artifact_id"] == plan.manifest.artifact_id

tests/edge/local/test_session.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The event stream: backpressure that blocks, and an epoch that fences."""
+
+import asyncio
+
+import pytest
+
+from vllm_omni.edge.local.session import (
+    STATE_LAYOUT_VERSION,
+    BoundedEventStream,
+    ChunkEvent,
+    StateHandle,
+    StreamClosed,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _event(seq: int, *, epoch: int = 0, text: str = "x", kind: str = "token") -> ChunkEvent:
+    return ChunkEvent(
+        request_id="r", stage_id=0, seq=seq, epoch=epoch, kind=kind,
+        payload={"text": text}, started_unix=0.0, emitted_unix=0.0,
+    )
+
+
+# -------------------------------------------------------------- state handle
+def test_a_handle_carries_the_artifact_it_was_made_for():
+    handle = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc")
+    assert handle.layout_version == STATE_LAYOUT_VERSION
+    assert handle.accepts(handle)
+
+
+def test_state_from_another_checkpoint_is_not_accepted():
+    """Same shapes from a different quantization are not the same state."""
+    a = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc")
+    b = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="def")
+    assert not a.accepts(b)
+
+
+def test_state_from_a_retired_epoch_is_not_accepted():
+    handle = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc")
+    assert not handle.next_epoch().accepts(handle)
+
+
+def test_state_is_never_migratable_in_m0():
+    handle = StateHandle(session_id="s", backend="vllm:cuda", artifact_id="abc")
+    assert handle.migratable is False
+    assert handle.replayable is True
+
+
+# ---------------------------------------------------------------- ordering
+async def test_events_arrive_in_order():
+    stream = BoundedEventStream()
+    for i in range(5):
+        await stream.put(_event(i))
+    await stream.close()
+    assert [e.seq async for e in stream] == [0, 1, 2, 3, 4]
+
+
+async def test_closing_ends_the_iteration():
+    stream = BoundedEventStream()
+    await stream.put(_event(0))
+    await stream.close()
+    seen = [e async for e in stream]
+    assert len(seen) == 1
+
+
+async def test_producing_after_close_raises():
+    stream = BoundedEventStream()
+    await stream.close()
+    with pytest.raises(StreamClosed):
+        await stream.put(_event(0))
+
+
+# ------------------------------------------------------------ backpressure
+async def test_a_full_queue_blocks_the_producer():
+    """The bound has to actually stop the producer. A queue that grows instead
+    is how a slow consumer turns into an OOM."""
+    stream = BoundedEventStream(max_chunks=2, max_bytes=1 << 20)
+    await stream.put(_event(0))
+    await stream.put(_event(1))
+    blocked = asyncio.create_task(stream.put(_event(2)))
+    await asyncio.sleep(0.05)
+    assert not blocked.done()
+    await stream.get()          # consumer drains one; credit returns
+    await asyncio.wait_for(blocked, timeout=1.0)
+    assert stream.high_water_chunks == 2
+
+
+async def test_the_byte_bound_also_blocks():
+    stream = BoundedEventStream(max_chunks=100, max_bytes=64)
+    await stream.put(_event(0, text="a" * 40))
+    blocked = asyncio.create_task(stream.put(_event(1, text="b" * 40)))
+    await asyncio.sleep(0.05)
+    assert not blocked.done()
+    await stream.get()
+    await asyncio.wait_for(blocked, timeout=1.0)
+
+
+async def test_an_oversized_event_still_gets_through_on_an_empty_queue():
+    """Otherwise one chunk larger than the whole bound deadlocks forever."""
+    stream = BoundedEventStream(max_chunks=4, max_bytes=8)
+    await asyncio.wait_for(stream.put(_event(0, text="x" * 400)), timeout=1.0)
+    assert (await stream.get()).seq == 0
+
+
+async def test_closing_wakes_a_blocked_producer():
+    stream = BoundedEventStream(max_chunks=1, max_bytes=1 << 20)
+    await stream.put(_event(0))
+    blocked = asyncio.create_task(stream.put(_event(1)))
+    await asyncio.sleep(0.05)
+    await stream.close()
+    with pytest.raises(StreamClosed):
+        await asyncio.wait_for(blocked, timeout=1.0)
+
+
+# ------------------------------------------------------------------- epochs
+async def test_a_retired_epoch_cannot_be_produced_into():
+    stream = BoundedEventStream()
+    stream.retire_epoch(0)
+    await stream.put(_event(0, epoch=0))
+    await stream.close()
+    assert [e async for e in stream] == []
+    assert stream.dropped_stale == 1
+
+
+async def test_events_already_queued_from_a_retired_epoch_are_not_delivered():
+    """The real cancellation race: the backend had tokens in flight when the
+    cancel landed, and they must not reach the consumer."""
+    stream = BoundedEventStream()
+    for i in range(3):
+        await stream.put(_event(i, epoch=0))
+    stream.retire_epoch(0)
+    await stream.put(_event(0, epoch=1, text="new"))
+    await stream.close()
+    delivered = [e async for e in stream]
+    assert [e.epoch for e in delivered] == [1]
+    assert stream.dropped_stale == 3
+
+
+async def test_retiring_is_monotonic():
+    stream = BoundedEventStream()
+    stream.retire_epoch(3)
+    stream.retire_epoch(0)
+    assert stream.epoch == 4
+
+
+async def test_stats_report_the_bounds_and_what_was_dropped():
+    stream = BoundedEventStream(max_chunks=8, max_bytes=256)
+    await stream.put(_event(0))
+    stats = stream.stats()
+    assert stats["max_chunks"] == 8
+    assert stats["max_bytes"] == 256
+    assert stats["high_water_chunks"] == 1
+    assert stats["dropped_stale"] == 0

tests/edge/test_hardware_probe.py

@@ -115,3 +115,141 @@ def test_live_probe_runs_on_this_host():
     p = hp.probe(use_torch=False)
     assert p.n_logical >= 1 and p.ram_total_bytes > 0
     assert hp.hardware_class(p) in hp.HW_CLASSES
+
+
+# --------------------------------------------------------- accelerators
+def _ryzen_ai_root(tmp_path: Path, *, wsl: bool, driver: bool = True) -> Path:
+    """A Ryzen AI laptop: XDNA2 NPU, Radeon iGPU, and one discrete card.
+
+    Under WSL none of it has a device node; on native Linux the NPU shows up
+    at ``/dev/accel/accel0`` and the GPUs at ``/dev/dri/renderD*``.
+    """
+    root = _snapshot(tmp_path, X86_CPUINFO, cpus={0: (5100, 0), 1: (5100, 1)})
+    (root / "proc" / "version").write_text(
+        "Linux version 6.18.33.2-microsoft-standard-WSL2" if wsl
+        else "Linux version 6.18.0-generic (gcc 14)"
+    )
+    npu = root / "sys" / "bus" / "pci" / "devices" / "0000:c5:00.1"
+    npu.mkdir(parents=True)
+    (npu / "vendor").write_text("0x1022\n")
+    (npu / "device").write_text("0x17f0\n")
+    if wsl:
+        return root
+    if driver:
+        (npu / "driver").mkdir()
+        node = root / "sys" / "class" / "accel" / "accel0" / "device"
+        node.mkdir(parents=True)
+        (node / "vendor").write_text("0x1022\n")
+        (node / "device").write_text("0x17f0\n")
+        (node / "driver").mkdir()
+        (root / "dev" / "accel").mkdir(parents=True)
+        (root / "dev" / "accel" / "accel0").write_text("")
+    (root / "dev" / "dri").mkdir(parents=True)
+    for name, vendor in (("renderD128", "0x1002"), ("renderD129", "0x10de")):
+        (root / "dev" / "dri" / name).write_text("")
+        d = root / "sys" / "class" / "drm" / name / "device"
+        d.mkdir(parents=True)
+        (d / "vendor").write_text(vendor + "\n")
+        (d / "device").write_text("0x150e\n")
+    return root
+
+
+def test_native_linux_finds_the_npu_and_both_gpus(tmp_path):
+    p = hp.probe(_ryzen_ai_root(tmp_path, wsl=False), use_torch=False)
+    assert p.host_os == hp.HOST_OS_LINUX
+    (npu,) = p.npus
+    assert npu["usable"] and npu["name"].startswith("AMD XDNA2")
+    assert [a["kind"] for a in p.accelerators].count(hp.ACCEL_GPU_INTEGRATED) == 1
+    assert p.accelerators_of(hp.ACCEL_GPU_DISCRETE, usable=True)
+    assert p.unusable_accelerators == []
+
+
+def test_an_npu_with_no_driver_is_reported_present_and_unusable(tmp_path):
+    """"No /dev/accel node" and "no such hardware" need different answers."""
+    p = hp.probe(_ryzen_ai_root(tmp_path, wsl=False, driver=False), use_torch=False)
+    (npu,) = p.npus
+    assert npu["usable"] is False
+    assert "amdxdna" in npu["reason"]
+
+
+def test_wsl_is_detected_and_sees_none_of_the_host_devices(tmp_path):
+    """The iGPU and the NPU are in the machine and unreachable from Linux;
+    the probe must not report either as absent or as available."""
+    p = hp.probe(_ryzen_ai_root(tmp_path, wsl=True), use_torch=False, probe_host=False)
+    assert p.host_os == hp.HOST_OS_WSL2
+    # The PCI sweep still finds the NPU; there is simply no way to open it.
+    (npu,) = p.npus
+    assert npu["usable"] is False
+    assert p.accelerators_of(hp.ACCEL_GPU_INTEGRATED, usable=True) == []
+
+
+def test_a_snapshot_root_never_shells_out_to_the_host(tmp_path):
+    """probe_host defaults on, but a captured snapshot must stay hermetic."""
+    p = hp.probe(_ryzen_ai_root(tmp_path, wsl=True), use_torch=False, probe_host=True)
+    assert all(a["source"] != "wsl-interop" for a in p.accelerators)
+
+
+def test_accelerators_survive_the_json_roundtrip(tmp_path):
+    p = hp.probe(_ryzen_ai_root(tmp_path, wsl=False), use_torch=False)
+    back = hp.HardwareProfile.from_json(p.to_json())
+    assert back.accelerators == p.accelerators
+    assert back.host_os == p.host_os
+
+
+def test_describe_accelerators_says_why_each_one_is_out_of_reach(tmp_path):
+    p = hp.probe(_ryzen_ai_root(tmp_path, wsl=True), use_torch=False, probe_host=False)
+    text = hp.describe_accelerators(p)
+    assert "not here" in text and "amdxdna" in text
+
+
+# ------------------------------------------------------- reachability routes
+def test_the_npu_reason_names_the_recipe_not_just_the_symptom():
+    """The iGPU and the NPU are both off-limits to a default WSL process, but
+    for different reasons with different fixes. A blanket "unreachable" loses
+    the only part a reader can act on -- and for the NPU the actionable part is
+    a three-part recipe, each part of which fails silently on its own."""
+    assert "MCDM" in hp._WSL_NPU_ROUTE                     # why WSL cannot see it
+    assert "onnxruntime_vitisai_ep" in hp._WSL_NPU_ROUTE   # which of the two EPs
+    assert "DLL search path" in hp._WSL_NPU_ROUTE          # how to make it load
+    assert "A16W8" in hp._WSL_NPU_ROUTE                    # what it will accept
+    assert "torch-directml" in hp._WSL_IGPU_ROUTE          # the iGPU's other fix
+
+
+def test_the_npu_recipe_module_states_the_three_preconditions():
+    """Each of the three can fail on its own and the symptom is identical:
+    a session that runs correctly with every node on the CPU."""
+    from vllm_omni.edge import npu_ryzenai as npu
+
+    assert npu.EP_GENERAL == "onnxruntime_vitisai_ep.dll"
+    assert npu.EP_LIGHT != npu.EP_GENERAL
+    assert npu.PARTITION_TARGET == "AMD_AIE2P_4x8_CMC_Overlay"
+    kw = npu.quantization_kwargs()
+    # A16W8: 16-bit activations, 8-bit weights. A8W8 is silently rejected.
+    assert "16" in str(kw["activation_type"])
+    assert "8" in str(kw["weight_type"]) and "16" not in str(kw["weight_type"])
+
+
+def test_the_igpu_is_usable_exactly_when_directml_is_importable(monkeypatch):
+    """torch-directml pins torch==2.4.1, so it never shares the engine's venv.
+    ``usable`` must answer for *this* process, not for the machine."""
+    calls = {"n": 0}
+
+    def fake_run(*a, **k):
+        calls["n"] += 1
+        class R: stdout = "AMD Radeon(TM) 890M Graphics\n"
+        return R()
+
+    import subprocess
+    monkeypatch.setattr(subprocess, "run", fake_run)
+    monkeypatch.setattr(hp, "directml_available", lambda: False)
+    (igpu,) = [a for a in hp.wsl_host_accelerators() if a["kind"] == hp.ACCEL_GPU_INTEGRATED]
+    assert igpu["usable"] is False and igpu["route"] == "directml"
+
+    monkeypatch.setattr(hp, "directml_available", lambda: True)
+    (igpu,) = [a for a in hp.wsl_host_accelerators() if a["kind"] == hp.ACCEL_GPU_INTEGRATED]
+    assert igpu["usable"] is True
+
+
+def test_directml_available_is_false_without_the_package(monkeypatch):
+    monkeypatch.setitem(__import__("sys").modules, "torch_directml", None)
+    assert hp.directml_available() is False

tests/edge/test_kv_budget.py

@@ -105,3 +105,100 @@ def test_dtype_bytes_covers_the_kv_dtypes_vllm_accepts():
 def test_summary_names_the_shape_it_priced():
     s = kv_budget_for(_spark(), max_model_len=32768).summary()
     assert "flat" in s and "21 sliding@512" in s and "MiB" in s
+
+
+def _qwen3_8(**kw) -> _Cfg:
+    """Qwen3.8-27B: 64 layers, 3 linear (gated DeltaNet) to 1 full, 4 KV
+    heads x 256. The linear layers hold a recurrent state and no KV."""
+    layer_types = ["linear_attention"] * 3 + ["full_attention"]
+    base = dict(
+        num_hidden_layers=64, num_attention_heads=24, num_key_value_heads=4,
+        head_dim=256, hidden_size=5120, layer_types=layer_types * 16,
+        linear_num_key_heads=16, linear_num_value_heads=48,
+        linear_key_head_dim=128, linear_value_head_dim=128,
+        linear_conv_kernel_dim=4, mamba_ssm_dtype="float32",
+    )
+    base.update(kw)
+    return _Cfg(**base)
+
+
+class _Outer:
+    """A multimodal config: the decoder geometry lives one level down."""
+
+    def __init__(self, text):
+        self.text_config = text
+        self.model_type = "qwen3_5"
+
+    def get_text_config(self):
+        return self.text_config
+
+
+def test_linear_attention_layers_are_charged_no_kv_at_all():
+    """Not "a shorter cache" -- none. A flat 64-layer price is 4x the truth."""
+    b = kv_budget_for(_qwen3_8(), max_model_len=8192)
+    assert (b.linear_layers, b.full_layers, b.sliding_layers) == (48, 16, 0)
+    # 16 full layers x 4 KV heads x 256 x 2 (K and V) x 2 bytes = 64 KiB.
+    assert b.bytes_per_token_flat == 16 * 4 * 256 * 2 * 2 == 64 * 1024
+    assert b.token_layers_flat == 16 * 8192
+
+
+def test_the_recurrent_state_is_sized_from_vllms_own_shapes():
+    """conv (kernel-1, 2*k_heads*k_dim + v_heads*v_dim); temporal
+    (v_heads, v_dim, k_dim) in mamba_ssm_dtype."""
+    b = kv_budget_for(_qwen3_8(), max_model_len=8192)
+    conv_dim = 128 * 16 * 2 + 128 * 48
+    per_layer = (4 - 1) * conv_dim * 2 + 48 * 128 * 128 * 4
+    assert b.state_bytes_per_seq == 48 * per_layer
+    assert b.state_bytes_per_seq / 2**20 == pytest.approx(147.0, abs=0.5)
+    assert b.state_bytes_total == b.state_bytes_per_seq  # max_num_seqs=1
+
+
+def test_the_state_is_inside_the_budget_and_scales_with_sequences():
+    """It comes out of the same pool as the KV cache, so a budget that
+    reports only the cache is short before the first token."""
+    one = kv_budget_for(_qwen3_8(), max_model_len=8192, max_num_seqs=1)
+    two = kv_budget_for(_qwen3_8(), max_model_len=8192, max_num_seqs=2)
+    assert two.state_bytes_total == 2 * one.state_bytes_total
+    assert one.bytes_total > one.state_bytes_total
+    kv_only = int(one.token_layers_flat * one.bytes_per_token_flat / 16 * 1.25)
+    assert one.bytes_total == kv_only + one.state_bytes_total
+
+
+def test_the_state_dominates_a_short_context():
+    """147 MiB of state against 64 KiB/token: the cache only outgrows it past
+    ~2350 tokens, so over an edge deployment's range the constant is the
+    bigger half. The crossover is on raw bytes -- ``headroom`` is a safety
+    factor on the cache, not a cost the state has to beat."""
+    assert kv_budget_for(
+        _qwen3_8(), max_model_len=2048
+    ).state_crossover_tokens == pytest.approx(2352, abs=8)
+    short = kv_budget_for(_qwen3_8(), max_model_len=1024)
+    assert short.state_bytes_total > short.bytes_total - short.state_bytes_total
+    long = kv_budget_for(_qwen3_8(), max_model_len=32768)
+    assert long.state_bytes_total < long.bytes_total - long.state_bytes_total
+
+
+def test_a_model_without_linear_layers_reports_no_state():
+    b = kv_budget_for(_spark(), max_model_len=2048)
+    assert b.linear_layers == 0
+    assert b.state_bytes_per_seq == 0
+    assert b.state_crossover_tokens == float("inf")
+
+
+def test_a_multimodal_config_is_resolved_to_its_text_config():
+    """Qwen3_5Config has no num_hidden_layers; this used to raise."""
+    nested = kv_budget_for(_Outer(_qwen3_8()), max_model_len=8192)
+    direct = kv_budget_for(_qwen3_8(), max_model_len=8192)
+    assert nested == direct
+
+
+def test_linear_layers_without_their_geometry_raise_rather_than_guess():
+    cfg = _qwen3_8()
+    del cfg.linear_num_value_heads
+    with pytest.raises(ValueError, match="linear_num_value_heads"):
+        kv_budget_for(cfg, max_model_len=8192)
+
+
+def test_summary_names_the_state_it_priced():
+    s = kv_budget_for(_qwen3_8(), max_model_len=8192).summary()
+    assert "recurrent state for 48 linear layers" in s

tests/edge/test_placement.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Placing a checkpoint that does not fit in the VRAM."""
+
+import json
+import struct
+
+import pytest
+
+from vllm_omni.edge.placement import (
+    GiB,
+    ParamGroup,
+    checkpoint_param_groups,
+    group_key,
+    is_cold,
+    plan_placement,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+LINK = 13.9e9   # measured pinned H2D on the reference laptop, PCIe 4.0 x8
+VRAM = 472e9    # measured VRAM bandwidth on the same machine
+
+
+def _plan(groups, avail_gib, **kw):
+    return plan_placement(
+        groups, device_bytes_available=int(avail_gib * GiB),
+        link_bytes_per_s=LINK, vram_bytes_per_s=VRAM, **kw)
+
+
+# ------------------------------------------------------------ classification
+def test_embed_tokens_is_cold_and_lm_head_is_not():
+    """Same shape, same dtype, opposite decision: the embedding is a gather
+    and the head is a GEMM over the whole table."""
+    assert is_cold("model.language_model.embed_tokens.weight")
+    assert not is_cold("lm_head.weight")
+
+
+def test_cold_matching_is_by_segment_not_substring():
+    assert is_cold("model.visual.blocks.0.attn.qkv.weight")
+    assert not is_cold("model.visualiser.weight")
+
+
+def test_a_decoder_layer_is_never_folded_into_one_group():
+    """A 3-segment prefix makes all 64 layers a single 24 GB group, and a
+    planner that moves groups whole then has no choice between none and all."""
+    assert group_key("model.language_model.layers.17.mlp.down_proj.weight") == (
+        "model.language_model.layers.17")
+    assert group_key("model.visual.blocks.3.attn.qkv.weight").startswith("model.visual")
+
+
+# ------------------------------------------------------------------ planning
+def test_a_checkpoint_that_fits_is_left_alone():
+    groups = [ParamGroup("model.layers.0", 8 * GiB, False)]
+    p = _plan(groups, 20)
+    assert p.strategy == "resident"
+    assert p.host_bytes == 0 and p.est_link_s_per_token == 0.0
+    assert p.engine_kwargs == {}
+
+
+def test_the_deficit_comes_out_of_the_cold_set_first():
+    """4 GiB over, and there is 6 GiB of cold weight to spend: nothing that a
+    decode step reads should end up on the host."""
+    groups = [
+        ParamGroup("model.layers.0", 18 * GiB, False),
+        ParamGroup("model.embed_tokens", 6 * GiB, True),
+    ]
+    p = _plan(groups, 20)
+    assert p.strategy == "cold_offload"
+    assert p.host_bytes == 6 * GiB
+    assert p.host_read_per_token == 0
+    assert p.est_link_s_per_token == 0.0
+
+
+def test_an_offloaded_embedding_still_pays_one_row_per_token():
+    groups = [
+        ParamGroup("model.layers.0", 18 * GiB, False),
+        ParamGroup("model.embed_tokens", 6 * GiB, True),
+    ]
+    p = _plan(groups, 20, embedding_row_bytes=5120 * 2)
+    assert p.host_read_per_token == 5120 * 2
+    # 10 KiB over a 13.9 GB/s link is well under a microsecond.
+    assert p.est_link_s_per_token < 1e-6
+
+
+def test_hot_weight_is_only_spent_once_the_cold_set_runs_out():
+    groups = [
+        ParamGroup("model.embed_tokens", 2 * GiB, True),
+        *[ParamGroup(f"model.layers.{i}", 1 * GiB, False) for i in range(20)],
+    ]
+    p = _plan(groups, 16)
+    assert p.strategy == "hot_offload"
+    # 6 GiB over; 2 GiB is free, so exactly 4 layers should cross.
+    assert p.host_read_per_token == 4 * GiB
+    assert p.host_bytes == 6 * GiB
+
+
+def test_the_estimate_prices_host_bytes_at_the_link_and_the_rest_at_vram():
+    groups = [
+        ParamGroup("model.embed_tokens", 2 * GiB, True),
+        *[ParamGroup(f"model.layers.{i}", 1 * GiB, False) for i in range(20)],
+    ]
+    p = _plan(groups, 16)
+    assert p.est_link_s_per_token == pytest.approx(4 * GiB / LINK)
+    assert p.est_vram_s_per_token == pytest.approx(16 * GiB / VRAM)
+    # The link is 34x slower, so 4 GiB on the host outweighs 16 GiB resident.
+    assert p.est_link_s_per_token > 8 * p.est_vram_s_per_token
+
+
+def test_a_checkpoint_too_big_for_host_and_device_says_so():
+    groups = [ParamGroup("model.layers.0", 40 * GiB, False)]
+    p = _plan(groups, 8)
+    # Everything hot can still move to the host, so this one is feasible but
+    # miserable; infeasible needs the groups themselves to run out.
+    assert p.strategy == "hot_offload"
+    assert p.deficit_bytes == 32 * GiB
+
+
+def test_engine_kwargs_name_only_segments_vllm_can_match():
+    groups = [
+        ParamGroup("model.embed_tokens", 6 * GiB, True),
+        ParamGroup("model.layers.0", 18 * GiB, False),
+    ]
+    kw = _plan(groups, 20).engine_kwargs
+    assert kw["cpu_offload_params"] == ["embed_tokens"]
+    assert kw["cpu_offload_gb"] >= 6.0
+    assert "offload_group_size" not in kw  # nothing hot was offloaded
+
+
+# --------------------------------------------------------------- checkpoints
+def _write_shard(path, tensors):
+    """A minimal safetensors file: 8-byte header length, JSON header, data."""
+    header, off = {}, 0
+    for name, nbytes in tensors:
+        header[name] = {"dtype": "BF16", "shape": [nbytes // 2],
+                        "data_offsets": [off, off + nbytes]}
+        off += nbytes
+    blob = json.dumps(header).encode()
+    path.write_bytes(struct.pack("<Q", len(blob)) + blob + b"\0" * off)
+
+
+def test_param_groups_are_read_from_headers_without_loading_weights(tmp_path):
+    _write_shard(tmp_path / "a.safetensors", [
+        ("model.language_model.embed_tokens.weight", 4096),
+        ("lm_head.weight", 4096),
+        ("model.language_model.layers.0.mlp.down_proj.weight", 2048),
+        ("model.language_model.layers.1.mlp.down_proj.weight", 2048),
+    ])
+    groups = {g.name: g for g in checkpoint_param_groups(tmp_path)}
+    assert groups["model.language_model.embed_tokens"].cold
+    assert not groups["lm_head.weight"].cold
+    assert set(groups) >= {"model.language_model.layers.0",
+                           "model.language_model.layers.1"}
+    assert groups["model.language_model.layers.0"].bytes == 2048
+
+
+def test_a_group_is_cold_only_if_every_tensor_in_it_is(tmp_path):
+    _write_shard(tmp_path / "a.safetensors", [
+        ("model.visual.blocks.0.attn.qkv.weight", 1024),
+        ("model.visual.blocks.0.lm_head.weight", 1024),  # contrived, but hot
+    ])
+    (group,) = [g for g in checkpoint_param_groups(tmp_path)
+                if g.name.startswith("model.visual")]
+    assert not group.cold
+
+
+def test_a_directory_with_no_shards_raises(tmp_path):
+    with pytest.raises(FileNotFoundError, match="safetensors"):
+        checkpoint_param_groups(tmp_path)
+
+
+# ------------------------------------------- the cost model, against hardware
+# Qwen3.8-27B-FP8, 4096-token context, batch 1, greedy, eager, idle machine.
+# Four points from benchmarks/edge_harness; see docs/edge/laptop-heterogeneous.md.
+MEASURED_S_PER_TOKEN = {10: 0.952, 12: 1.124, 14: 1.316, 16: 1.515}
+QWEN38_FP8_GB = 30.87
+
+
+def test_the_estimate_reproduces_the_measured_offload_curve():
+    """The whole design rests on `resident/472 + offloaded/10.6`. It is worth
+    one test that the arithmetic matches the hardware it claims to describe.
+
+    It runs slightly *pessimistic* at every point, which is the right
+    direction for a number used to decide whether a build will fit.
+    """
+    # 0.25 GB chunks, so the planner offloads close to exactly the target
+    # rather than having to take one oversized group whole -- which is also
+    # what the real checkpoint looks like at 0.38 GB per decoder layer.
+    CHUNK = 0.25
+    n = round(QWEN38_FP8_GB / CHUNK)
+    for offloaded_gb, measured in MEASURED_S_PER_TOKEN.items():
+        groups = [ParamGroup(f"layers.{i}", int(CHUNK * 1e9), False) for i in range(n)]
+        p = plan_placement(
+            groups,
+            device_bytes_available=int((QWEN38_FP8_GB - offloaded_gb) * 1e9),
+            link_bytes_per_s=10.6e9, vram_bytes_per_s=472e9,
+        )
+        modelled = p.est_link_s_per_token + p.est_vram_s_per_token
+        assert modelled == pytest.approx(measured, rel=0.05), (
+            f"{offloaded_gb} GB: modelled {modelled:.3f}s vs measured {measured:.3f}s")
+        assert modelled >= measured, "the estimate must not flatter the build"
+
+
+def test_the_link_term_dominates_at_every_measured_point():
+    """At the tightest fit the 10 GB on the host is ~94% of the step. If this
+    ever stops being true the advice in the design document changes."""
+    groups = [
+        ParamGroup("resident", int(20.87e9), False),
+        ParamGroup("onhost", int(10e9), False),
+    ]
+    p = plan_placement(groups, device_bytes_available=int(20.87e9),
+                       link_bytes_per_s=10.6e9, vram_bytes_per_s=472e9)
+    share = p.est_link_s_per_token / (p.est_link_s_per_token + p.est_vram_s_per_token)
+    assert share > 0.9

tests/edge/test_qnn_export.py

@@ -105,6 +105,26 @@ def test_masked_cache_slots_do_not_change_the_result():
     assert torch.allclose(first, module(*args)[0], atol=1e-5)
 
 
+def test_llama_style_config_skips_qk_norm():
+    """MiniCPM-o's TTS head is a Llama decoder: no per-head q/k RMSNorm, so those params must not be applied."""
+    cfg = qx.DecodeStepConfig(**{**{f.name: getattr(CFG, f.name) for f in CFG.__dataclass_fields__.values()}, "qk_norm": False})
+    module = qx.DecodeStep(cfg).eval()
+    for p in module.parameters():
+        torch.nn.init.normal_(p, std=0.15)
+    args = qx.example_inputs(cfg, L)
+    before = module(*args)[0].clone()
+    with torch.no_grad():  # q_norm/k_norm must be dead weights when qk_norm is off
+        module.layers[0].q_norm.mul_(3.0)
+        module.layers[0].k_norm.mul_(0.1)
+    assert torch.equal(before, module(*args)[0])
+
+
+def test_minicpmo_config_maps_the_tts_head():
+    cfg = qx.config_from_minicpmo({"tts_config": {"num_hidden_layers": 20, "hidden_size": 768, "num_attention_heads": 12,
+                                                  "num_key_value_heads": 12, "intermediate_size": 3072, "num_audio_tokens": 6562}})
+    assert cfg.qk_norm is False and cfg.head_dim == 64 and cfg.vocab_size == 6562 and cfg.group_size == 1
+
+
 def test_placement_table_prefers_npu_for_the_step_and_gpu_for_the_vocoder():
     p = qx.PLACEMENT_MS_PER_FRAME
     assert p["talker_step"]["npu"] < p["talker_step"]["gpu"]

tests/model_executor/models/spark2_5/test_spark2_5.py

@@ -243,3 +243,84 @@ def test_rope_table_is_capped_to_the_servable_context(monkeypatch):
     if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0":
         max_position = min(max_position, 2048)
     assert max_position == 1048576, "the escape hatch must restore the full table"
+
+
+# ------------------------------------------------------- published sizes
+# Spark-X2.5 ships at two sizes and the vLLM implementation is meant to be
+# size-generic. It is, but only one of them had ever been run: the constants
+# below are the published configs, so a change that quietly assumes 1.7B's
+# 8 query heads or 28 layers fails here instead of on a checkpoint.
+PUBLISHED = {
+    "XHToken/Spark-X2.5-1.7B": dict(
+        hidden_size=2048, intermediate_size=6656, num_hidden_layers=28,
+        num_attention_heads=8, num_key_value_heads=2,
+    ),
+    "XHToken/Spark-X2.5-4B": dict(
+        hidden_size=2560, intermediate_size=10240, num_hidden_layers=36,
+        num_attention_heads=16, num_key_value_heads=4,
+    ),
+}
+_COMMON = dict(
+    head_dim=256, vocab_size=131072, sliding_window=512, hidden_act="gelu",
+    headwise_attn_output_gate=True, tie_word_embeddings=True,
+    max_position_embeddings=1048576,
+    rope_parameters={
+        "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25},
+        "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0},
+    },
+)
+
+
+def _published(name: str) -> Spark2_5Config:
+    kw = dict(_COMMON, **PUBLISHED[name])
+    n = kw["num_hidden_layers"]
+    kw["layer_types"] = (["sliding_attention"] * 3 + ["full_attention"]) * (n // 4)
+    return Spark2_5Config(**kw)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("name", sorted(PUBLISHED))
+def test_published_configs_validate_and_keep_the_3_to_1_hybrid(name):
+    config = _published(name)
+    assert len(config.layer_types) == config.num_hidden_layers
+    assert config.layer_types.count("full_attention") == config.num_hidden_layers // 4
+    # Both sizes keep head_dim 256 and widen by head count, not by head width.
+    assert config.head_dim == 256
+    assert config.num_attention_heads % config.num_key_value_heads == 0
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("name", sorted(PUBLISHED))
+def test_split_fused_qkv_holds_at_both_published_head_counts(name):
+    """4B is 16 q heads to 4 KV heads where 1.7B is 8 to 2. The split is by
+    dimension, not by a ratio, so getting one right does not imply the other."""
+    config = _published(name)
+    q_dim = config.num_attention_heads * config.head_dim
+    kv_dim = config.num_key_value_heads * config.head_dim
+    fused = torch.cat([
+        torch.full((q_dim, 8), 1.0),
+        torch.full((kv_dim, 8), 2.0),
+        torch.full((kv_dim, 8), 3.0),
+    ], dim=0)
+    out = dict(_split_fused_qkv(
+        [("model.layers.0.self_attn.q_k_v_proj.weight", fused)], config))
+    assert out["model.layers.0.self_attn.q_proj.weight"].shape[0] == q_dim
+    assert out["model.layers.0.self_attn.k_proj.weight"].shape[0] == kv_dim
+    assert torch.all(out["model.layers.0.self_attn.q_proj.weight"] == 1.0)
+    assert torch.all(out["model.layers.0.self_attn.v_proj.weight"] == 3.0)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_4b_is_priced_as_a_hybrid_not_as_36_full_layers():
+    """The KV budget reads the same config the model does; 4B is 27 sliding
+    layers to 9 full ones, which is a 3.4x discount a flat price would miss."""
+    from vllm_omni.edge.kv_budget import kv_budget_for
+
+    b = kv_budget_for(_published("XHToken/Spark-X2.5-4B"), max_model_len=32768)
+    assert (b.full_layers, b.sliding_layers, b.linear_layers) == (9, 27, 0)
+    # 36 layers x 4 KV heads x 256 x 2 (K and V) x 2 bytes = 144 KiB/token.
+    assert b.bytes_per_token_flat == 36 * 4 * 256 * 2 * 2 == 144 * 1024
+    assert b.saving_vs_flat > 3.3

tests/worker/test_forward_kwargs_filtering.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Omni-only forward kwargs must not be forced on a model that cannot take them.
+
+The AR runner hands every stage ``sampling_metadata``, ``logits_index`` and
+``sampler``, which omni-native models absorb through ``**kwargs``. A model vLLM
+owns does not: ``LlamaForCausalLM.forward`` is
+``(input_ids, positions, intermediate_tensors, inputs_embeds)`` and nothing
+else, so passing the extras is a ``TypeError`` on the very first forward.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+
+class _OmniNativeModel(torch.nn.Module):
+    def forward(self, input_ids, positions, intermediate_tensors=None, inputs_embeds=None, **kwargs):
+        return input_ids
+
+
+class _VLLMNativeModel(torch.nn.Module):
+    """The shape of vLLM's own causal LMs: no ``**kwargs``."""
+
+    def forward(self, input_ids, positions, intermediate_tensors=None, inputs_embeds=None):
+        return input_ids
+
+
+class _Probe:
+    """Just enough of a runner to exercise the helper."""
+
+    def __init__(self, model):
+        self.model = model
+
+    _accepted_forward_kwargs = OmniGPUModelRunner._accepted_forward_kwargs
+
+
+def test_a_model_taking_kwargs_is_left_alone():
+    assert _Probe(_OmniNativeModel())._accepted_forward_kwargs() is None
+
+
+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"}
+    )
+    for omni_only in ("sampling_metadata", "logits_index", "sampler"):
+        assert omni_only not in accepted
+
+
+def test_the_answer_is_cached():
+    """``inspect.signature`` on every decode step would be absurd."""
+    probe = _Probe(_VLLMNativeModel())
+    first = probe._accepted_forward_kwargs()
+    probe.model = _OmniNativeModel()  # would change the answer if recomputed
+    assert probe._accepted_forward_kwargs() is first
+
+
+def test_an_uninspectable_forward_falls_back_to_passing_everything():
+    """Degrade to the previous behaviour rather than silently dropping
+    arguments a model may actually need."""
+
+    class Weird(torch.nn.Module):
+        forward = staticmethod(print)  # no usable signature for our purposes
+
+    accepted = _Probe(Weird())._accepted_forward_kwargs()
+    assert accepted is None or "sampling_metadata" not in accepted