[Bugfix] Preserve host transport fallback when discovery cleanup fails

Signed-off-by: Sy03 <1370724210@qq.com>

docs/design/feature/stage_host_transport.md

@@ -0,0 +1,84 @@
+# Local stage host transport
+
+`SharedMemoryConnector` uses a persistent host ring for ordinary local stage
+payloads by default. Models keep the existing `put`/`get` and typed payload
+contracts; no model-specific sender, receiver or deployment change is required.
+This applies to both V1 and MRV2 paths using this connector, including CPU codec
+codes and CPU hidden states. Remote connectors use their existing transports.
+
+Each producer process and destination stage lazily allocates a 4 MiB ring.
+Small CPU tensor trees use a msgpack header and raw tensor regions. Plain
+serialized payloads can use the same ring. Frames larger than a quarter of the
+ring, and frames arriving while it is full, use the existing per-key SHM path.
+The fallback preserves delivery without blocking a model step on ring capacity.
+Large or unsupported tensor trees do not allocate an unused contiguous snapshot.
+
+## Ownership and delivery
+
+The producer copies the payload before `put` returns. Readers hold the ring's
+file lock while copying into independently owned CPU storage, then mark the
+frame consumed. Multiple receiving processes may compete for a key; one reader
+claims it. Existing TP fanout remains responsible for distributing that owned
+payload to peer ranks. This transport does not provide multicast.
+
+Small tensor leaves share one owned allocation. Large multi-leaf messages use
+separate allocations, so retaining small metadata does not retain a large
+hidden state. No view of ring storage escapes the receive boundary.
+
+Cancellation retires exact keys or tracked numeric chunk suffixes. Replacing
+an unclaimed key retires its previous frame, including transitions between
+the ring and per-key SHM. Descriptors include monotonic frame positions; an old
+descriptor cannot read a newer publication after wraparound. Keys must have
+one active publisher, as assumed by the existing stage request routing.
+
+Only producers register ring allocations with Python's shared-memory resource
+tracker. Readers map files without taking ownership. Producer process identity
+includes its start time, and available Linux pidfds detect producer exit, so
+old frames cannot be delivered after a producer restart. Normal close removes
+owned allocations and discovery markers. Resource tracking covers allocation
+cleanup when the owning application exits; an abnormal exit can leave small
+discovery markers, which contain no payload data.
+Live receivers release cached mappings after a producer exits and prune stale
+markers during discovery within their deployment scope.
+
+## Configuration and observation
+
+Connector extras are shared by the endpoints of an edge. For example:
+
+```yaml
+connectors:
+  shm:
+    name: SharedMemoryConnector
+    extra:
+      host_ring_bytes: 4194304
+```
+
+Set `host_ring_bytes: 0` at both endpoints to use the per-key path, including
+controlled comparisons or deployments with older receivers. Nonzero capacities
+must be between 4 KiB and 64 MiB. All ring endpoints must support this protocol.
+The launcher supplies a common deployment scope; discovery scans only that
+scope's directory. Receive attempts do not wait for another process's file lock.
+An unchanged empty channel skips its frame lock. Channel registration and
+retirement advance a serialized discovery stamp, so empty polling does not
+rescan the directory and new producers remain discoverable without a timer.
+Ring endpoints must share a Linux PID namespace. Deployments that share SHM
+across different PID namespaces must use `host_ring_bytes: 0`.
+
+`health()` exposes `host_ring_puts`, `host_ring_gets` and `host_ring_fallbacks`
+alongside the existing counters. Ring setup failure logs once per edge and uses
+per-key SHM. CPU payload size and the cost of serialization determine the benefit;
+reduced transfer latency does not imply higher model throughput.
+
+## Scope
+
+This mechanism reduces host serialization copies and per-message segment work.
+It does not retain worker outputs on the GPU, allocate a CUDA arena, or replace
+GPU stream synchronization. It also does not impose a global byte budget on
+queued snapshots, downstream model caches or client audio delivery. The bounded
+ring with a compatibility fallback is not end-to-end streaming backpressure.
+
+Regression tests are in
+`tests/distributed/omni_connectors/test_shm_ring.py`, collected by the existing
+`core_model and cpu` CI lane. They cover ownership, tensor types, typed payloads,
+pressure fallback, cancellation, overwrite, deadlines, competing processes and
+producer restart.

docs/design/minicpmo_encoder_cuda_graph.md

@@ -1,61 +0,0 @@
-# MiniCPM-o 4.5 input encoder CUDA graphs
-
-The vision transformer and stateless audio encoder reuse vLLM's encoder graph
-manager for exact input shapes. Each encoder retains at most four graphs by
-default, and captures on the second call with the same shape and CUDA stream.
-Startup profiling can consume slots. Graphs are never evicted: once the cache
-is full, other shapes run eagerly for the lifetime of the model.
-
-Use `--hf-overrides` to configure both encoders:
-
-```json
-{
-  "encoder_cuda_graph": true,
-  "encoder_cuda_graph_max_graphs": 4,
-  "encoder_cuda_graph_min_capture_calls": 2,
-  "encoder_cuda_graph_min_free_bytes": 1073741824,
-  "encoder_cuda_graph_share_pools": true
-}
-```
-
-`max_graphs` is a nonnegative integer per encoder; zero disables capture.
-`min_capture_calls` is an integer of at least two. Increasing it reduces capture
-of short-lived shapes. Admission history is bounded to four times `max_graphs`,
-so an evicted history entry must accumulate its calls again. Increasing the
-graph cap can improve coverage but retains more GPU memory. Existing graphs
-are not replaced, avoiding repeated capture costs when traffic changes.
-
-By default, graphs in the same encoder on the same device and replay stream
-share a memory pool and capture stream. Different encoders and replay streams
-remain isolated. Graph inputs and output buffers are allocated outside capture;
-only transient capture allocations share the pool. Returned outputs remain
-cloned. A late capture waits for earlier replay work and output clones before
-reusing the pool. Set `encoder_cuda_graph_share_pools` to `false` for an A/B
-comparison with a separate pool and capture stream per graph.
-
-The graph count is **not a byte limit**. Sharing pools reuses intermediate
-storage but does not eliminate per-shape static input and output buffers.
-Before a new capture, the adapter checks device-free memory against
-`min_free_bytes`, a nonnegative integer (default: 1 GiB; zero disables the check).
-Below this floor the call runs eagerly, and a later call may try admission
-again. Existing graphs can still replay. This check is conservative about
-reusable allocator memory and does not predict capture allocations, reserve
-memory, or guarantee that capture fits. Size the floor for the deployment;
-disable encoder graphs when their additional memory cannot be accommodated.
-
-Each adapter exposes `get_cumulative_stats()`: graph hits, total misses, hit
-rate, graph count, and separate ineligible, capacity, warmup and memory misses.
-These include calls rejected before the upstream manager, so a full cache's
-lost coverage is visible. Vision uses `vpm._encoder_graph`; audio uses
-`_audio_encoder_graph` on the model. The counters are per adapter, not global.
-
-Capture failures propagate the original exception and make that adapter
-unusable. Subsequent calls raise without retrying capture or running CUDA eager
-work. Restart the worker, optionally with encoder graphs disabled. This is a
-fatal contract, not automatic recovery: a failed CUDA capture can leave the
-device context unusable. `capture_failures` records this terminal state.
-
-`--enforce-eager` always disables these graphs. CPU, grad/autocast, stateful
-streaming audio and the other model-specific eager fallbacks remain unchanged.
-The adapter clones graph outputs through upstream postprocessing so retained
-embeddings survive later replays.

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -31,10 +31,10 @@ from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-@pytest.fixture
-def shm_sender(build_adapter):
+@pytest.fixture(params=[0, 4 << 20], ids=["per-key-shm", "host-ring"])
+def shm_sender(build_adapter, request):
     adapter, _ = build_adapter(stage_id=0)
-    connector = SharedMemoryConnector({"stage_id": 0})
+    connector = SharedMemoryConnector({"stage_id": 0, "host_ring_bytes": request.param})
     adapter.connector = connector
     adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
         codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
@@ -111,13 +111,14 @@ def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
 
     adapter.release_shm_resources(ext_id)
     # The release is queued behind sends; nothing is unlinked on the caller.
-    assert f"{ext_id}_0_1" in connector._pending_keys
+    assert adapter._reclaimed_shm_total == 0
     while adapter._pending_save_reqs:
         adapter._send_single_request(adapter._pending_save_reqs.popleft())
 
     assert adapter._reclaimed_shm_total == 2
     assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
     assert connector.get("0", "1", f"{ext_id}_0_1") is None
+    assert connector.get("0", "1", f"{ext_id}_0_2") is None
     assert connector.get("0", "1", sibling_key)[0] == "sibling"
 
 

tests/distributed/omni_connectors/test_shm_connector.py

@@ -19,7 +19,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 @pytest.fixture()
 def connector():
-    c = SharedMemoryConnector({})
+    # These cases inspect the per-key segment/lock protocol. The enabled ring
+    # protocol and its pressure fallback are covered in test_shm_ring.py.
+    c = SharedMemoryConnector({"host_ring_bytes": 0})
     yield c
     c.close()
 

tests/distributed/omni_connectors/test_shm_ring.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cross-process ownership, pressure and compatibility of the default host path."""
+
+import fcntl
+import multiprocessing
+import os
+import time
+import uuid
+from dataclasses import asdict
+from multiprocessing import shared_memory
+from pathlib import Path
+
+import msgspec
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.outputs import CompletionOutput, RequestOutput
+
+from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
+from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+from vllm_omni.distributed.omni_connectors.utils.tensor_frame import prepare_tensor_frame
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def edge():
+    scope = uuid.uuid4().hex
+    sender = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope, "host_ring_bytes": 4096})
+    receiver = SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope, "host_ring_bytes": 4096})
+    yield sender, receiver
+    receiver.close()
+    sender.close()
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+        torch.complex64,
+    ],
+)
+def test_owned_tensor_tree_preserves_values_shapes_and_types(edge, dtype):
+    sender, receiver = edge
+    tensor = torch.arange(24).to(dtype).reshape(4, 6).T
+    data = {
+        "codes": {"audio": tensor},
+        "other": (tensor[0], [torch.empty(0, dtype=dtype), torch.tensor(1)]),
+        "__tensor__": False,
+        "meta": {"finished": torch.tensor(False)},
+    }
+    ok, size, metadata = sender.put("0", "1", "tree", data)
+    assert ok and "host_ring" in metadata
+    output, read_size = receiver.get("0", "1", "tree", metadata)
+    assert read_size == size and output["__tensor__"] is False
+    assert type(output["other"]) is list and type(output["other"][1]) is list
+    assert output["other"][1][0].shape == (0,)
+    assert output["other"][1][1].shape == ()
+    assert output["codes"]["audio"].dtype == dtype
+    assert torch.equal(output["codes"]["audio"], tensor)
+    tensor.fill_(0)
+    assert torch.count_nonzero(output["codes"]["audio"]) > 0
+    for index in range(100):
+        sender.put("0", "1", f"wrap{index}", {"codes": torch.full((56,), index)})
+        assert receiver.get("0", "1", f"wrap{index}")[0]["codes"][0] == index
+    assert torch.count_nonzero(output["codes"]["audio"]) > 0
+
+
+def test_conjugate_tensor_and_scalar_finish(edge):
+    sender, receiver = edge
+    tensor = torch.tensor([1 + 2j, 3 - 4j]).conj()
+    sender.put("0", "1", "conjugate", {"condition": tensor, "finished": torch.tensor(True)})
+    result = receiver.get("0", "1", "conjugate")[0]
+    assert torch.equal(result["condition"], tensor)
+    assert result["finished"].shape == () and result["finished"].item() is True
+
+
+def test_shared_typed_payload_keeps_omitted_defaults_and_cpu_tensors(edge):
+    sender, receiver = edge
+    payload = OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.arange(28).reshape(1, 28)),
+        meta=MetaStruct(finished=torch.tensor(False), left_context_size=0),
+        kv_metadata={"binary_handle": b"\x00\xff\x10"},
+    )
+    ok, _, metadata = sender.put("0", "1", "typed", payload)
+    assert ok and "host_ring" in metadata
+    result = receiver.get("0", "1", "typed")[0]
+    assert set(result) == {"codes", "meta", "kv_metadata"}
+    assert set(result["codes"]) == {"audio"}
+    assert torch.equal(result["codes"]["audio"], payload.codes.audio)
+    assert result["meta"]["left_context_size"] == 0
+    assert result["kv_metadata"]["binary_handle"] == b"\x00\xff\x10"
+
+
+def test_source_can_be_overwritten_as_soon_as_put_returns(edge):
+    sender, receiver = edge
+    source = torch.arange(56)
+    sender.put("0", "1", "source-reuse", source)
+    source.fill_(-1)
+    assert receiver.get("0", "1", "source-reuse")[0].tolist() == list(range(56))
+
+
+@pytest.mark.parametrize("dtype", [np.bool_, np.int32, np.float16, np.float64, np.complex64])
+def test_mixed_numpy_image_and_tensor_keep_shared_wire_semantics(edge, dtype):
+    sender, receiver = edge
+    array = np.arange(24).astype(dtype).reshape(4, 6)[::-1, ::2]
+    image = Image.fromarray(np.arange(18, dtype=np.uint8).reshape(2, 3, 3))
+    payload = {"codes": torch.arange(28), "array": array, "image": image, "slice": slice(1, 7, 2)}
+    _, _, metadata = sender.put("0", "1", "mixed-native", payload)
+    assert "host_ring" in metadata
+    output = receiver.get("0", "1", "mixed-native", metadata)[0]
+    assert np.array_equal(output["array"], array)
+    assert output["array"].dtype == array.dtype
+    assert not output["array"].flags.writeable
+    assert output["image"].mode == image.mode and output["image"].tobytes() == image.tobytes()
+    assert output["slice"] == [1, 7, 2]
+    assert torch.equal(output["codes"], payload["codes"])
+    array.fill(0)
+    image.paste(0, (0, 0, 3, 2))
+    assert np.count_nonzero(output["array"]) and any(output["image"].tobytes())
+
+
+def test_large_array_rejects_unused_framing_snapshot(monkeypatch):
+    def unexpected_snapshot(value):
+        pytest.fail("a too-large array must not be snapshotted while testing ring admission")
+
+    monkeypatch.setattr(OmniSerializer.encoder, "_enc_hook", unexpected_snapshot)
+    assert prepare_tensor_frame({"array": np.zeros(1024)}, max_bytes=1024) is None
+
+
+def test_native_scalar_restore_preserves_nested_markers_and_input_containers():
+    array = np.arange(3)
+    image = Image.fromarray(np.array([[1, 2]], dtype=np.uint8))
+    tensor = torch.arange(2)
+    payload = {"ids": list(range(8192)), "mixed": [1.5, None, (array, {"image": image, "tensor": tensor})]}
+    decoded = msgspec.msgpack.decode(OmniSerializer.serialize(payload))
+    output = OmniSerializer.restore(decoded)
+    assert output is not decoded and output["ids"] is not decoded["ids"]
+    assert output["ids"] == payload["ids"]
+    assert np.array_equal(output["mixed"][2][0], array)
+    assert output["mixed"][2][1]["image"].tobytes() == image.tobytes()
+    assert torch.equal(output["mixed"][2][1]["tensor"], tensor)
+    assert isinstance(decoded["mixed"][2][0], dict)
+
+
+def test_fallback_miss_does_not_register_an_unclaimed_allocation(edge, monkeypatch, caplog):
+    sender, receiver = edge
+    key = "fallback-claim-" + uuid.uuid4().hex
+    _, _, metadata = sender.put("0", "1", key, torch.arange(4096))
+    assert "shm" in metadata
+    original = shared_memory.SharedMemory
+    calls = []
+
+    def tracked_open(*args, **kwargs):
+        calls.append((args, kwargs))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(shared_memory, "SharedMemory", tracked_open)
+    with open(f"/dev/shm/shm_{key}_lockfile.lock", "rb+") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        assert receiver.get("0", "1", key) is None
+        assert not calls
+    assert torch.equal(receiver.get("0", "1", key)[0], torch.arange(4096))
+    assert receiver.get("0", "1", key, metadata) is None
+    assert not any(record.levelname == "ERROR" for record in caplog.records)
+
+
+def test_tensor_frames_preserve_shared_output_reconstruction(edge):
+    sender, receiver = edge
+    completion = asdict(CompletionOutput(index=0, text="", token_ids=[1], cumulative_logprob=0.0, logprobs=None))
+    completion["multimodal_output"] = {"codes": torch.arange(28)}
+    payload = {
+        "request_id": "output",
+        "prompt": None,
+        "prompt_token_ids": [1],
+        "prompt_logprobs": None,
+        "outputs": [completion],
+        "finished": False,
+    }
+    _, _, metadata = sender.put("0", "1", "output", payload)
+    assert "host_ring" in metadata
+    result = receiver.get("0", "1", "output", metadata)[0]
+    assert isinstance(result, RequestOutput)
+    assert isinstance(result.outputs[0], CompletionOutput)
+    assert torch.equal(result.outputs[0].multimodal_output["codes"], torch.arange(28))
+
+
+def test_large_tensor_does_not_share_owner_with_small_metadata():
+    with SharedMemoryConnector({}) as connector:
+        payload = {"hidden": torch.ones(16384), "finished": torch.tensor(False)}
+        assert connector.put("0", "1", "large-owner", payload)[0]
+        output = connector.get("0", "1", "large-owner")[0]
+        assert output["finished"].untyped_storage().nbytes() == 1
+        assert output["hidden"].untyped_storage().nbytes() == 65536
+
+
+def test_full_ring_and_large_frame_use_existing_shm(edge):
+    sender, receiver = edge
+    payload = {"codes": torch.arange(56)}
+    metadata = {}
+    for index in range(12):
+        ok, _size, metadata[index] = sender.put("0", "1", f"full{index}", payload)
+        assert ok
+    assert any("host_ring" in item for item in metadata.values())
+    assert any("shm" in item for item in metadata.values())
+    # Unclaimed earlier frames pin the ring; later requests still arrive.
+    for index in reversed(range(12)):
+        assert torch.equal(receiver.get("0", "1", f"full{index}")[0]["codes"], payload["codes"])
+    ok, _size, large_metadata = sender.put("0", "1", "large", torch.arange(4096))
+    assert ok and "shm" in large_metadata
+    assert torch.equal(receiver.get("0", "1", "large", large_metadata)[0], torch.arange(4096))
+    assert "host_ring" in sender.put("0", "1", "recovered", payload)[2]
+    assert receiver.get("0", "1", "recovered") is not None
+
+
+def test_reput_across_ring_and_fallback_and_stale_descriptor(edge):
+    sender, receiver = edge
+    _, _, old = sender.put("0", "1", "reuse", torch.tensor([11]))
+    _, _, current = sender.put("0", "1", "reuse", torch.tensor([22]))
+    assert receiver.get("0", "1", "reuse", old) is None
+    assert receiver.get("0", "1", "reuse", current)[0].item() == 22
+    sender.put("0", "1", "reuse", torch.tensor([33]))
+    sender.put("0", "1", "reuse", torch.full((4096,), 44))
+    assert receiver.get("0", "1", "reuse")[0][0].item() == 44
+    sender.put("0", "1", "reuse", torch.full((4096,), 55))
+    sender.put("0", "1", "reuse", torch.tensor([66]))
+    assert receiver.get("0", "1", "reuse")[0].item() == 66
+    assert not os.path.exists("/dev/shm/reuse")
+
+
+def test_cancel_without_receiver_scan_returns_credit_and_prefix_is_exact(edge):
+    sender, receiver = edge
+    for key in ("req_0_0", "req_0_1", "req_0_10other", "request_0_0"):
+        assert sender.put("0", "1", key, torch.arange(56))[0]
+    assert sender.cleanup_prefix("req_0_") == 2
+    assert receiver.get("0", "1", "req_0_0") is None
+    assert receiver.get("0", "1", "req_0_10other") is not None
+    assert receiver.get("0", "1", "request_0_0") is not None
+    assert "host_ring" in sender.put("0", "1", "next", torch.arange(56))[2]
+
+
+def test_flat_scope_and_metadata_work_across_different_parents(monkeypatch):
+    monkeypatch.setattr(os, "getppid", lambda: 100)
+    sender = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": "flat-scope"})
+    monkeypatch.setattr(os, "getppid", lambda: 200)
+    receiver = SharedMemoryConnector({"stage_id": 1, "wakeup_scope": "flat-scope"})
+    try:
+        assert sender._wake_directory == receiver._wake_directory
+        assert sender.put("0", "1", "scope", torch.tensor([7]))[0]
+        assert receiver.get("0", "1", "scope")[0].item() == 7
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def test_empty_poll_reuses_discovery_and_never_waits_for_a_ring_lock(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "prime-poll", torch.arange(3))
+    assert receiver.get("0", "1", "prime-poll") is not None
+
+    def unexpected_work(*args, **kwargs):
+        pytest.fail("an empty, unchanged channel must not rescan discovery or acquire its frame lock")
+
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", unexpected_work)
+    monkeypatch.setattr(fcntl, "flock", unexpected_work)
+    for index in range(128):
+        assert receiver.get("0", "1", f"missing-poll-{index}") is None
+    monkeypatch.undo()
+    sender.put("0", "1", "new-poll", torch.arange(5))
+    assert receiver.get("0", "1", "new-poll")[0].tolist() == list(range(5))
+
+
+def test_new_producer_is_discovered_after_a_miss_even_when_the_clock_does_not_advance(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "first-producer", torch.arange(3))
+    assert receiver.get("0", "1", "first-producer") is not None
+    assert receiver.get("0", "1", "second-producer") is None
+    old_version = receiver._host_ring._discovery_versions["0", "1"]
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.time.time_ns", lambda: 0)
+    with SharedMemoryConnector(sender.config) as second:
+        second.put("0", "1", "second-producer", torch.arange(5))
+        assert receiver.get("0", "1", "second-producer")[0].tolist() == list(range(5))
+        new_version = receiver._host_ring._discovery_versions["0", "1"]
+        assert new_version[2] == old_version[2] + 1
+
+
+def test_one_receiver_discovers_each_input_edge_at_the_same_registry_version(edge):
+    sender, receiver = edge
+    with SharedMemoryConnector({**sender.config, "stage_id": 2}) as second:
+        sender.put("0", "1", "fanin-first", torch.arange(3))
+        second.put("2", "1", "fanin-second", torch.arange(5))
+        assert receiver.get("0", "1", "fanin-first")[0].tolist() == list(range(3))
+        assert receiver.get("2", "1", "fanin-second")[0].tolist() == list(range(5))
+
+
+def test_discovery_attaches_all_producers_even_when_the_first_one_has_the_requested_key(edge, monkeypatch):
+    sender, receiver = edge
+    with SharedMemoryConnector(sender.config) as second:
+        first_metadata = sender.put("0", "1", "multi-first", torch.arange(3))[2]
+        second_metadata = second.put("0", "1", "multi-second", torch.arange(5))[2]
+        paths = [
+            f"/dev/shm/{receiver._host_ring.directory}/{metadata['host_ring']['name']}"
+            for metadata in (first_metadata, second_metadata)
+        ]
+        monkeypatch.setattr(
+            "vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", lambda _: iter(paths)
+        )
+        assert receiver.get("0", "1", "multi-first")[0].tolist() == list(range(3))
+        assert receiver.get("0", "1", "multi-second")[0].tolist() == list(range(5))
+
+
+def test_closed_channel_releases_reader_cache_while_the_producer_process_is_still_alive(edge):
+    sender, receiver = edge
+    sender.put("0", "1", "close-cache", torch.arange(3))
+    assert receiver.get("0", "1", "close-cache") is not None
+    channel = next(iter(receiver._host_ring.readers.values()))
+    fd = channel.fd
+    sender.close()
+    assert receiver.get("0", "1", "not-published-after-close") is None
+    assert not receiver._host_ring.readers
+    with pytest.raises(OSError):
+        os.fstat(fd)
+
+
+def test_disabled_ring_and_unsupported_tree_keep_legacy_serializer(edge):
+    sender, receiver = edge
+    with SharedMemoryConnector({"host_ring_bytes": 0}) as legacy:
+        _, _, metadata = legacy.put("0", "1", "legacy", {"codes": torch.arange(3)})
+        assert "shm" in metadata
+        assert receiver.get("0", "1", "legacy", metadata)[0]["codes"].tolist() == [0, 1, 2]
+    assert prepare_tensor_frame({"codes": torch.arange(3), "opaque": object()}) is None
+    assert sender.put("0", "1", "falsey", False)[0]
+    assert receiver.get("0", "1", "falsey")[0] is False
+
+
+def test_tensor_free_lists_reuse_native_serialization():
+    payload = {"prompt_token_ids": list(range(8192)), "condition": [0.25] * 8192}
+    with SharedMemoryConnector({}) as connector:
+        frame = prepare_tensor_frame(payload)
+        assert isinstance(frame, bytes)
+        assert frame == connector.serialize_obj(payload)
+        assert connector.put("0", "1", "native-lists", payload)[0]
+        assert connector.get("0", "1", "native-lists")[0] == payload
+
+
+def test_scalar_lists_and_literal_extensions_preserve_wire_semantics():
+    literal = msgspec.msgpack.Ext(42, bytes(20))
+    payload = {"condition": [0.25] * 8192, "codes": torch.arange(28), "extension": literal}
+    with SharedMemoryConnector({}) as connector:
+        _, _, metadata = connector.put("0", "1", "mixed-lists", payload)
+        assert "host_ring" in metadata
+        result = connector.get("0", "1", "mixed-lists")[0]
+        assert result["condition"] == payload["condition"]
+        assert torch.equal(result["codes"], payload["codes"])
+        assert result["extension"].code == literal.code
+        assert result["extension"].data == literal.data
+
+
+def _locked_writer(name, ready, release):
+    fd = os.open(f"/dev/shm/{name}", os.O_RDWR)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        ready.set()
+        assert release.wait(10)
+    finally:
+        os.close(fd)
+
+
+@pytest.mark.parametrize("use_metadata", [False, True])
+def test_deadline_receive_never_waits_for_another_process_ring_lock(edge, use_metadata):
+    sender, receiver = edge
+    _, _, metadata = sender.put("0", "1", "locked", torch.arange(3))
+    ctx = multiprocessing.get_context("spawn")
+    ready, release = ctx.Event(), ctx.Event()
+    writer = ctx.Process(target=_locked_writer, args=(metadata["host_ring"]["name"], ready, release))
+    writer.start()
+    try:
+        assert ready.wait(15)
+        start = time.monotonic()
+        assert (
+            receiver.get_with_deadline("0", "1", "locked", metadata if use_metadata else None, deadline=start + 1)
+            is None
+        )
+        assert time.monotonic() - start < 1
+    finally:
+        release.set()
+        writer.join(15)
+        if writer.is_alive():
+            writer.kill()
+            writer.join()
+    assert writer.exitcode == 0
+    assert receiver.get("0", "1", "locked", metadata if use_metadata else None)[0].tolist() == [0, 1, 2]
+
+
+def _competing_reader(scope, barrier, results):
+    with SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope}) as receiver:
+        barrier.wait(timeout=20)
+        result = receiver.get("0", "1", "shared")
+        results.put(None if result is None else result[0]["codes"].tolist())
+
+
+def test_two_reader_processes_claim_exactly_once_and_close_does_not_unlink_owner():
+    scope = uuid.uuid4().hex
+    ctx = multiprocessing.get_context("spawn")
+    barrier, results = ctx.Barrier(3), ctx.Queue()
+    with SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope}) as sender:
+        _, _, metadata = sender.put("0", "1", "shared", {"codes": torch.arange(56)})
+        readers = [ctx.Process(target=_competing_reader, args=(scope, barrier, results)) for _ in range(2)]
+        for reader in readers:
+            reader.start()
+        try:
+            barrier.wait(timeout=30)
+            outputs = [results.get(timeout=20) for _ in readers]
+            assert outputs.count(None) == 1
+            assert next(output for output in outputs if output is not None) == list(range(56))
+            assert os.path.exists(f"/dev/shm/{metadata['host_ring']['name']}")
+        finally:
+            for reader in readers:
+                reader.join(15)
+                if reader.is_alive():
+                    reader.kill()
+                    reader.join()
+            results.close()
+        assert all(reader.exitcode == 0 for reader in readers)
+
+
+def test_ring_setup_failure_keeps_connector_usable(edge, monkeypatch):
+    sender, receiver = edge
+
+    def no_space(*args):
+        raise OSError(28, "No space left on device")
+
+    monkeypatch.setattr(os, "posix_fallocate", no_space)
+    ok, _, metadata = sender.put("0", "1", "no-ring-space", torch.arange(3))
+    assert ok and "shm" in metadata
+    assert receiver.get("0", "1", "no-ring-space")[0].tolist() == [0, 1, 2]
+
+
+def test_discovery_stamp_failure_keeps_fallback_and_cleanup_usable(edge, monkeypatch):
+    sender, receiver = edge
+
+    def stamp_denied(*args, **kwargs):
+        raise PermissionError("directory timestamp update denied")
+
+    monkeypatch.setattr(os, "utime", stamp_denied)
+    ok, _, metadata = sender.put("0", "1", "no-ring-stamp", torch.arange(3))
+    assert ok and "shm" in metadata
+    assert not sender._host_ring.producers
+    assert not list(Path(f"/dev/shm/{sender._host_ring.directory}").iterdir())
+    assert receiver.get("0", "1", "no-ring-stamp")[0].tolist() == [0, 1, 2]
+    sender.close()
+
+
+def test_close_is_idempotent_and_removes_owned_ring(edge):
+    sender, receiver = edge
+    _, _, metadata = sender.put("0", "1", "close", torch.arange(3))
+    receiver.get("0", "1", "close")
+    sender.close()
+    sender.close()
+    receiver.close()
+    receiver.close()
+    assert not os.path.exists(f"/dev/shm/{metadata['host_ring']['name']}")
+
+
+def _orphan_producer(scope, pipe):
+    producer = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope})
+    metadata = producer.put("0", "1", "restart", torch.tensor([11]))[2]
+    pipe.send(metadata)
+    pipe.recv()
+    pipe.close()
+    os._exit(0)
+
+
+def test_dead_producer_payload_cannot_reach_a_reused_key():
+    scope = uuid.uuid4().hex
+    ctx = multiprocessing.get_context("spawn")
+    parent, child = ctx.Pipe()
+    process = ctx.Process(target=_orphan_producer, args=(scope, child))
+    process.start()
+    child.close()
+    old = parent.recv()
+    try:
+        with SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope}) as receiver:
+            # Cache the live channel without claiming the pending frame.
+            assert receiver.get("0", "1", "not-published") is None
+            channel = next(iter(receiver._host_ring.readers.values()))
+            fd = channel.fd
+            parent.send("exit")
+            process.join(15)
+            assert process.exitcode == 0
+            assert receiver.get("0", "1", "restart", old) is None
+            assert not receiver._host_ring.readers
+            with pytest.raises(OSError):
+                os.fstat(fd)
+            with SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope}) as replacement:
+                replacement.put("0", "1", "restart", torch.tensor([22]))
+                assert receiver.get("0", "1", "restart")[0].item() == 22
+    finally:
+        if process.is_alive():
+            process.kill()
+            process.join()
+        # The multiprocessing parent still owns this test's resource tracker.
+        # Retire only the intentionally orphaned allocation and its exact marker.
+        segment = shared_memory.SharedMemory(name=old["host_ring"]["name"])
+        segment.close()
+        segment.unlink()
+        directory = old["host_ring"]["name"].rsplit("_", 4)[0]
+        try:
+            os.unlink(f"/dev/shm/{directory}/{old['host_ring']['name']}")
+        except FileNotFoundError:
+            pass
+        try:
+            os.rmdir(f"/dev/shm/{directory}")
+        except FileNotFoundError:
+            pass
+        parent.close()

tests/entrypoints/openai_api/test_omni_sleep_wakeup.py

@@ -125,10 +125,9 @@ def test_sleep_client_error_does_not_mark_stages(sleep_capable_engine, mocker):
     app = _make_app(sleep_capable_engine)
     client = TestClient(app)
 
-    response = client.post("/v1/omni/sleep", json={"stage_ids": [9], "level": 1})
+    with pytest.raises(ValueError, match="unknown stage id 9"):
+        client.post("/v1/omni/sleep", json={"stage_ids": [9], "level": 1})
 
-    assert response.status_code == 400
-    assert "unknown stage id 9" in response.json()["detail"]
     assert app.state.sleeping_stages == set()
 
 
@@ -214,19 +213,6 @@ def test_wakeup_failure_returns_error_and_keeps_sleeping_set(sleep_capable_engin
     assert app.state.sleeping_stages == {0}
 
 
-def test_wakeup_client_error_returns_400_and_keeps_sleeping_set(sleep_capable_engine, mocker):
-    sleep_capable_engine.wake_up = mocker.AsyncMock(side_effect=ValueError("unknown stage id 9"))
-    app = _make_app(sleep_capable_engine)
-    app.state.sleeping_stages = {0, 9}
-    client = TestClient(app)
-
-    response = client.post("/v1/omni/wakeup", json={"stage_ids": [9]})
-
-    assert response.status_code == 400
-    assert "unknown stage id 9" in response.json()["detail"]
-    assert app.state.sleeping_stages == {0, 9}
-
-
 def test_wakeup_engine_not_support(sleep_incapable_engine):
     app = _make_app(sleep_incapable_engine)
     app.state.sleeping_stages = {0}

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -353,28 +353,6 @@ def test_sleep_waits_for_in_flight_generate_admission():
     asyncio.run(run())
 
 
-@pytest.mark.cpu
-def test_sleep_rejects_invalid_stage_ids_without_blocking_admission():
-    """Invalid stage_ids must raise before _paused is set.
-
-    Raising afterwards leaves no sleeping tags behind, so wake_up() early-returns
-    "already warm" and nothing ever clears _paused: generate() hangs forever.
-    """
-
-    async def run() -> None:
-        omni = _make_omni(stage_types=["llm", "diffusion"])
-
-        with pytest.raises(ValueError, match=r"Invalid stage_ids \[99\]"):
-            await omni.sleep(stage_ids=[99], level=1, mode="abort")
-
-        assert omni._paused is False
-        assert omni._hold_admission_until_resume is False
-        omni._clear_frontend_mm_cache.assert_not_awaited()
-        omni.collective_rpc.assert_not_awaited()
-
-    asyncio.run(run())
-
-
 @pytest.mark.cpu
 def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
     async def run() -> None:

tests/model_executor/models/minicpmo_4_5/test_encoder_cuda_graph.py

@@ -1,379 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import pytest
-import torch
-
-from vllm_omni.model_executor.models.minicpmo_4_5.encoder_cuda_graph import EncoderCudaGraph
-from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
-    SiglipVisionConfig,
-    SiglipVisionTransformer,
-)
-from vllm_omni.platforms import current_omni_platform
-
-pytestmark = [pytest.mark.core_model]
-
-
-def _make_graph(forward, **kwargs):
-    from vllm.config import ModelConfig, VllmConfig
-    from vllm.config.multimodal import MultiModalConfig
-
-    config = VllmConfig()
-    # The manager only needs multimodal limits; no checkpoint is loaded here.
-    config.model_config = ModelConfig.__new__(ModelConfig)
-    config.model_config.multimodal_config = MultiModalConfig()
-    return EncoderCudaGraph(forward, config, **kwargs)
-
-
-@pytest.mark.cpu
-def test_cpu_and_grad_paths_remain_eager():
-    graph = _make_graph(lambda x, mask: x.sin() if mask is None else x.sin() + mask)
-    x = torch.randn(2, 3, requires_grad=True)
-    graph(x, None).sum().backward()
-    torch.testing.assert_close(x.grad, x.detach().cos())
-    assert not graph.graphs
-    assert not graph._seen
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_replay_refreshes_inputs_preserves_outputs_and_bounds_shapes():
-    graph = _make_graph(lambda x, mask: x.sin() if mask is None else x.sin() + mask, max_graphs=2)
-    x = torch.randn(2, 8, device="cuda")
-    mask = torch.randn_like(x)
-    graph(x, mask)
-    old = graph(x, mask)
-    expected_old = old.clone()
-    x.mul_(2)
-    mask.add_(3)
-    torch.testing.assert_close(graph(x, mask), x.sin() + mask)
-    torch.testing.assert_close(old, expected_old)
-    graph(x, None)
-    torch.testing.assert_close(graph(x, None), x.sin())
-    assert len(graph.graphs) == 2
-    from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
-
-    assert all(isinstance(entry, EncoderCudaGraphManager) for entry in graph.graphs.values())
-    assert sum(entry.graph_hits for entry in graph.graphs.values()) == 3
-    assert graph.vllm_config.compilation_config.encoder_cudagraph_token_budgets == []
-    for size in range(3, 24):
-        other = torch.randn(size, 8, device="cuda")
-        torch.testing.assert_close(graph(other, None), other.sin())
-    assert len(graph.graphs) == 2
-    assert len(graph._seen) <= 8
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_vision_graph_handles_changed_mask_and_retained_embeddings():
-    config = SiglipVisionConfig(
-        hidden_size=64,
-        intermediate_size=128,
-        num_hidden_layers=2,
-        num_attention_heads=4,
-        image_size=28,
-        patch_size=14,
-    )
-    config._attn_implementation = "eager"
-    model = SiglipVisionTransformer(config).eval().cuda()
-    pixels = torch.randn(2, 3, 14, 56, device="cuda")
-    sizes = torch.tensor([[2, 2], [1, 2]], dtype=torch.int32)
-    mask = torch.tensor([[[1, 1, 1, 1]], [[1, 1, 0, 0]]], device="cuda", dtype=torch.bool)
-    reference = model(pixels, mask, sizes).last_hidden_state
-    model._encoder_graph = _make_graph(model._encode_last_hidden_state)
-    model(pixels, mask, sizes)
-    output = model(pixels, mask, sizes).last_hidden_state
-    torch.testing.assert_close(output, reference)
-    assert len(model._encoder_graph.graphs) == 1
-    # Same padded shape, different valid positions and image content.
-    pixels.add_(1)
-    sizes[1] = torch.tensor([1, 3])
-    mask[1, 0, 2] = True
-    actual = model(pixels, mask, sizes).last_hidden_state
-    graph = model._encoder_graph
-    model._encoder_graph = None
-    expected = model(pixels, mask, sizes).last_hidden_state
-    torch.testing.assert_close(actual, expected)
-    torch.testing.assert_close(output, reference)
-    model._encoder_graph = graph
-    detailed = model(pixels, mask, sizes, output_hidden_states=True)
-    assert len(detailed.hidden_states) == 3
-    assert len(graph.graphs) == 1
-
-
-def _audio_model(device):
-    from types import SimpleNamespace
-
-    from transformers import WhisperConfig
-
-    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
-        MiniCPMO45OmniLLMForConditionalGeneration,
-        MiniCPMWhisperEncoder,
-        MultiModalProjector,
-    )
-
-    config = WhisperConfig(
-        d_model=32,
-        encoder_layers=2,
-        encoder_attention_heads=4,
-        encoder_ffn_dim=64,
-        num_mel_bins=80,
-        max_source_positions=1500,
-    )
-    config._attn_implementation = "sdpa"
-    model = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration)
-    torch.nn.Module.__init__(model)
-    model.config = SimpleNamespace(audio_chunk_length=1.0, audio_pool_step=5)
-    model.apm = MiniCPMWhisperEncoder(config)
-    model.audio_projection_layer = MultiModalProjector(32, 48)
-    model.audio_avg_pooler = torch.nn.AvgPool1d(5, stride=5)
-    model.audio_encoder_layer = -1
-    model.audio_past_key_values = None
-    return model.eval().to(device)
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_audio_graph_refreshes_mask_and_matches_eager():
-    model = _audio_model("cuda")
-    data = {
-        "audio_features": torch.randn(2, 80, 100, device="cuda"),
-        "audio_feature_lens": torch.tensor([[100], [80]], device="cuda"),
-    }
-    expected = model.get_audio_hidden_states(data)
-    graph = _make_graph(model._encode_audio_features)
-    model._audio_encoder_graph = graph
-    model.get_audio_hidden_states(data)
-    actual = model.get_audio_hidden_states(data)
-    for a, b in zip(actual, expected, strict=True):
-        torch.testing.assert_close(a, b)
-    assert len(graph.graphs) == 1
-    data["audio_features"].add_(1)
-    data["audio_feature_lens"][1, 0] = 30
-    changed = model.get_audio_hidden_states(data)
-    model._audio_encoder_graph = None
-    for a, b in zip(changed, model.get_audio_hidden_states(data), strict=True):
-        torch.testing.assert_close(a, b)
-    for a, b in zip(actual, expected, strict=True):
-        torch.testing.assert_close(a, b)
-
-
-@pytest.mark.cpu
-@torch.inference_mode()
-def test_streaming_audio_does_not_replay_stateless_graph():
-    model = _audio_model("cpu")
-
-    def forbidden(*args):
-        raise AssertionError("streaming KV must not use the stateless graph")
-
-    model._audio_encoder_graph = forbidden
-    data = {"audio_features": torch.randn(1, 80, 100), "audio_feature_lens": torch.tensor([[100]])}
-    for _ in range(2):
-        output = model.get_audio_embedding_streaming(data)
-        assert output[0][0].shape == (10, 48)
-    assert model.audio_past_key_values.get_seq_length() == 100
-
-
-@pytest.mark.cpu
-@torch.inference_mode()
-def test_fp16_audio_keeps_host_overflow_check_eager():
-    model = _audio_model("cpu").half()
-
-    def forbidden(*args):
-        raise AssertionError("FP16 overflow guard cannot run inside capture")
-
-    model._audio_encoder_graph = forbidden
-    data = {
-        "audio_features": torch.randn(1, 80, 100, dtype=torch.float16),
-        "audio_feature_lens": torch.tensor([[100]]),
-    }
-    assert model.get_audio_hidden_states(data)[0].shape == (10, 48)
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_one_off_shapes_do_not_exhaust_capture_admission():
-    graph = _make_graph(torch.sin, max_graphs=1)
-    for size in range(1, 20):
-        graph(torch.zeros(size, device="cuda"))
-    assert len(graph._seen) == 4
-    assert not graph.graphs
-    x = torch.randn(32, device="cuda")
-    graph(x)
-    torch.testing.assert_close(graph(x), x.sin())
-    assert len(graph.graphs) == 1
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_managers_do_not_share_buffers_between_streams():
-    graph = _make_graph(torch.sin)
-    streams = [torch.cuda.Stream(), torch.cuda.Stream()]
-    outputs = []
-    for i, stream in enumerate(streams):
-        with torch.cuda.stream(stream):
-            value = torch.full((4, 8), float(i + 1), device="cuda")
-            graph(value)
-            output = graph(value)
-            graph(value + 1)
-            outputs.append((output, value.sin()))
-    for stream in streams:
-        torch.cuda.current_stream().wait_stream(stream)
-    assert len(graph.graphs) == 2
-    for actual, expected in outputs:
-        torch.testing.assert_close(actual, expected)
-    assert not graph.vllm_config.compilation_config.encoder_cudagraph_token_budgets
-
-
-@pytest.mark.cpu
-@pytest.mark.parametrize(
-    "options",
-    [
-        {"max_graphs": -1},
-        {"max_graphs": True},
-        {"min_capture_calls": 1},
-        {"min_capture_calls": 2.5},
-        {"min_free_bytes": -1},
-    ],
-)
-def test_invalid_admission_options(options):
-    with pytest.raises(ValueError, match="must be an integer"):
-        _make_graph(torch.sin, **options)
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_configurable_admission_and_capacity_miss_accounting():
-    graph = _make_graph(torch.sin, max_graphs=1, min_capture_calls=3)
-    x = torch.randn(2, 8, device="cuda")
-    for _ in range(2):
-        torch.testing.assert_close(graph(x), x.sin())
-        assert not graph.graphs
-    torch.testing.assert_close(graph(x), x.sin())
-    torch.testing.assert_close(graph(x + 1), (x + 1).sin())
-    for _ in range(5):
-        other = torch.randn(3, 8, device="cuda")
-        torch.testing.assert_close(graph(other), other.sin())
-    stats = graph.get_cumulative_stats()
-    assert stats["num_graphs"] == 1
-    assert stats["graph_hits"] == 2
-    assert stats["graph_misses"] == 7
-    assert stats["capacity_misses"] == 5
-    assert stats["warmup_misses"] == 2
-    assert stats["hit_rate"] == pytest.approx(2 / 9)
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_memory_admission_defers_capture_but_preserves_existing_replay(monkeypatch):
-    graph = _make_graph(torch.sin, min_free_bytes=1024)
-    x = torch.randn(2, 8, device="cuda")
-    graph(x)
-    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda device: 1023)
-    for _ in range(3):
-        torch.testing.assert_close(graph(x), x.sin())
-    assert not graph.graphs
-    assert graph.get_cumulative_stats()["memory_misses"] == 3
-    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda device: 1024)
-    old = graph(x)
-    assert len(graph.graphs) == 1
-    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda device: 0)
-    torch.testing.assert_close(graph(x + 1), (x + 1).sin())
-    torch.testing.assert_close(old, x.sin())
-    assert graph.get_cumulative_stats()["graph_hits"] == 2
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_capture_failure_is_fatal_without_retry_or_eager_fallback(monkeypatch):
-    from unittest.mock import Mock
-
-    forward = Mock(side_effect=torch.sin)
-    graph = _make_graph(forward)
-    x = torch.randn(2, 8, device="cuda")
-    graph(x)
-    capture = Mock(side_effect=RuntimeError("injected capture failure"))
-    monkeypatch.setattr(graph, "_capture", capture)
-    with pytest.raises(RuntimeError, match="injected capture failure"):
-        graph(x)
-    for value in (x, torch.randn(3, 8, device="cuda"), torch.randn(2, 8)):
-        with pytest.raises(RuntimeError, match="restart the worker"):
-            graph(value)
-    capture.assert_called_once()
-    forward.assert_called_once()
-    assert not graph.graphs
-    assert not graph._seen
-    assert not graph._seen_calls
-    assert graph.get_cumulative_stats()["capture_failures"] == 1
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_zero_capacity_disables_capture():
-    graph = _make_graph(torch.sin, max_graphs=0)
-    x = torch.randn(2, 8, device="cuda")
-    for _ in range(3):
-        torch.testing.assert_close(graph(x), x.sin())
-    assert not graph.graphs
-    assert not graph._seen
-    assert graph.get_cumulative_stats()["capacity_misses"] == 3
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.parametrize("share_pools", [False, True])
-@torch.inference_mode()
-def test_pool_reuse_preserves_outputs_across_late_capture_and_arbitrary_replay(share_pools):
-    graph = _make_graph(lambda x: (x.sin() + 2).cos(), share_pools=share_pools)
-    retained = []
-    for size in (64, 128, 32, 96):
-        x = torch.randn(size, 64, device="cuda")
-        graph(x)
-        retained.append((graph(x), (x.sin() + 2).cos()))
-        # Replay older shapes between captures, including smaller -> larger.
-        for previous in (64, size, 64):
-            value = torch.randn(previous, 64, device="cuda")
-            torch.testing.assert_close(graph(value), (value.sin() + 2).cos())
-        for actual, expected in retained:
-            torch.testing.assert_close(actual, expected)
-    pools = {entry.graph_pool for entry in graph.graphs.values()}
-    assert len(pools) == (1 if share_pools else 4)
-    assert len(graph._capture_pools) == (1 if share_pools else 0)
-    for size in (96, 32, 128, 64, 128, 32):
-        value = torch.randn(size, 64, device="cuda")
-        torch.testing.assert_close(graph(value), (value.sin() + 2).cos())
-    for actual, expected in retained:
-        torch.testing.assert_close(actual, expected)
-
-
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@torch.inference_mode()
-def test_shared_pools_are_isolated_by_encoder_and_replay_stream():
-    graphs = [_make_graph(torch.sin), _make_graph(torch.cos)]
-    streams = [torch.cuda.Stream(), torch.cuda.Stream()]
-    retained = []
-    for graph in graphs:
-        for stream in streams:
-            with torch.cuda.stream(stream):
-                for size in (16, 32):
-                    x = torch.randn(size, 64, device="cuda")
-                    graph(x)
-                    retained.append((graph(x), graph.forward(x)))
-                    graph(x + 1)
-    for stream in streams:
-        torch.cuda.current_stream().wait_stream(stream)
-    pools = {entry.graph_pool for graph in graphs for entry in graph.graphs.values()}
-    assert len(pools) == 4
-    assert all(len(graph._capture_pools) == 2 for graph in graphs)
-    for actual, expected in retained:
-        torch.testing.assert_close(actual, expected)

tests/model_executor/models/minicpmo_4_5/test_vision_batching.py

@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 from dataclasses import dataclass
 
 import pytest
@@ -100,8 +97,6 @@ class _FakeVisionModel:
 
 
 class _FakeAudioModel:
-    _encode_audio_features = MiniCPMO45OmniLLMForConditionalGeneration._encode_audio_features
-
     def __init__(self, audio_encoder: _FakeAudioEncoder, audio_encoder_layer: int) -> None:
         self.config = _AudioConfig()
         self.apm = audio_encoder

tests/model_tests/diffusion/diff_model_builders.py

@@ -278,20 +278,6 @@ def tiny_sd3_builder() -> str:
     )
 
 
-def tiny_sdxl_builder() -> str:
-    # The UNet's architecture (block_out_channels, transformer_layers_per_block, etc.)
-    # is hardcoded in SDXLUNet2DConditionModel rather than read from config, so it
-    # can't be shrunk here; only the two CLIP text encoders are reduced.
-    return build_tiny_from_configs(
-        "StableDiffusionXLPipeline",
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        transform={
-            "text_encoder": _shrink_flux_clip_text_encoder,
-            "text_encoder_2": _shrink_flux_clip_text_encoder,
-        },
-    )
-
-
 def tiny_flux2_builder() -> str:
     def shrink_text_encoder(config: dict) -> dict:
         config["tie_word_embeddings"] = False

tests/model_tests/diffusion/model_settings.py

@@ -181,10 +181,4 @@ DIFFUSION_TEST_SETTINGS = {
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
         model_type_marker=ModelTypeMarker.DIFFUSION,
     ),
-    "StableDiffusionXLPipeline": DiffusionModelTestOpts(
-        model="stabilityai/stable-diffusion-xl-base-1.0",
-        builder=diff_model_builders.tiny_sdxl_builder,
-        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
-        extra_test_groups=[[DiffusionAccs.CPU_OFFLOAD]],
-    ),
 }

tests/model_tests/diffusion/test_alignment.py

@@ -84,6 +84,7 @@ EXCLUDED_MODELS = [
     "HiDreamImagePipeline",
     "HiDreamO1ImagePipeline",
     "DreamZeroPipeline",
+    "StableDiffusionXLPipeline",
     "Gr00tN1d7Pipeline",
     "Pi0Pipeline",
     "Pi05Pipeline",

vllm_omni/distributed/omni_connectors/connectors/base.py

@@ -56,8 +56,8 @@ class OmniConnectorBase(ABC):
         pass
 
     @abstractmethod
-    def cleanup(self, request_id: str) -> None:
-        """Clean up resources for a request."""
+    def cleanup(self, request_id: str) -> bool | None:
+        """Clean up resources; backends may report whether a payload was retired."""
         pass
 
     def get_with_deadline(

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -17,7 +17,10 @@ from typing import Any
 from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes
 
 from ..utils.logging import get_connector_logger
+from ..utils.serialization import OmniSerializer
+from ..utils.tensor_frame import TensorFrame, prepare_tensor_frame
 from .base import OmniConnectorBase
+from .shm_ring import HostRingTransport
 
 logger = get_connector_logger(__name__)
 
@@ -55,7 +58,8 @@ class SharedMemoryConnector(OmniConnectorBase):
     def __init__(self, config: dict[str, Any]):
         self.config = config
         self.stage_id = config.get("stage_id", -1)
-        scope = config.get("extra", {}).get("wakeup_scope")
+        extra = config.get("extra", config)
+        scope = extra.get("wakeup_scope")
         self._wake_directory = (
             f"/dev/shm/omni_shm_wake_{os.getuid()}_{hashlib.sha256(str(scope).encode()).hexdigest()[:24]}"
             if scope is not None
@@ -67,7 +71,13 @@ class SharedMemoryConnector(OmniConnectorBase):
             "puts": 0,
             "gets": 0,
             "bytes_transferred": 0,
+            "host_ring_puts": 0,
+            "host_ring_gets": 0,
+            "host_ring_fallbacks": 0,
         }
+        # One bounded host fast path per edge, allocated only on the first put.
+        # A zero capacity preserves the per-key protocol for controlled comparisons.
+        self._host_ring = HostRingTransport(self._wake_directory, int(extra.get("host_ring_bytes", 4 << 20)))
         # Receiver side: FIFO read end (plus a write end of our own, so the
         # FIFO never reports EOF when no sender has it open) and a counter of
         # drained wakeups. Each receiver owns its path, including across restarts.
@@ -200,8 +210,32 @@ class SharedMemoryConnector(OmniConnectorBase):
         data: Any,
     ) -> tuple[bool, int, dict[str, Any] | None]:
         try:
-            self.reap_consumed()
-            payload = self.serialize_obj(data)
+            # Ring put() reclaims its own edge under the publication lock.
+            # Keep the per-key sweep here without taking every ring lock twice.
+            self._reap_segments()
+            use_ring = self._host_ring.capacity and str(from_stage).isdigit() and str(to_stage).isdigit()
+            frame = (
+                prepare_tensor_frame(data, max_bytes=self._host_ring.capacity // 4 - 16 - len(put_key.encode()))
+                if use_ring
+                else None
+            )
+            payload: TensorFrame | bytes = frame if frame is not None else self.serialize_obj(data)
+            ring_meta = self._host_ring.put(str(from_stage), str(to_stage), put_key, payload)
+            if ring_meta is not None:
+                # A previous put may have used the per-key fallback.
+                with self._pending_keys_lock:
+                    had_segment = put_key in self._pending_keys
+                if had_segment:
+                    self._cleanup_segment(put_key)
+                size = ring_meta["size"]
+                self._wake_receiver(to_stage)
+                self._metrics["puts"] += 1
+                self._metrics["host_ring_puts"] += 1
+                self._metrics["bytes_transferred"] += size
+                return True, size, {"host_ring": ring_meta, "size": size}
+            self._metrics["host_ring_fallbacks"] += 1
+            if isinstance(payload, TensorFrame):
+                payload = self.serialize_obj(data)
             size = len(payload)
 
             lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock"
@@ -214,6 +248,8 @@ class SharedMemoryConnector(OmniConnectorBase):
             metadata = {"shm": meta, "size": size}
             with self._pending_keys_lock:
                 self._pending_keys[put_key] = None
+            # A successful fallback replaces an unclaimed ring version, too.
+            self._host_ring.cancel(key=put_key)
             self._wake_receiver(to_stage)
 
             self._metrics["puts"] += 1
@@ -236,7 +272,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             obj = self.deserialize_obj(data_bytes)
             result = (obj, int(shm_handle.get("size", 0)))
             return result
-        except BlockingIOError:
+        except (BlockingIOError, FileNotFoundError):
             return None
         except Exception as e:
             logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
@@ -250,13 +286,15 @@ class SharedMemoryConnector(OmniConnectorBase):
 
     def _get_by_key(self, get_key: str) -> tuple[Any, int] | None:
         """Read a SHM segment addressed purely by *get_key*."""
-        shm = None
         try:
-            shm = shm_pkg.SharedMemory(name=get_key)
-            if shm is None or shm.size == 0:
+            # Inspect size without registering a reader before it owns the
+            # delivery lock. Competing readers must not re-register an already
+            # consumed allocation with the application's resource tracker.
+            size = os.stat(f"/dev/shm/{get_key}").st_size
+            if size == 0:
                 return None
             lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock"
-            shm_handle = {"name": get_key, "size": shm.size}
+            shm_handle = {"name": get_key, "size": size}
             result = self._get_data_with_lock(lock_file, shm_handle)
             if result is not None:
                 with self._pending_keys_lock:
@@ -264,20 +302,9 @@ class SharedMemoryConnector(OmniConnectorBase):
             return result
         except FileNotFoundError:
             return None
-        except ValueError as e:
-            # A receiver can observe a newly-created POSIX SHM object before
-            # the writer has finished sizing it. Treat that as "not ready yet"
-            # so async polling can retry without a traceback.
-            if "empty file" in str(e):
-                return None
-            logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True)
-            return None
         except Exception:
             logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True)
             return None
-        finally:
-            if shm:
-                shm.close()
 
     def get(
         self,
@@ -286,6 +313,23 @@ class SharedMemoryConnector(OmniConnectorBase):
         get_key: str,
         metadata=None,
     ) -> tuple[Any, int] | None:
+        if isinstance(metadata, dict) and get_key in metadata:
+            metadata = metadata[get_key]
+        ring_meta = metadata.get("host_ring") if isinstance(metadata, dict) else None
+        ring_result = self._host_ring.get(str(from_stage), str(to_stage), get_key, ring_meta)
+        if ring_result is not None:
+            obj, size = ring_result
+            if isinstance(obj, bytes):
+                obj = self.deserialize_obj(obj)
+            else:
+                # Preserve the centralized serializer's output reconstruction,
+                # including RequestOutput/CompletionOutput dictionaries.
+                obj = OmniSerializer.restore(obj)
+            self._metrics["gets"] += 1
+            self._metrics["host_ring_gets"] += 1
+            return obj, size
+        if ring_meta is not None:
+            return None
         if metadata is not None:
             if isinstance(metadata, dict) and get_key in metadata:
                 metadata = metadata.get(get_key)
@@ -312,7 +356,10 @@ class SharedMemoryConnector(OmniConnectorBase):
 
         Returns True when an unconsumed segment was actually unlinked.
         """
-        key = request_id
+        ring_cancelled = self._host_ring.cancel(key=request_id)
+        return self._cleanup_segment(request_id) or bool(ring_cancelled)
+
+    def _cleanup_segment(self, key: str) -> bool:
         unlinked = False
         with self._pending_keys_lock:
             self._pending_keys.pop(key, None)
@@ -343,10 +390,14 @@ class SharedMemoryConnector(OmniConnectorBase):
         """
         with self._pending_keys_lock:
             keys = [k for k in self._pending_keys if k.startswith(key_prefix) and k[len(key_prefix) :].isdigit()]
-        return sum(self.cleanup(key) for key in keys)
+        return sum(self._cleanup_segment(key) for key in keys) + self._host_ring.cancel(prefix=key_prefix)
 
     def close(self) -> None:
         """Unlink all remaining tracked SHM segments."""
+        # A constructor failure may be followed by the base destructor.
+        if not hasattr(self, "_host_ring"):
+            return
+        self._host_ring.close()
         with self._pending_keys_lock:
             keys = list(self._pending_keys)
         for key in keys:
@@ -355,6 +406,10 @@ class SharedMemoryConnector(OmniConnectorBase):
 
     def reap_consumed(self) -> None:
         """Bounded round-robin sweep; receivers unlink SHM in another process."""
+        self._host_ring.reap()
+        self._reap_segments()
+
+    def _reap_segments(self) -> None:
         with self._pending_keys_lock:
             for _ in range(min(64, len(self._pending_keys))):
                 key, _ = self._pending_keys.popitem(last=False)

vllm_omni/distributed/omni_connectors/connectors/shm_ring.py

@@ -0,0 +1,489 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded host frames shared by local stage processes.
+
+Each producer/destination edge owns one mmap, with a lock on that same file.
+Readers copy before marking a frame consumed, and all readers participate in
+the same lock. Pressure returns False to the existing per-key SHM path; this
+fast path does not implement end-to-end stream admission or GPU ownership.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import glob
+import hashlib
+import mmap
+import os
+import select
+import struct
+import threading
+import time
+import uuid
+from collections import deque
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from multiprocessing import shared_memory as shm_pkg
+
+from ..utils.logging import get_connector_logger
+from ..utils.tensor_frame import TensorFrame, read_tensor_frame
+
+logger = get_connector_logger(__name__)
+
+_HEADER_SIZE = 64
+_HEADER = struct.Struct("<8sQQQQQ")  # magic, capacity, head, tail, creator pid, process start
+_MAGIC = b"OMSHMR01"
+_FRAME = struct.Struct("<IIBBHI")  # total, payload size, kind, state, key length, reserved
+_ALIGN = 16
+_READY, _DONE = 0, 1
+_BYTES, _TENSORS, _PADDING = 0, 1, 2
+_STATE_OFFSET = 9
+_DEFAULT_BYTES = 4 << 20
+
+
+def _align(size: int) -> int:
+    return (size + _ALIGN - 1) & ~(_ALIGN - 1)
+
+
+def _process_start(pid: int) -> int:
+    with open(f"/proc/{pid}/stat") as file:
+        return int(file.read().rsplit(")", 1)[1].split()[19])
+
+
+@contextmanager
+def _registry_update(directory: str) -> Iterator[int]:
+    """Publish a strictly changing discovery stamp at channel birth/retirement.
+
+    Kernel directory mtimes can coincide within a clock tick. Serialize marker
+    updates and preserve the previous stamp before the filesystem changes it.
+    Receivers can then cache discovery without a timer or a per-message scan.
+    """
+    fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        previous = os.fstat(fd).st_mtime_ns
+        yield fd
+        stamp = max(previous + 1, time.time_ns())
+        os.utime(fd, ns=(stamp, stamp))
+    finally:
+        os.close(fd)
+
+
+@dataclass
+class _Channel:
+    name: str
+    fd: int
+    mapping: mmap.mmap
+    capacity: int
+    owner: bool
+    creator_pid: int
+    creator_start: int
+    segment: shm_pkg.SharedMemory | None = None
+    registry: str | None = None
+    pidfd: int | None = None
+    cursor: int = 0
+    index: dict[str, int] = field(default_factory=dict)
+    positions: deque[tuple[int, str]] = field(default_factory=deque)
+
+    @classmethod
+    def open(cls, name: str, capacity: int | None = None) -> _Channel:
+        owner = capacity is not None
+        segment = None
+        mapping = None
+        pidfd = None
+        fd = -1
+        try:
+            if capacity is not None:
+                # Only the producer registers the allocation with Python's
+                # resource tracker. Readers map the existing file directly.
+                segment = shm_pkg.SharedMemory(name=name, create=True, size=_HEADER_SIZE + capacity)
+            fd = os.open(f"/dev/shm/{name}", os.O_RDWR | os.O_NOFOLLOW)
+            fcntl.flock(fd, fcntl.LOCK_EX | (0 if owner else fcntl.LOCK_NB))
+            try:
+                if capacity is not None:
+                    os.posix_fallocate(fd, 0, _HEADER_SIZE + capacity)
+                size = os.fstat(fd).st_size
+                if size < _HEADER_SIZE:
+                    raise ValueError("incomplete SHM ring header")
+                mapping = mmap.mmap(fd, size)
+                if capacity is not None:
+                    # Fault once per edge, outside the per-message path.
+                    mapping[:] = bytes(size)
+                    _HEADER.pack_into(mapping, 0, _MAGIC, capacity, 0, 0, os.getpid(), _process_start(os.getpid()))
+                magic, actual_capacity, head, tail, creator_pid, creator_start = _HEADER.unpack_from(mapping)
+                if magic != _MAGIC or actual_capacity != size - _HEADER_SIZE or not 0 <= tail <= head:
+                    raise ValueError("invalid SHM ring header")
+                if not owner:
+                    try:
+                        pidfd = os.pidfd_open(creator_pid)
+                    except (AttributeError, OSError):
+                        # Older Linux kernels use the process-start check on receive.
+                        pidfd = None
+                    if _process_start(creator_pid) != creator_start:
+                        raise FileNotFoundError("SHM ring producer was replaced")
+                return cls(
+                    name, fd, mapping, actual_capacity, owner, creator_pid, creator_start, segment=segment, pidfd=pidfd
+                )
+            finally:
+                fcntl.flock(fd, fcntl.LOCK_UN)
+        except BaseException:
+            if mapping is not None:
+                mapping.close()
+            if fd >= 0:
+                os.close(fd)
+            if segment is not None:
+                segment.close()
+                segment.unlink()
+            if pidfd is not None:
+                os.close(pidfd)
+            raise
+
+    @contextmanager
+    def locked(self, *, blocking: bool = True) -> Iterator[None]:
+        fcntl.flock(self.fd, fcntl.LOCK_EX | (0 if blocking else fcntl.LOCK_NB))
+        try:
+            yield
+        finally:
+            fcntl.flock(self.fd, fcntl.LOCK_UN)
+
+    def bounds(self) -> tuple[int, int]:
+        _magic, _capacity, head, tail, _pid, _start = _HEADER.unpack_from(self.mapping)
+        return head, tail
+
+    def producer_alive(self) -> bool:
+        if self.pidfd is not None:
+            return not select.select([self.pidfd], [], [], 0)[0]
+        try:
+            return _process_start(self.creator_pid) == self.creator_start
+        except FileNotFoundError:
+            return False
+
+    def frame(self, position: int) -> tuple[int, int, int, int, int, int]:
+        frame = _FRAME.unpack_from(self.mapping, _HEADER_SIZE + position % self.capacity)
+        total, length, kind, _state, key_length, _reserved = frame
+        if total < _FRAME.size or total % _ALIGN or position % self.capacity + total > self.capacity:
+            raise ValueError("invalid SHM ring frame size")
+        if kind != _PADDING and _FRAME.size + length + key_length > total:
+            raise ValueError("invalid SHM ring payload bounds")
+        return frame
+
+    def scan(self) -> None:
+        head, tail = self.bounds()
+        self.cursor = max(self.cursor, tail)
+        while self.cursor < head:
+            position = self.cursor
+            total, _length, kind, state, key_length, _reserved = self.frame(position)
+            self.cursor += total
+            if kind != _PADDING and state == _READY:
+                start = _HEADER_SIZE + position % self.capacity + _FRAME.size
+                key = self.mapping[start : start + key_length].decode()
+                self.index[key] = position
+                self.positions.append((position, key))
+        # Remove positions retired by another reader or overwritten after wrap.
+        while self.positions and self.positions[0][0] < tail:
+            position, key = self.positions.popleft()
+            if self.index.get(key) == position:
+                del self.index[key]
+
+    def consume(self, position: int) -> None:
+        self.mapping[_HEADER_SIZE + position % self.capacity + _STATE_OFFSET] = _DONE
+
+    def reclaim(self) -> None:
+        head, tail = self.bounds()
+        while tail < head:
+            total, _length, kind, state, _key_length, _reserved = self.frame(tail)
+            if kind != _PADDING and state != _DONE:
+                break
+            tail += total
+        struct.pack_into("<Q", self.mapping, 24, tail)
+
+    def cancel(self, key: str | None = None, prefix: str | None = None) -> int:
+        self.scan()
+        count = 0
+        candidates = (
+            (((key, self.index[key]),) if key in self.index else ()) if key is not None else tuple(self.index.items())
+        )
+        for candidate, position in candidates:
+            matches = candidate == key or (
+                prefix is not None and candidate.startswith(prefix) and candidate[len(prefix) :].isdigit()
+            )
+            if matches:
+                if self.frame(position)[3] == _READY:
+                    self.consume(position)
+                    count += 1
+                del self.index[candidate]
+        self.reclaim()
+        return count
+
+    def close(self) -> None:
+        self.mapping.close()
+        os.close(self.fd)
+        if self.pidfd is not None:
+            os.close(self.pidfd)
+        if self.segment is not None:
+            self.segment.close()
+            try:
+                self.segment.unlink()
+            except FileNotFoundError:
+                pass
+        if self.registry is not None:
+            try:
+                with _registry_update(os.path.dirname(self.registry)) as directory_fd:
+                    os.unlink(os.path.basename(self.registry), dir_fd=directory_fd)
+            except FileNotFoundError:
+                pass
+            except OSError as error:
+                # The allocation is already retired. A failed discovery stamp
+                # must not mask setup failure or prevent the per-key fallback.
+                logger.debug("Failed to retire host ring marker %s: %s", self.registry, error)
+
+
+class HostRingTransport:
+    """Private SharedMemoryConnector fast path; no new model-specific API."""
+
+    def __init__(self, scope: str, capacity: int = _DEFAULT_BYTES):
+        if capacity < 0 or capacity and not 4096 <= capacity <= 64 << 20:
+            raise ValueError("host_ring_bytes must be zero or between 4096 and 67108864")
+        self.capacity = _align(capacity)
+        scope_id = hashlib.sha256(scope.encode()).hexdigest()[:24]
+        self.directory = f"omni_shm_ring_v1_{os.getuid()}_{scope_id}"
+        self.prefix = f"{self.directory}_"
+        self.producers: dict[tuple[str, str], _Channel] = {}
+        self._disabled_edges: set[tuple[str, str]] = set()
+        self.readers: dict[str, _Channel] = {}
+        self._discovery_versions: dict[tuple[str, str], tuple[int, int, int]] = {}
+        self._lock = threading.RLock()
+        self._closed = False
+
+    def put(self, source: str, destination: str, key: str, payload: bytes | TensorFrame) -> dict | None:
+        """Return an immutable descriptor, or None without consuming ring credit."""
+        encoded_key = key.encode()
+        length = payload.size if isinstance(payload, TensorFrame) else len(payload)
+        total = _align(_FRAME.size + len(encoded_key) + length)
+        if not self.capacity or not source.isdigit() or not destination.isdigit():
+            return None
+        # Keep one exceptional payload from monopolizing a channel.
+        if total > self.capacity // 4 or len(encoded_key) > 65535:
+            return None
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("host ring is closed")
+            if (source, destination) in self._disabled_edges:
+                return None
+            channel = self.producers.get((source, destination))
+            if channel is None:
+                name = f"{self.prefix}{source}_{destination}_{os.getpid()}_{uuid.uuid4().hex}"
+                try:
+                    os.makedirs(f"/dev/shm/{self.directory}", mode=0o700, exist_ok=True)
+                    channel = _Channel.open(name, self.capacity)
+                    channel.registry = f"/dev/shm/{self.directory}/{name}"
+                    # Discovery is confined to this deployment directory. The
+                    # marker appears only after the allocation is ready.
+                    with _registry_update(os.path.dirname(channel.registry)) as directory_fd:
+                        marker = os.open(name, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600, dir_fd=directory_fd)
+                        os.close(marker)
+                except OSError as error:
+                    if channel is not None:
+                        channel.close()
+                    self._disabled_edges.add((source, destination))
+                    logger.warning(
+                        "Host ring unavailable on edge %s -> %s; using per-key SHM: %s", source, destination, error
+                    )
+                    return None
+                self.producers[source, destination] = channel
+            with channel.locked():
+                channel.reclaim()
+                head, tail = channel.bounds()
+                position = head % channel.capacity
+                padding = channel.capacity - position if position + total > channel.capacity else 0
+                if padding + total > channel.capacity - (head - tail):
+                    return None
+                if padding:
+                    _FRAME.pack_into(channel.mapping, _HEADER_SIZE + position, padding, 0, _PADDING, _DONE, 0, 0)
+                    head += padding
+                    position = 0
+                start = _HEADER_SIZE + position + _FRAME.size
+                channel.mapping[start : start + len(encoded_key)] = encoded_key
+                start += len(encoded_key)
+                if isinstance(payload, TensorFrame):
+                    view = memoryview(channel.mapping)
+                    try:
+                        payload.write(view, start)
+                    finally:
+                        view.release()
+                else:
+                    channel.mapping[start : start + length] = payload
+                _FRAME.pack_into(
+                    channel.mapping,
+                    _HEADER_SIZE + position,
+                    total,
+                    length,
+                    _TENSORS if isinstance(payload, TensorFrame) else _BYTES,
+                    _READY,
+                    len(encoded_key),
+                    0,
+                )
+                # Retire an earlier version while holding the same lock as all readers.
+                channel.cancel(key=key)
+                struct.pack_into("<Q", channel.mapping, 16, head + total)
+                channel.index[key] = head
+                channel.positions.append((head, key))
+                channel.cursor = head + total
+            return {"name": channel.name, "position": head, "size": length}
+
+    def _reader(self, name: str) -> _Channel:
+        if not name.startswith(f"omni_shm_ring_v1_{os.getuid()}_") or not name.replace("_", "").isalnum():
+            raise ValueError("invalid local SHM ring name")
+        channel = self.readers.get(name)
+        if channel is None:
+            channel = _Channel.open(name)
+            self.readers[name] = channel
+        return channel
+
+    def _receive(self, channel: _Channel, key: str, descriptor: dict | None) -> tuple[object, int] | None:
+        if not channel.producer_alive():
+            self.readers.pop(channel.name, None)
+            channel.close()
+            self._discovery_versions.clear()
+            return None
+        # This optimistic check only rejects a miss; successful reads still
+        # take the shared lock. An unchanged publication cursor cannot have
+        # introduced a new key since this receiver's last locked scan.
+        if key not in channel.index and struct.unpack_from("<Q", channel.mapping, 16)[0] == channel.cursor:
+            return None
+        try:
+            fcntl.flock(channel.fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            return None
+        try:
+            channel.scan()
+            position = channel.index.get(key)
+            if position is None or descriptor is not None and position != descriptor["position"]:
+                return None
+            _total, length, kind, state, key_length, _reserved = channel.frame(position)
+            del channel.index[key]
+            if state != _READY:
+                return None
+            start = _HEADER_SIZE + position % channel.capacity + _FRAME.size + key_length
+            view = memoryview(channel.mapping)[start : start + length]
+            try:
+                # No reference to ring storage escapes this locked boundary.
+                result = read_tensor_frame(view) if kind == _TENSORS else bytes(view)
+            finally:
+                view.release()
+                channel.consume(position)
+            return result, length
+        finally:
+            fcntl.flock(channel.fd, fcntl.LOCK_UN)
+
+    def get(self, source: str, destination: str, key: str, descriptor: dict | None) -> tuple[object, int] | None:
+        with self._lock:
+            if self._closed:
+                return None
+            if descriptor is not None:
+                try:
+                    result = self._receive(self._reader(descriptor["name"]), key, descriptor)
+                    if result is None:
+                        self._reap_readers()
+                    return result
+                except (FileNotFoundError, BlockingIOError):
+                    return None
+            if not self.capacity or not source.isdigit() or not destination.isdigit():
+                return None
+            edge_prefix = f"{self.prefix}{source}_{destination}_"
+            for name, channel in tuple(self.readers.items()):
+                if name.startswith(edge_prefix):
+                    result = self._receive(channel, key, None)
+                    if result is not None:
+                        return result
+            directory = f"/dev/shm/{self.directory}"
+            try:
+                status = os.stat(directory)
+            except FileNotFoundError:
+                self._discovery_versions.clear()
+                self._reap_readers()
+                return None
+            version = (status.st_dev, status.st_ino, status.st_mtime_ns)
+            edge = (source, destination)
+            if version == self._discovery_versions.get(edge):
+                return None
+            self._reap_readers()
+            # A new producer advances the stamp before publishing any frames,
+            # retaining first-message discovery without scanning every miss.
+            discovered_result = None
+            discovery_complete = True
+            for path in glob.iglob(f"/dev/shm/{self.directory}/{edge_prefix}*"):
+                name = os.path.basename(path)
+                if name in self.readers:
+                    continue
+                try:
+                    channel = self._reader(name)
+                except FileNotFoundError:
+                    # A published marker whose allocation or producer is gone
+                    # cannot become valid again: allocation names are unique.
+                    try:
+                        with _registry_update(directory) as directory_fd:
+                            os.unlink(name, dir_fd=directory_fd)
+                    except FileNotFoundError:
+                        pass
+                    except OSError as error:
+                        logger.debug("Failed to prune host ring marker %s: %s", path, error)
+                    continue
+                except (BlockingIOError, ValueError):
+                    # Attachment may race a writer holding the frame lock.
+                    # Retry discovery after it releases, even at this stamp.
+                    discovery_complete = False
+                    continue
+                if discovered_result is None:
+                    discovered_result = self._receive(channel, key, None)
+            # Attach every producer before caching this edge's stamp. Returning
+            # as soon as one frame is found would hide its unvisited peers.
+            if discovery_complete:
+                self._discovery_versions[edge] = version
+            else:
+                self._discovery_versions.pop(edge, None)
+            return discovered_result
+
+    def cancel(self, key: str | None = None, prefix: str | None = None) -> int:
+        with self._lock:
+            count = 0
+            for channel in self.producers.values():
+                with channel.locked():
+                    count += channel.cancel(key, prefix)
+            return count
+
+    def reap(self) -> None:
+        with self._lock:
+            for channel in self.producers.values():
+                with channel.locked():
+                    channel.reclaim()
+            self._reap_readers()
+
+    def _reap_readers(self) -> None:
+        for name, channel in tuple(self.readers.items()):
+            if not channel.producer_alive() or not os.path.exists(f"/dev/shm/{name}"):
+                channel.close()
+                del self.readers[name]
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            for channel in self.producers.values():
+                with channel.locked():
+                    channel.scan()
+                    for position in channel.index.values():
+                        channel.consume(position)
+                channel.close()
+            for channel in self.readers.values():
+                channel.close()
+            self.producers.clear()
+            self.readers.clear()
+            self._discovery_versions.clear()
+            try:
+                os.rmdir(f"/dev/shm/{self.directory}")
+            except (FileNotFoundError, OSError):
+                # Other producers in this deployment may still own rings.
+                pass

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from collections.abc import Mapping
 from dataclasses import asdict, fields, is_dataclass
@@ -205,7 +205,7 @@ class OmniMsgpackDecoder:
                 return self._decode_pil_image(obj)
 
             # Process values recursively first
-            processed = {k: self._post_process(v) for k, v in obj.items()}
+            processed = {k: self._post_process(v) if isinstance(v, (dict, list, tuple)) else v for k, v in obj.items()}
 
             # Check if this looks like an OmniRequestOutput (check before RequestOutput
             # since OmniRequestOutput may also have some RequestOutput-like fields)
@@ -223,10 +223,13 @@ class OmniMsgpackDecoder:
             return processed
 
         if isinstance(obj, list):
-            return [self._post_process(item) for item in obj]
+            # Native msgpack already owns scalar values. Only containers can
+            # contain wire markers; avoid a recursive call for every token or
+            # float in long model payloads, while preserving new-list semantics.
+            return [self._post_process(item) if isinstance(item, (dict, list, tuple)) else item for item in obj]
 
         if isinstance(obj, tuple):
-            return tuple(self._post_process(item) for item in obj)
+            return tuple(self._post_process(item) if isinstance(item, (dict, list, tuple)) else item for item in obj)
 
         return obj
 
@@ -350,6 +353,10 @@ class OmniSerde:
         """Deserialize bytes to an object."""
         return self.decoder.decode(data)
 
+    def restore(self, obj: object) -> object:
+        """Restore wire objects from an already decoded transport tree."""
+        return self.decoder._post_process(obj)
+
 
 # Global instance for simple interface
 OmniSerializer = OmniSerde()

vllm_omni/distributed/omni_connectors/utils/tensor_frame.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Owned CPU tensor trees with a msgpack header and raw tensor regions.
+
+Tensor references use message-scoped msgpack extensions, preserving literal
+extensions and native scalar containers without reserved dictionary keys. It never
+exposes a view of transport storage to a caller. Unsupported objects retain
+the ordinary OmniSerializer path.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import math
+import struct
+import uuid
+from dataclasses import dataclass
+
+import msgspec
+import numpy as np
+import torch
+from PIL import Image
+
+from .serialization import OmniSerializer
+
+_LENGTH = struct.Struct("<I")
+_ENCODER = msgspec.msgpack.Encoder()
+_REFERENCE = struct.Struct("<16sI")
+_REFERENCE_CODE = 42
+_DECODER = msgspec.msgpack.Decoder()
+
+
+def _align(size: int) -> int:
+    return (size + 15) & ~15
+
+
+class _UnsupportedTreeError(Exception):
+    pass
+
+
+@dataclass
+class TensorFrame:
+    header: bytes
+    tensors: list[torch.Tensor]
+    offsets: list[int]
+    size: int
+
+    def write(self, buffer: memoryview, start: int) -> None:
+        """Copy directly into caller-owned writable storage, while it is locked."""
+        _LENGTH.pack_into(buffer, start, len(self.header))
+        buffer[start + 4 : start + 4 + len(self.header)] = self.header
+        data_start = start + _align(4 + len(self.header))
+        # c_char is an anchor only; all bounds are checked by the ring writer.
+        anchor = ctypes.c_char.from_buffer(buffer)
+        base = ctypes.addressof(anchor) + data_start
+        for tensor, offset in zip(self.tensors, self.offsets, strict=True):
+            if tensor.nbytes:
+                ctypes.memmove(base + offset, tensor.data_ptr(), tensor.nbytes)
+
+
+def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> TensorFrame | bytes | None:
+    """Encode native containers once, or prepare their dense CPU tensor regions.
+
+    Scalar/token lists stay in msgpack's C encoder. Tensor-free messages reuse
+    those encoded bytes, rather than constructing and discarding a tagged tree.
+    Sources remain caller-owned until the synchronous put() returns.
+    """
+    tensors: list[torch.Tensor] = []
+    offsets: list[int] = []
+    descriptors: list[tuple[str, tuple[int, ...], int]] = []
+    data_size = 0
+    nonce: bytes | None = None
+
+    def encode_tensor(value: object) -> object:
+        nonlocal data_size, nonce
+        if type(value) is not torch.Tensor:
+            # Reuse the common wire representations for mixed model payloads.
+            # Large arrays/images keep the ordinary path without first making
+            # an extra byte snapshot just to discover that the frame cannot fit.
+            if max_bytes is not None:
+                if isinstance(value, np.ndarray) and value.nbytes > max_bytes:
+                    raise _UnsupportedTreeError
+                if isinstance(value, Image.Image) and value.width * value.height * len(value.getbands()) > max_bytes:
+                    raise _UnsupportedTreeError
+            return OmniSerializer.encoder._enc_hook(value)
+        if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized:
+            raise _UnsupportedTreeError
+        offset = _align(data_size)
+        index = len(tensors)
+        tensors.append(value)
+        offsets.append(offset)
+        descriptors.append((str(value.dtype).removeprefix("torch."), tuple(value.shape), offset))
+        data_size = offset + value.nbytes
+        if max_bytes is not None and data_size > max_bytes:
+            raise _UnsupportedTreeError
+        if nonce is None:
+            nonce = uuid.uuid4().bytes
+        return msgspec.msgpack.Ext(_REFERENCE_CODE, _REFERENCE.pack(nonce, index))
+
+    try:
+        tree = msgspec.msgpack.Encoder(enc_hook=encode_tensor).encode(payload)
+        if not tensors:
+            return tree
+        header = _ENCODER.encode([descriptors, tree, nonce])
+    except (_UnsupportedTreeError, TypeError, OverflowError, RecursionError):
+        return None
+    size = _align(4 + len(header)) + data_size
+    if max_bytes is not None and size > max_bytes:
+        return None
+    # Only materialize non-contiguous/flagged storage after the entire message
+    # is known to fit. Unsupported and large frames avoid an unused copy.
+    for index, tensor in enumerate(tensors):
+        if tensor.is_conj():
+            tensor = tensor.resolve_conj()
+        if tensor.is_neg():
+            tensor = tensor.resolve_neg()
+        if not tensor.is_contiguous():
+            tensor = tensor.contiguous()
+        tensors[index] = tensor
+    return TensorFrame(header, tensors, offsets, size)
+
+
+def read_tensor_frame(buffer: memoryview) -> object:
+    """Decode into independent CPU allocations before returning ring credit."""
+    if len(buffer) < 4:
+        raise ValueError("truncated tensor frame")
+    header_size = _LENGTH.unpack_from(buffer)[0]
+    data_start = _align(4 + header_size)
+    if data_start > len(buffer):
+        raise ValueError("tensor frame header exceeds payload")
+    header_view = buffer[4 : 4 + header_size]
+    try:
+        descriptors, tree, nonce = _DECODER.decode(header_view)
+    finally:
+        header_view.release()
+    # All leaves share one independently owned allocation. This amortizes CPU
+    # allocation/copy dispatch across small codec groups and hidden-state leaves.
+    data_size = len(buffer) - data_start
+    # A small metadata tensor must not keep a large, already consumed hidden
+    # state alive. Pack small trees; give large multi-leaf trees separate owners.
+    storage = torch.empty(data_size, dtype=torch.uint8) if data_size <= 65536 or len(descriptors) == 1 else None
+    anchor = ctypes.c_char.from_buffer(buffer)
+    if data_size and storage is not None:
+        ctypes.memmove(storage.data_ptr(), ctypes.addressof(anchor) + data_start, data_size)
+    tensors = []
+    for name, shape, offset in descriptors:
+        dtype = getattr(torch, name, None)
+        if not isinstance(dtype, torch.dtype):
+            raise ValueError("unknown tensor frame dtype")
+        if any(type(dim) is not int or dim < 0 for dim in shape) or type(offset) is not int or offset < 0:
+            raise ValueError("invalid tensor frame region")
+        count = math.prod(shape)
+        size = count * dtype.itemsize
+        start = data_start + offset
+        if start + size > len(buffer):
+            raise ValueError("tensor frame region exceeds payload")
+        if size and storage is not None:
+            tensor = storage[offset : offset + size].view(dtype).reshape(shape)
+        else:
+            tensor = torch.empty(shape, dtype=dtype)
+            if size:
+                ctypes.memmove(tensor.data_ptr(), ctypes.addressof(anchor) + start, size)
+        tensors.append(tensor)
+
+    def restore_reference(code: int, data: memoryview) -> object:
+        if code == _REFERENCE_CODE and len(data) == _REFERENCE.size:
+            reference_nonce, index = _REFERENCE.unpack(data)
+            if reference_nonce == nonce:
+                if index >= len(tensors):
+                    raise ValueError("invalid tensor frame reference")
+                return tensors[index]
+        # Literal user extensions keep their existing msgpack wire semantics.
+        return msgspec.msgpack.Ext(code, bytes(data))
+
+    return msgspec.msgpack.decode(tree, ext_hook=restore_reference)

vllm_omni/entrypoints/async_omni.py

@@ -959,12 +959,6 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         engines have no EngineCore pause to hold, so ``wake_up`` restores
         admission and ``sleep → wake → generate`` keeps working.
         """
-        # Validate before touching admission state. Raising after
-        # ``_paused = True`` would wedge generate() for good: nothing was
-        # slept, so wake_up() has no tags to restore and never reaches the
-        # block that clears ``_paused``.
-        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
-
         # Block admission before any sleep RPC so generate() waits on
         # _pause_cond during the drain/offload window. Wait until generate()
         # coroutines that already passed the pause check have submitted (or
@@ -972,6 +966,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         async with self._pause_cond:
             self._paused = True
             await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0)
+
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
         # EngineCore.sleep resets receiver caches itself; only clear P0 here.
         if 0 in ar_stage_ids:
             await self._clear_frontend_mm_cache()

vllm_omni/entrypoints/openai/api_server.py

@@ -3062,8 +3062,6 @@ async def omni_sleep(request: OmniSleepRequest, raw_request: Request):
         raise HTTPException(status_code=501, detail="Engine does not support sleep")
     try:
         acks = await engine_client.sleep(stage_ids=request.stage_ids, level=request.level)
-    except ValueError as e:
-        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) from e
     except RuntimeError as e:
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to sleep: {e}") from e
     finally:
@@ -3088,8 +3086,6 @@ async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request):
         acks = await engine_client.wake_up(stage_ids=request.stage_ids)
     except NotImplementedError:
         raise
-    except ValueError as e:
-        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) from e
     except RuntimeError as e:
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to wake up: {e}") from e
     for sid in request.stage_ids:

vllm_omni/model_executor/models/minicpmo_4_5/encoder_cuda_graph.py

@@ -1,262 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Exact-shape MiniCPM adapters for vLLM encoder CUDA graph management."""
-
-from collections import Counter
-from collections.abc import Callable, Hashable
-from copy import copy
-from math import prod
-
-import torch
-from vllm.config import VllmConfig
-from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
-from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
-from vllm.v1.worker.encoder_cudagraph_defs import (
-    EncoderCudaGraphCaptureInputs,
-    EncoderCudaGraphConfig,
-    EncoderCudaGraphReplayBuffers,
-    EncoderItemSpec,
-)
-
-from vllm_omni.platforms import current_omni_platform
-
-
-def _copy_exact(target: torch.Tensor, source: torch.Tensor) -> None:
-    target.copy_(source)
-
-
-class _ExactShapeEncoder(SupportsEncoderCudaGraph):
-    """One already-packed local encoder batch is one indivisible manager item.
-
-    Packing and distributed placement belong to the caller. Keeping the batch
-    intact preserves attention masks and avoids introducing padding semantics.
-    """
-
-    def __init__(
-        self, forward: Callable[..., torch.Tensor], inputs: tuple[torch.Tensor | None, ...], output_shape: torch.Size
-    ):
-        self.forward = forward
-        self.templates = tuple(None if x is None else (x.shape, x.dtype, x.device) for x in inputs)
-        self.output_shape = output_shape
-        self.tokens = prod(output_shape[:-1]) or 1
-        self.keys = [str(i) for i, x in enumerate(inputs) if x is not None]
-
-    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
-        return EncoderCudaGraphConfig(
-            modalities=["image", "audio"],
-            buffer_keys=self.keys,
-            out_hidden_size=self.output_shape[-1],
-            padding_logics=dict.fromkeys(self.keys, _copy_exact),
-        )
-
-    def get_encoder_cudagraph_budget_range(self, vllm_config: VllmConfig) -> tuple[int, int]:
-        return self.tokens, self.tokens
-
-    def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, torch.Tensor]) -> list[EncoderItemSpec]:
-        return [EncoderItemSpec(input_size=self.tokens, output_tokens=self.tokens)]
-
-    def select_encoder_cudagraph_items(
-        self, mm_kwargs: dict[str, torch.Tensor], indices: list[int]
-    ) -> dict[str, torch.Tensor]:
-        assert indices == [0]
-        return dict(mm_kwargs)
-
-    def prepare_encoder_cudagraph_capture_inputs(
-        self,
-        token_budget: int,
-        max_batch_size: int,
-        max_frames_per_batch: int,
-        device: torch.device,
-        dtype: torch.dtype,
-        path: str = "default",
-        axis_keys: tuple[Hashable, ...] | None = (),
-    ) -> EncoderCudaGraphCaptureInputs:
-        return EncoderCudaGraphCaptureInputs(
-            {
-                str(i): torch.zeros(shape, dtype=input_dtype, device=input_device)
-                for i, template in enumerate(self.templates)
-                if template is not None
-                for shape, input_dtype, input_device in [template]
-            }
-        )
-
-    def prepare_encoder_cudagraph_replay_buffers(
-        self,
-        mm_kwargs: dict[str, torch.Tensor],
-        max_batch_size: int,
-        max_frames_per_batch: int,
-        path: str = "default",
-    ) -> EncoderCudaGraphReplayBuffers:
-        return EncoderCudaGraphReplayBuffers(mm_kwargs)
-
-    def encoder_cudagraph_forward(self, inputs: dict[str, torch.Tensor], path: str = "default") -> torch.Tensor:
-        return self.forward(*(inputs.get(str(i)) for i in range(len(self.templates))))
-
-    def encoder_eager_forward(self, mm_kwargs: dict[str, torch.Tensor], path: str = "default") -> torch.Tensor:
-        return self.encoder_cudagraph_forward(mm_kwargs, path)
-
-    def postprocess_encoder_output(
-        self,
-        outputs: dict[str, torch.Tensor],
-        indices: list[int],
-        per_item_out_tokens: list[int],
-        dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
-        clone: bool = False,
-        batch_mm_kwargs: dict[str, torch.Tensor] | None = None,
-    ) -> None:
-        output = outputs["default"]
-        # Callers retain embeddings across subsequent replays.
-        dest[0] = output.clone() if clone else output
-
-
-class EncoderCudaGraph:
-    """Capture repeated shapes without padding or changing attention semantics.
-
-    The owner must bypass this object for training, stateful KV caches and
-    data-dependent Python branches. Entries on the same device and replay
-    stream share a graph pool and capture stream. Different replay streams
-    and encoder instances remain isolated. Graphs must not run concurrently
-    on a shared pool, and their intermediate tensors must not escape capture;
-    callers get a clone so a later replay cannot overwrite retained embeddings.
-    Unseen shapes after the cap run eagerly instead of growing GPU memory.
-    Admission is first-come, with no eviction or automatic recapture. Startup
-    profiling can consume slots. The free-memory floor is a pre-capture guard,
-    not a bound on pool bytes or a guarantee that capture will fit.
-
-    Capture errors are fatal for this instance: propagate the original error
-    and reject subsequent calls without retrying or falling back to CUDA eager
-    execution, since a failed capture may leave the CUDA context unusable.
-    """
-
-    def __init__(
-        self,
-        forward: Callable[..., torch.Tensor],
-        vllm_config: VllmConfig,
-        *,
-        max_graphs: int = 4,
-        min_capture_calls: int = 2,
-        min_free_bytes: int = 1 << 30,
-        share_pools: bool = True,
-    ):
-        for name, value, minimum in (
-            ("max_graphs", max_graphs, 0),
-            ("min_capture_calls", min_capture_calls, 2),
-            ("min_free_bytes", min_free_bytes, 0),
-        ):
-            if type(value) is not int or value < minimum:
-                raise ValueError(f"{name} must be an integer >= {minimum}")
-        self.forward = forward
-        self.vllm_config = vllm_config
-        self.max_graphs = max_graphs
-        self.min_capture_calls = min_capture_calls
-        self.min_free_bytes = min_free_bytes
-        self.share_pools = share_pools
-        self._capture_pools: dict[tuple, tuple[tuple[int, int], torch.cuda.Stream]] = {}
-        self.graphs: dict[tuple, EncoderCudaGraphManager] = {}
-        self._seen: dict[tuple, torch.Size] = {}
-        self._seen_calls: Counter[tuple] = Counter()
-        self._misses: Counter[str] = Counter()
-        self._capture_failed = False
-
-    def get_cumulative_stats(self) -> dict[str, int | float]:
-        """Include eager admission misses that never reach an upstream manager."""
-        hits = sum(entry.graph_hits for entry in self.graphs.values())
-        misses = sum(self._misses.values()) + sum(entry.graph_misses for entry in self.graphs.values())
-        return {
-            "graph_hits": hits,
-            "graph_misses": misses,
-            "hit_rate": hits / (hits + misses) if hits + misses else 0.0,
-            "num_graphs": len(self.graphs),
-            "ineligible_misses": self._misses["ineligible"],
-            "capacity_misses": self._misses["capacity"],
-            "warmup_misses": self._misses["warmup"],
-            "memory_misses": self._misses["memory"],
-            "capture_failures": int(self._capture_failed),
-        }
-
-    def __call__(self, *inputs: torch.Tensor | None) -> torch.Tensor:
-        if self._capture_failed:
-            raise RuntimeError("Encoder CUDA graph capture previously failed; restart the worker before reuse")
-        tensors = [value for value in inputs if value is not None]
-        if (
-            not tensors
-            or any(value.device.type != "cuda" for value in tensors)
-            or len({value.device for value in tensors}) != 1
-            or torch.is_grad_enabled()
-            or torch.is_autocast_enabled("cuda")
-            or torch.cuda.is_current_stream_capturing()
-        ):
-            self._misses["ineligible"] += 1
-            return self.forward(*inputs)
-        # Do not share mutable replay buffers across independent CUDA streams.
-        key = (
-            torch.cuda.current_stream(tensors[0].device).cuda_stream,
-            tuple(None if x is None else (x.shape, x.dtype, x.device) for x in inputs),
-        )
-        entry = self.graphs.get(key)
-        if entry is None:
-            if len(self.graphs) >= self.max_graphs:
-                self._misses["capacity"] += 1
-                return self.forward(*inputs)
-            if key not in self._seen:
-                # Avoid paying capture latency for one-off shapes. Bound even
-                # the admission metadata for streams with arbitrary resolutions.
-                if len(self._seen) >= self.max_graphs * 4:
-                    oldest = next(iter(self._seen))
-                    self._seen.pop(oldest)
-                    self._seen_calls.pop(oldest)
-                output = self.forward(*inputs)
-                self._seen[key] = output.shape
-                self._seen_calls[key] = 1
-                self._misses["warmup"] += 1
-                return output
-            self._seen_calls[key] = min(self._seen_calls[key] + 1, self.min_capture_calls)
-            if self._seen_calls[key] < self.min_capture_calls:
-                self._misses["warmup"] += 1
-                return self.forward(*inputs)
-            if self.min_free_bytes and current_omni_platform.get_free_memory(tensors[0].device) < self.min_free_bytes:
-                self._misses["memory"] += 1
-                return self.forward(*inputs)
-            try:
-                entry = self._capture(inputs, self._seen[key])
-            except Exception:
-                self._capture_failed = True
-                self._seen.clear()
-                self._seen_calls.clear()
-                raise
-            self.graphs[key] = entry
-            self._seen.pop(key)
-            self._seen_calls.pop(key)
-        return entry.execute({str(i): x for i, x in enumerate(inputs) if x is not None})[0]
-
-    def _capture(self, inputs: tuple[torch.Tensor | None, ...], output_shape: torch.Size) -> EncoderCudaGraphManager:
-        tensor = next(x for x in inputs if x is not None)
-        adapter = _ExactShapeEncoder(self.forward, inputs, output_shape)
-        # Do not mutate the runner configuration. This adapter has already
-        # packed and placed a local batch, so manager-level DP must be disabled.
-        config = copy(self.vllm_config)
-        config.compilation_config = copy(config.compilation_config)
-        config.compilation_config.encoder_cudagraph_token_budgets = [adapter.tokens]
-        config.compilation_config.encoder_cudagraph_max_vision_items_per_batch = 1
-        config.compilation_config.encoder_cudagraph_max_frames_per_batch = 0
-        config.parallel_config = copy(config.parallel_config)
-        config.parallel_config.tensor_parallel_size = 1
-        manager = EncoderCudaGraphManager(config, tensor.device, tensor.dtype, adapter)
-        replay_stream = torch.cuda.current_stream(tensor.device)
-        pool_key = (tensor.device, replay_stream.cuda_stream)
-        if self.share_pools and pool_key in self._capture_pools:
-            pool, stream = self._capture_pools[pool_key]
-        else:
-            pool = torch.cuda.graph_pool_handle()
-            stream = torch.cuda.Stream(device=tensor.device)
-        # A late capture must wait for prior replays and output clones before
-        # reusing their pool. Reusing the capture stream also lets the allocator
-        # reuse freed blocks from earlier captures. Manager input/output buffers
-        # are allocated outside capture; only transient intermediates share it.
-        stream.wait_stream(replay_stream)
-        with torch.cuda.stream(stream):
-            manager.capture(graph_pool=pool)
-        replay_stream.wait_stream(stream)
-        if self.share_pools:
-            self._capture_pools[pool_key] = (pool, stream)
-        return manager

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -55,8 +55,6 @@ from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPo
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.whisper.modeling_whisper import ACT2FN
 
-from vllm_omni.model_executor.models.minicpmo_4_5.encoder_cuda_graph import EncoderCudaGraph
-
 try:
     from transformers.models.whisper.modeling_whisper import WHISPER_ATTENTION_CLASSES
 except ImportError:
@@ -1582,7 +1580,6 @@ class SiglipVisionTransformer(SiglipPreTrainedModel):
         self.encoder = SiglipEncoder(config)
         self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
         self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
-        self._encoder_graph: EncoderCudaGraph | None = None
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -1590,18 +1587,6 @@ class SiglipVisionTransformer(SiglipPreTrainedModel):
     def get_input_embeddings(self) -> nn.Module:
         return self.embeddings.patch_embedding
 
-    def _encode_last_hidden_state(
-        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None
-    ) -> torch.Tensor:
-        output = self.encoder(
-            inputs_embeds=hidden_states,
-            attention_mask=attention_mask,
-            output_attentions=False,
-            output_hidden_states=False,
-            return_dict=True,
-        )
-        return self.post_layernorm(output.last_hidden_state)
-
     @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipVisionConfig)
     def forward(
@@ -1651,21 +1636,6 @@ class SiglipVisionTransformer(SiglipPreTrainedModel):
                 else patch_attention_mask
             )
 
-        # Position/mask construction above performs host decisions. Capture
-        # only the transformer stack; padded FlashAttention unpadding also
-        # reads dynamic lengths on the host and remains eager.
-        if (
-            self._encoder_graph is not None
-            and not self.training
-            and not output_attentions
-            and not output_hidden_states
-            and not (self._use_flash_attention_2 and attention_mask is not None)
-        ):
-            last_hidden_state = self._encoder_graph(hidden_states, attention_mask)
-            if not return_dict:
-                return (last_hidden_state, None)
-            return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state)
-
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,
             attention_mask=attention_mask,
@@ -3945,18 +3915,6 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
 
         self.config = config
         self.multimodal_config = multimodal_config
-        # Model-local opt-out for A/B measurements; --enforce-eager always wins.
-        encoder_graphs = (
-            bool(getattr(config, "encoder_cuda_graph", True)) and not vllm_config.model_config.enforce_eager
-        )
-        # Per-encoder limits, configurable via --hf-overrides. No eviction:
-        # increasing the cap trades retained GPU memory for shape coverage.
-        encoder_graph_options = {
-            "max_graphs": getattr(config, "encoder_cuda_graph_max_graphs", 4),
-            "min_capture_calls": getattr(config, "encoder_cuda_graph_min_capture_calls", 2),
-            "min_free_bytes": getattr(config, "encoder_cuda_graph_min_free_bytes", 1 << 30),
-            "share_pools": getattr(config, "encoder_cuda_graph_share_pools", True),
-        }
 
         # Initialize image processor
         self.image_processor = MiniCPMVImageProcessor(
@@ -3976,10 +3934,6 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
                 config.vision_config._attn_implementation = "eager"
 
             self.vpm = SiglipVisionTransformer(config.vision_config)
-            if encoder_graphs:
-                self.vpm._encoder_graph = EncoderCudaGraph(
-                    self.vpm._encode_last_hidden_state, vllm_config, **encoder_graph_options
-                )
             # Drop last layer if configured
             if config.drop_vision_last_layer:
                 self.vpm.encoder.layers = self.vpm.encoder.layers[:-1]
@@ -4047,11 +4001,6 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             self.audio_encoder_layer = None
             self.audio_past_key_values = None
 
-        self._audio_encoder_graph = (
-            EncoderCudaGraph(self._encode_audio_features, vllm_config, **encoder_graph_options)
-            if encoder_graphs
-            else None
-        )
         self.mm_token_ids = set[int]()
         self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors
 
@@ -4475,36 +4424,6 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             audio_attention_mask_ = torch.logical_or(audio_attention_mask_, torch.logical_not(chunk_mask))
 
         audio_attention_mask[audio_attention_mask_] = float("-inf")
-        graph = getattr(self, "_audio_encoder_graph", None)
-        if (
-            graph is not None
-            and not self.training
-            and self.audio_encoder_layer == -1
-            # Whisper's FP16 overflow guard reads tensor booleans on the host.
-            and self.apm.conv1.weight.dtype != torch.float16
-        ):
-            audio_embeds = graph(wavforms, audio_attention_mask)
-        else:
-            audio_embeds = self._encode_audio_features(wavforms, audio_attention_mask)
-
-        _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens)
-
-        num_audio_tokens = feature_lens_after_pooling
-
-        final_audio_embeds = list[torch.Tensor]()
-        idx = 0
-        for i in range(len(audio_feature_lens_raw)):
-            target_audio_embeds_lst = list[torch.Tensor]()
-            for _ in range(len(audio_feature_lens_raw[i])):
-                target_audio_embeds_lst.append(audio_embeds[idx, : num_audio_tokens[idx], :])
-                idx += 1
-
-            final_audio_embeds.append(torch.cat(target_audio_embeds_lst))
-
-        return final_audio_embeds
-
-    def _encode_audio_features(self, wavforms: torch.Tensor, audio_attention_mask: torch.Tensor) -> torch.Tensor:
-        """Stateless audio encoder, projection and pooling (no streaming KV)."""
         selects_final_layer = self.audio_encoder_layer == -1
         audio_outputs = self.apm(
             wavforms,
@@ -4524,7 +4443,21 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         audio_embeds = self.audio_avg_pooler(audio_embeds)
         audio_embeds = audio_embeds.transpose(1, 2)
 
-        return audio_embeds
+        _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens)
+
+        num_audio_tokens = feature_lens_after_pooling
+
+        final_audio_embeds = list[torch.Tensor]()
+        idx = 0
+        for i in range(len(audio_feature_lens_raw)):
+            target_audio_embeds_lst = list[torch.Tensor]()
+            for _ in range(len(audio_feature_lens_raw[i])):
+                target_audio_embeds_lst.append(audio_embeds[idx, : num_audio_tokens[idx], :])
+                idx += 1
+
+            final_audio_embeds.append(torch.cat(target_audio_embeds_lst))
+
+        return final_audio_embeds
 
     def get_audio_embedding_streaming(
         self,