diff --git a/tests/distributed/omni_connectors/test_cuda_ipc_connector.py b/tests/distributed/omni_connectors/test_cuda_ipc_connector.py new file mode 100644 index 000000000..1ce9b3b66 --- /dev/null +++ b/tests/distributed/omni_connectors/test_cuda_ipc_connector.py @@ -0,0 +1,101 @@ +"""CudaIpcConnector: leaf walking (CPU) and a spawned producer/consumer round trip (CUDA).""" + +from __future__ import annotations + +import multiprocessing as mp +import os +import uuid + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.distributed.omni_connectors.connectors import cuda_ipc_connector as cic +from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_walk_rewrites_tensor_leaves_only(): + seen = [] + + def leaf(t): + seen.append(t) + return "L" + + data = {"a": torch.zeros(2), "b": [1, torch.ones(1), (torch.ones(1), "s")], "c": {"d": None}} + out = cic._walk(data, leaf) + assert out == {"a": "L", "b": [1, "L", ("L", "s")], "c": {"d": None}} + assert len(seen) == 3 + # marker dicts are treated as leaves too + marker = {cic.MARKER: 1, "args": b"", "event": b"", "dtype": "torch.float32", "shape": [1]} + assert cic._walk({"m": marker}, lambda m: "M") == {"m": "M"} + assert cic.is_marker(marker) and not cic.is_marker({"x": 1}) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_cpu_leaves_pass_through_shared_memory(tmp_path): + assert "CudaIpcConnector" in OmniConnectorFactory._registry + conn = cic.CudaIpcConnector({"stage_id": 0, "extra": {"cuda_ipc_enabled": False}}) + key = f"cic_cpu_{uuid.uuid4().hex[:8]}" + payload = {"codes": {"audio": torch.arange(6).reshape(3, 2)}, "meta": {"finished": True}} + ok, size, meta = conn.put("stage0", "stage1", key, payload) + assert ok and size > 0 and "cuda_ipc" not in (meta or {}) + got, _ = conn.get("stage0", "stage1", key, meta) + assert torch.equal(got["codes"]["audio"], payload["codes"]["audio"]) and got["meta"]["finished"] is True + conn.cleanup(key) + assert "cuda_ipc" in conn.health() + conn.close() + + +def _producer(key: str, q: mp.Queue): + conn = cic.CudaIpcConnector({"stage_id": 0}) + t = torch.randint(0, 2048, (25, 16), device="cuda", dtype=torch.int64) + hs = torch.randn(4, 8, device="cuda", dtype=torch.float16) + ok, size, meta = conn.put( + "stage0", "stage1", key, {"codes": {"audio": t}, "hidden": hs, "meta": {"n": 25}, "cpu": torch.ones(2)} + ) + q.put(("put", ok, size, meta, t.cpu(), hs.cpu())) + # wait for consumer ack, then sweep on cleanup + deadline = 60 + import time + + while deadline > 0 and not os.path.exists(cic._ack_path(key)): + time.sleep(0.1) + deadline -= 0.1 + conn.cleanup(key) + q.put(("done", len(conn._exported))) + + +def _consumer(key: str, meta, q: mp.Queue): + conn = cic.CudaIpcConnector({"stage_id": 1}) + got, size = conn.get("stage0", "stage1", key, meta) + q.put(("got", got["codes"]["audio"].cpu(), got["hidden"].cpu(), got["meta"], got["cpu"], conn.health()["rebuilt"])) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_cuda_ipc_round_trip_between_processes(): + ctx = mp.get_context("spawn") + key = f"cic_gpu_{uuid.uuid4().hex[:8]}" + q = ctx.Queue() + prod = ctx.Process(target=_producer, args=(key, q)) + prod.start() + tag, ok, size, meta, t_ref, hs_ref = q.get(timeout=120) + assert tag == "put" and ok and meta.get("cuda_ipc") == 2 + cons = ctx.Process(target=_consumer, args=(key, meta, q)) + cons.start() + tag, codes, hidden, m, cpu_t, rebuilt = q.get(timeout=120) + assert tag == "got" and rebuilt == 2 + assert ( + torch.equal(codes, t_ref) + and torch.equal(hidden, hs_ref) + and m == {"n": 25} + and torch.equal(cpu_t, torch.ones(2)) + ) + tag, live = q.get(timeout=120) + assert tag == "done" and live == 0 + prod.join(30) + cons.join(30) + assert not os.path.exists(cic._ack_path(key)) diff --git a/vllm_omni/deploy/qwen3_tts_cuda_ipc.yaml b/vllm_omni/deploy/qwen3_tts_cuda_ipc.yaml new file mode 100644 index 000000000..07b459522 --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_cuda_ipc.yaml @@ -0,0 +1,20 @@ +# Qwen3-TTS with the CUDA-IPC connector between talker and code2wav +# (co-located CUDA stages, same GPU). GPU tensor leaves in the stage payload +# are shared through CUDA IPC instead of D2H + msgpack; CPU leaves and the +# control-plane metadata still use shared memory. Measurement gate: adopt in +# the edge profile only if the per-chunk hop latency improves (see +# benchmarks/tts/edge_results/stream_latency and WP7 notes). +base_config: qwen3_tts.yaml +async_chunk: true +connectors: + connector_of_shared_memory: + name: CudaIpcConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 72 + initial_codec_chunk_frames: 1 + decode_batch_max_size: 4 diff --git a/vllm_omni/distributed/omni_connectors/connectors/cuda_ipc_connector.py b/vllm_omni/distributed/omni_connectors/connectors/cuda_ipc_connector.py new file mode 100644 index 000000000..73b11701e --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/cuda_ipc_connector.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Device-direct connector for co-located CUDA stages (single host, same GPU). + +Tensor leaves that live on a CUDA device are exported through CUDA IPC +(``torch.multiprocessing.reductions.reduce_tensor``) instead of being copied +to host memory and msgpack-serialized; everything else (CPU tensors, ids, +metadata, msgspec structs) still travels through the wrapped +:class:`SharedMemoryConnector`, so the control-plane contract (``put`` / +``get`` / ``cleanup``, key-addressed lookup, metadata pass-through) is +unchanged. + +Protocol per exported tensor: + +* producer records an inter-process CUDA event on the current stream, packs + ``(rebuild args, event handle, dtype, shape)`` into a marker dict and keeps + the source tensor alive in ``_exported[put_key]``; +* consumer rebuilds the tensor, waits on the event, **clones** it into its own + allocation and writes an ack file; the producer releases the source on its + next ``put``/``cleanup`` sweep (or when the registry exceeds ``max_live``). + +Requirements: producer and consumer are different processes on the same GPU +(``VLLM_WORKER_MULTIPROC_METHOD=spawn``). Tensors allocated inside CUDA-graph +private pools cannot be exported; they are cloned into a plain allocation first. +""" + +from __future__ import annotations + +import copy +import os +import pickle +import time +from collections import OrderedDict +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase +from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector + +logger = init_logger(__name__) + +MARKER = "__cuda_ipc__" +ACK_DIR = "/dev/shm" + + +def _ack_path(put_key: str) -> str: + return os.path.join(ACK_DIR, f"cuda_ipc_ack_{put_key}") + + +def is_marker(obj: Any) -> bool: + return isinstance(obj, dict) and obj.get(MARKER) == 1 + + +class _Exported: + __slots__ = ("tensor", "event", "t") + + def __init__(self, tensor: torch.Tensor, event: Any): + self.tensor, self.event, self.t = tensor, event, time.time() + + +def export_tensor(t: torch.Tensor) -> tuple[dict[str, Any], _Exported]: + """Export a CUDA tensor; returns (marker dict, keep-alive record).""" + from torch.multiprocessing.reductions import reduce_tensor + + src = t.detach() + if not src.is_contiguous(): + src = src.contiguous() + try: + _, args = reduce_tensor(src) + except RuntimeError: + # e.g. CUDA-graph private pool allocation: move to a plain allocation + src = src.clone() + _, args = reduce_tensor(src) + event = torch.cuda.Event(interprocess=True) + event.record(torch.cuda.current_stream(src.device)) + marker = { + MARKER: 1, + "args": pickle.dumps(args, protocol=pickle.HIGHEST_PROTOCOL), + "event": event.ipc_handle(), + "device": int(src.device.index or 0), + "dtype": str(src.dtype), + "shape": list(src.shape), + } + return marker, _Exported(src, event) + + +def rebuild_tensor(marker: dict[str, Any]) -> torch.Tensor: + """Rebuild an exported tensor in the consumer process and own a copy of it.""" + from torch.multiprocessing.reductions import rebuild_cuda_tensor + + args = pickle.loads(marker["args"]) + device = torch.device("cuda", int(marker.get("device", 0))) + with torch.cuda.device(device): + tensor = rebuild_cuda_tensor(*args) + event = torch.cuda.Event.from_ipc_handle(device, marker["event"]) + event.wait(torch.cuda.current_stream(device)) + owned = tensor.clone() + return owned + + +def _walk(obj: Any, fn) -> Any: + """Apply ``fn`` to tensor leaves / marker dicts inside dicts, lists, tuples and msgspec structs.""" + if isinstance(obj, torch.Tensor) or is_marker(obj): + return fn(obj) + if isinstance(obj, dict): + return {k: _walk(v, fn) for k, v in obj.items()} + if isinstance(obj, list): + return [_walk(v, fn) for v in obj] + if isinstance(obj, tuple): + return tuple(_walk(v, fn) for v in obj) + fields = getattr(obj, "__struct_fields__", None) + if fields: + out = copy.copy(obj) + for name in fields: + setattr(out, name, _walk(getattr(obj, name), fn)) + return out + return obj + + +class CudaIpcConnector(OmniConnectorBase): + """CUDA-IPC fast path for GPU tensor leaves, shared memory for the rest.""" + + supports_raw_data = False + + def __init__(self, config: dict[str, Any]): + self.config = config + self.stage_id = config.get("stage_id", -1) + extra = config.get("extra", {}) if isinstance(config.get("extra"), dict) else {} + self.max_live = int(extra.get("cuda_ipc_max_live", 256)) + self.enabled = bool(extra.get("cuda_ipc_enabled", True)) and torch.cuda.is_available() + self._inner = SharedMemoryConnector(config) + self._exported: OrderedDict[str, list[_Exported]] = OrderedDict() + self._metrics = {"ipc_tensors": 0, "ipc_bytes": 0, "fallback_leaves": 0, "rebuilt": 0} + + # ------------------------------------------------------------- producer + def _sweep(self) -> None: + for key in list(self._exported.keys()): + ack = _ack_path(key) + if os.path.exists(ack): + try: + os.remove(ack) + except OSError: + pass + self._exported.pop(key, None) + while len(self._exported) > self.max_live: + key, _ = self._exported.popitem(last=False) + logger.warning("CudaIpcConnector: dropping unacked export %s (registry > %d)", key, self.max_live) + + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + self._sweep() + keep: list[_Exported] = [] + + def leaf(t: Any) -> Any: + if isinstance(t, torch.Tensor) and t.is_cuda and self.enabled and t.numel() > 0: + try: + marker, rec = export_tensor(t) + except Exception as exc: # any IPC failure -> host path for this leaf + logger.debug("CudaIpcConnector: export failed (%s); falling back to shm", exc) + self._metrics["fallback_leaves"] += 1 + return t + keep.append(rec) + self._metrics["ipc_tensors"] += 1 + self._metrics["ipc_bytes"] += t.numel() * t.element_size() + return marker + return t + + rewritten = _walk(data, leaf) + ok, size, meta = self._inner.put(from_stage, to_stage, put_key, rewritten) + if ok and keep: + self._exported[put_key] = keep + meta = dict(meta or {}) + meta["cuda_ipc"] = len(keep) + return ok, size, meta + + # ------------------------------------------------------------- consumer + def get( + self, from_stage: str, to_stage: str, get_key: str, metadata: dict[str, Any] | None = None + ) -> tuple[Any, int] | None: + result = self._inner.get(from_stage, to_stage, get_key, metadata) + if result is None: + return None + obj, size = result + rebuilt = 0 + + def leaf(m: Any) -> Any: + nonlocal rebuilt + if is_marker(m): + rebuilt += 1 + return rebuild_tensor(m) + return m + + obj = _walk(obj, leaf) + if rebuilt: + self._metrics["rebuilt"] += rebuilt + try: + with open(_ack_path(get_key), "w") as f: + f.write(str(time.time())) + except OSError: + pass + return obj, size + + # ------------------------------------------------------------- lifecycle + def cleanup(self, request_id: str) -> None: + for key in list(self._exported.keys()): + if ( + key == request_id + or key.startswith(request_id + "_") + or key.endswith("_" + request_id) + or request_id in key + ): + self._exported.pop(key, None) + try: + os.remove(_ack_path(key)) + except OSError: + pass + self._inner.cleanup(request_id) + + def health(self) -> dict[str, Any]: + h = self._inner.health() + h.update({"cuda_ipc": self.enabled, "live_exports": len(self._exported), **self._metrics}) + return h + + def close(self) -> None: + self._exported.clear() + self._inner.close() diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index c49da1c6c..4ccbfc3a8 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -126,6 +126,13 @@ def _create_mori_transfer_engine_connector(config: dict[str, Any]) -> OmniConnec # Register connectors +def _create_cuda_ipc_connector(config: dict[str, Any]) -> OmniConnectorBase: + from .connectors.cuda_ipc_connector import CudaIpcConnector + + return CudaIpcConnector(config) + + +OmniConnectorFactory.register_connector("CudaIpcConnector", _create_cuda_ipc_connector) OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector) OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector) OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector)