update

ray + mooncake connector work on 2 nodes, but needes several updates\n1.start ray inside vllm-omni \n2.placement group optimize \n3.fix close

.gitignore

@@ -232,3 +232,4 @@ configs/development.yaml
 # Docker
 .dockerignore
 Dockerfile.dev
+discussion

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -283,6 +283,19 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="process",
+        choices=["process","ray"],
+        help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+
 
     return parser.parse_args()
 

examples/offline_inference/tmp_test/README.md

@@ -0,0 +1,34 @@
+# Offline Example of vLLM-omni for Qwen2.5-omni
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (Qwen2.5-omni)
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
+```bash
+tar -xf <Your Download Path>/seedtts_testset.tar
+cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
+python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
+  --input examples/offline_inference/qwen2_5_omni/meta.lst \
+  --output examples/offline_inference/qwen2_5_omni/top100.txt \
+  --topk 100
+```
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```

examples/offline_inference/tmp_test/end2end.py

@@ -0,0 +1,264 @@
+import argparse
+import os
+import os as _os_env_toggle
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
+
+SEED = 42
+# Set all random seeds
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Make PyTorch deterministic
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set environment variables for deterministic behavior
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to merged model directory (will be created if downloading).",
+    )
+    parser.add_argument("--thinker-model", type=str, default=None)
+    parser.add_argument("--talker-model", type=str, default=None)
+    parser.add_argument("--code2wav-model", type=str, default=None)
+    parser.add_argument(
+        "--hf-hub-id",
+        default="Qwen/Qwen2.5-Omni-7B",
+        help="Hugging Face repo id to download if needed.",
+    )
+    parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.")
+    parser.add_argument(
+        "--code2wav-dir",
+        default=None,
+        help="Path to code2wav folder (contains spk_dict.pt).",
+    )
+    parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).")
+    parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.")
+    parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
+    parser.add_argument("--max-model-len", type=int, default=32768)
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
+    )
+
+    parser.add_argument("--thinker-only", action="store_true")
+    parser.add_argument("--text-only", action="store_true")
+    parser.add_argument("--do-wave", action="store_true")
+    parser.add_argument(
+        "--prompt_type",
+        choices=[
+            "text",
+            "audio",
+            "audio-long",
+            "audio-long-chunks",
+            "audio-long-expand-chunks",
+            "image",
+            "video",
+            "video-frames",
+            "audio-in-video",
+            "audio-in-video-v2",
+            "audio-multi-round",
+            "badcase-vl",
+            "badcase-text",
+            "badcase-image-early-stop",
+            "badcase-two-audios",
+            "badcase-two-videos",
+            "badcase-multi-round",
+            "badcase-voice-type",
+            "badcase-voice-type-v2",
+            "badcase-audio-tower-1",
+            "badcase-audio-only",
+        ],
+        default="text",
+    )
+    parser.add_argument("--use-torchvision", action="store_true")
+    parser.add_argument("--tokenize", action="store_true")
+    parser.add_argument(
+        "--output-wav",
+        default="output.wav",
+        help="[Deprecated] Output wav directory (use --output-dir).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="outputs",
+        help="Output directory to save text and wav files together.",
+    )
+    parser.add_argument(
+        "--thinker-hidden-states-dir",
+        default="thinker_hidden_states",
+        help="Path to thinker hidden states directory.",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="process",
+        choices=["process","ray"],
+        help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred)."
+        )
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            args.prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
+    omni_llm = OmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+        worker_backend=args.worker_backend,
+        ray_address=args.ray_address
+    )
+    thinker_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+    talker_sampling_params = SamplingParams(
+        temperature=0.9,
+        top_p=0.8,
+        top_k=40,
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.05,
+        stop_token_ids=[8294],
+    )
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
+    import time
+    for i in range(1):
+        t1 = time.time()
+        prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+        print(f"prompt:{prompt}")
+        omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+        t2 = time.time()
+        print(f"==========> time:{t2-t1}")
+
+    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+    for stage_outputs in omni_outputs:
+        if stage_outputs.final_output_type == "text":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+        elif stage_outputs.final_output_type == "audio":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                audio_tensor = output.multimodal_output["audio"]
+                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/tmp_test/extract_prompts.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+import argparse
+from typing import Optional
+
+
+def extract_prompt(line: str) -> Optional[str]:
+    # 提取第一个 '|' 与第二个 '|' 之间的内容
+    i = line.find("|")
+    if i == -1:
+        return None
+    j = line.find("|", i + 1)
+    if j == -1:
+        return None
+    return line[i + 1 : j].strip()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
+    parser.add_argument("--output", "-o", required=True, help="Output file path")
+    parser.add_argument(
+        "--topk",
+        "-k",
+        type=int,
+        default=100,
+        help="Extract the top K prompts (default: 100)",
+    )
+    args = parser.parse_args()
+
+    prompts = []
+    with open(args.input, encoding="utf-8", errors="ignore") as f:
+        for line in f:
+            if len(prompts) >= args.topk:
+                break
+            p = extract_prompt(line.rstrip("\n"))
+            if p:
+                prompts.append(p)
+
+    with open(args.output, "w", encoding="utf-8") as f:
+        for p in prompts:
+            f.write(p + "\n")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/tmp_test/processing_omni.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+temporal_patch_size = 2
+spatial_patch_size = 14
+spatial_merge_size = 2
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = image_obj.convert("RGB")
+    # resize
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", MIN_PIXELS)
+        max_pixels = ele.get("max_pixels", MAX_PIXELS)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=size_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+
+def smart_nframes(
+    ele: dict,
+    total_frames: int,
+    video_fps: int | float,
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            support either `fps` or `nframes`:
+                - nframes: the number of frames to extract for model inputs.
+                - fps: the fps to extract frames for model inputs.
+                    - min_frames: the minimum number of frames of the video,
+                      only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video,
+                      only used when fps is provided.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    Returns:
+        int: the number of frames for video used for model inputs.
+    """
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(
+            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
+            FRAME_FACTOR,
+        )
+        nframes = total_frames / video_fps * fps
+        nframes = min(max(nframes, min_frames), max_frames)
+        nframes = round_by_factor(nframes, FRAME_FACTOR)
+    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using torchvision.io.read_video
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    total_duration = round(total_frames / video_fps, 3)
+    logger.info(
+        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+    video = video[idx]
+    return video, total_duration, nframes
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+def _read_video_decord(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using decord.VideoReader
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    import decord
+
+    video_path = ele["video"]
+    st = time.time()
+    vr = decord.VideoReader(video_path)
+    # TODO: support start_pts and end_pts
+    if "video_start" in ele or "video_end" in ele:
+        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    total_duration = round(total_frames / video_fps, 3)
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    return video, total_duration, nframes
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
+    #       file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::FRAME_FACTOR]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+        nframes, _, height, width = video.shape
+        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+        max_pixels = max(
+            min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
+            int(min_pixels * 1.05),
+        )
+        max_pixels = ele.get("max_pixels", max_pixels)
+        # min_pixels = (factor ** 2) * 52
+        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"],
+                ele["resized_width"],
+                factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=image_factor,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        ).float()
+        return video, total_dur, nframes, second_per_grid
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
+            for video_element in ele["video"]
+        ]
+        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        return images, None, None, None
+
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
+    vision_infos = extract_vision_info(conversations)
+    # Read images or videos
+    image_inputs = []
+    video_inputs = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            video_inputs.append(fetch_video(vision_info))
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+    return image_inputs, video_inputs

examples/offline_inference/tmp_test/ray_head.sh

@@ -0,0 +1 @@
+ray start --head --port=6399 --num-gpus=1

examples/offline_inference/tmp_test/ray_worker.sh

@@ -0,0 +1 @@
+ray start --address=10.90.67.86:6399

examples/offline_inference/tmp_test/run_multiple_prompts.sh

@@ -0,0 +1,8 @@
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --txt-prompts top100.txt

examples/offline_inference/tmp_test/run_single_prompt.sh

@@ -0,0 +1,11 @@
+export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
+CUDA_VISIBLE_DEVICES=4 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+				 --worker-backend ray \
+				 --ray-address auto \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

examples/offline_inference/tmp_test/utils.py

@@ -0,0 +1,312 @@
+import tempfile
+from typing import Optional, Union
+from urllib.request import urlopen
+
+import librosa
+import requests
+import resampy
+import soundfile as sf
+import torch
+import torchvision.io
+from processing_omni import fetch_image, fetch_video
+from transformers import AutoConfig, AutoProcessor
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+# Simple caches to avoid repeated heavy HF loads per prompt
+_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
+_CONFIG_CACHE: dict[str, "AutoConfig"] = {}
+
+
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def resample_wav_to_16khz(input_filepath):
+    data, original_sample_rate = sf.read(input_filepath)
+    # Only use the first channel
+    if len(data.shape) > 1:
+        data = data[:, 0]
+    # resample to 16kHz
+    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
+    return data_resampled
+
+
+def fetch_and_read_video(args, video_url: str, fps=2):
+    def read_video_with_torchvision(video_file_name: str):
+        video, audio, info = torchvision.io.read_video(
+            video_file_name,
+            start_pts=0.0,
+            end_pts=None,
+            pts_unit="sec",
+            output_format="TCHW",
+        )
+
+        total_frames, video_fps = video.size(0), info["video_fps"]
+        total_duration = round(total_frames / video_fps, 3)
+        nframes = int(total_frames / video_fps * fps)
+
+        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::2]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+
+        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
+        video = video[idx]
+
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid.item()]
+        else:
+            return video
+
+    def read_video_with_transformers(video_file_name: Union[str, list[str]]):
+        video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name})
+        if total_duration is None and nframes is None:
+            nframes = len(video)
+            total_duration = 0.5 * nframes
+            second_per_grid = 1.0
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid]
+        else:
+            return video
+
+    def read_video(video_file_name: str):
+        if args.use_torchvision:
+            return read_video_with_torchvision(video_file_name)
+        else:
+            return read_video_with_transformers(video_file_name)
+
+    if isinstance(video_url, str) and video_url.startswith("http"):
+        with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+            resp = requests.get(video_url)
+            assert resp.status_code == requests.codes.ok, (
+                f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
+            )
+
+            temp_video_file.write(urlopen(video_url).read())
+            temp_video_file_path = temp_video_file.name
+            video_file_name = temp_video_file_path
+            return read_video(video_file_name)
+    else:
+        video_file_name = video_url
+        return read_video(video_file_name)
+
+
+def make_inputs_qwen2_omni(
+    args,
+    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
+    use_audio_in_video: Optional[bool] = False,
+    tokenize: bool = False,
+) -> Union[OmniTokensPrompt, TextPrompt]:
+    from transformers import AutoConfig, AutoProcessor
+
+    # Cached processor/config to prevent per-prompt reloading and repeated warnings
+    if args.model not in _PROCESSOR_CACHE:
+        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
+    processor = _PROCESSOR_CACHE[args.model]
+
+    config = _CONFIG_CACHE.get(args.model)
+    if config is None:
+        try:
+            config = AutoConfig.from_pretrained(args.model)
+        except Exception:
+            config = None
+        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
+
+    # Decide legacy flag only once based on config (default True if unknown)
+    if getattr(args, "legacy_omni_video", None) is None:
+        if config is not None and hasattr(config, "architectures"):
+            args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
+        else:
+            args.legacy_omni_video = True
+
+    audios, images, videos = [], [], []
+    for message in messages:
+        if not isinstance(message["content"], list):
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": message["content"],
+                }
+            ]
+        index, num_contents = 0, len(message["content"])
+        while index < num_contents:
+            ele = message["content"][index]
+            if "type" not in ele:
+                if "text" in ele:
+                    ele["type"] = "text"
+                elif "audio" in ele:
+                    ele["type"] = "audio"
+                elif "audio_url" in ele:
+                    ele["type"] = "audio_url"
+                elif "image" in ele:
+                    ele["type"] = "image"
+                elif "image_url" in ele:
+                    ele["type"] = "image_url"
+                elif "video" in ele:
+                    ele["type"] = "video"
+                elif "video_url" in ele:
+                    ele["type"] = "video_url"
+                else:
+                    raise ValueError(f"Unknown ele: {ele}")
+
+            if ele["type"] == "audio" or ele["type"] == "audio_url":
+                if "audio_url" in ele:
+                    audio_key = "audio_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
+                        temp_audio_file.write(urlopen(ele[audio_key]).read())
+                        temp_audio_file_path = temp_audio_file.name
+                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
+                        ele["audio"] = temp_audio_file_path
+                elif "audio" in ele:
+                    audio_key = "audio"
+                    audios.append(resample_wav_to_16khz(ele[audio_key]))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+            elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"):
+                # use video as audio as well
+                if "video_url" in ele:
+                    audio_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele[audio_key]).read())
+                        temp_video_file_path = temp_video_file.name
+                        ele[audio_key] = temp_video_file_path
+                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
+                        videos.append(fetch_and_read_video(args, temp_video_file_path))
+                        ele["video"] = temp_video_file_path
+                elif "video" in ele:
+                    audio_key = "video"
+                    audios.append(librosa.load(ele[audio_key], sr=16000)[0])
+                    videos.append(fetch_and_read_video(args, audio_key))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+                # insert a audio after the video
+                message["content"].insert(
+                    index + 1,
+                    {
+                        "type": "audio",
+                        "audio": ele[audio_key],
+                    },
+                )
+                # no need to load the added audio again
+                index += 1
+            elif ele["type"] == "video" or ele["type"] == "video_url":
+                if "video_url" in ele:
+                    video_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele["video_url"]).read())
+                        temp_video_file_path = temp_video_file.name
+                        videos.append(fetch_and_read_video(args, temp_video_file))
+                        ele["video"] = temp_video_file_path
+                else:
+                    video_key = "video"
+                    videos.append(fetch_and_read_video(args, ele[video_key]))
+            elif ele["type"] == "image" or ele["type"] == "image_url":
+                images.append(fetch_image(ele))
+
+            # move to the next content
+            index += 1
+
+    prompt = processor.apply_chat_template(
+        messages,
+        tokenize=tokenize,
+        add_generation_prompt=True,
+        add_vision_id=True,
+    )
+
+    audios = audios if len(audios) > 0 else None
+    images = images if len(images) > 0 else None
+    videos = videos if len(videos) > 0 else None
+
+    multi_modal_data = {}
+    if audios:
+        multi_modal_data["audio"] = audios
+    if images:
+        multi_modal_data["image"] = images
+    if videos:
+        multi_modal_data["video"] = videos
+
+    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
+        prompt = prompt[0]
+
+    if tokenize:
+        return OmniTokensPrompt(
+            prompt_token_ids=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+    else:
+        return TextPrompt(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+
+
+def make_text_prompt(args, prompt):
+    messages = [
+        get_system_prompt(),
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+
+    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
+    return prompt
+
+
+def make_audio_in_video_v2_prompt(args):
+    messages = [
+        {
+            "role": "system",
+            "content": [
+                {
+                    "type": "text",
+                    "text": (
+                        "You are Qwen, a virtual human developed by the Qwen Team, "
+                        "Alibaba Group, capable of perceiving auditory and visual "
+                        "inputs, as well as generating text and speech."
+                    ),
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "video_url",
+                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
+                },
+            ],
+        },
+    ]
+    prompt = make_inputs_qwen2_omni(
+        args,
+        messages,
+        use_audio_in_video=True,
+        tokenize=args.tokenize,
+    )
+    return prompt
+
+
+def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]:
+    if args.prompt_type == "text":
+        prompt = make_text_prompt(args, prompt)
+    elif args.prompt_type == "audio-in-video-v2":
+        prompt = make_audio_in_video_v2_prompt(args)
+    else:
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt

vllm_omni/distributed/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .connectors import (
+    ConnectorSpec, OmniTransferConfig,
+    OmniConnectorBase, InMemoryOmniConnector,
+    OmniConnectorFactory, MooncakeConnector,
+    load_omni_transfer_config, create_simple_config, create_mooncake_config
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+
+    # Connectors
+    "OmniConnectorBase",
+    "InMemoryOmniConnector",
+    "OmniConnectorFactory",
+    "MooncakeConnector",
+
+    # Utilities
+    "load_omni_transfer_config",
+    "create_simple_config",
+    "create_mooncake_config",
+]

vllm_omni/distributed/connectors/__init__.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .config import ConnectorSpec, OmniTransferConfig
+from .base import OmniConnectorBase, InMemoryOmniConnector
+from .factory import OmniConnectorFactory
+from .mooncake_connector import MooncakeConnector
+from .utils import (
+    load_omni_transfer_config,
+    initialize_connectors_from_config,
+    create_simple_config,
+    create_mooncake_config,
+    get_connectors_config_for_stage,
+    initialize_orchestrator_connectors,
+    get_stage_connector_config,
+    build_stage_connectors,
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+
+    # Base classes and implementations
+    "OmniConnectorBase",
+    "InMemoryOmniConnector",
+
+    # Factory
+    "OmniConnectorFactory",
+
+    # Specific implementations
+    "MooncakeConnector",
+
+    # Utilities
+    "load_omni_transfer_config",
+    "initialize_connectors_from_config",
+    "create_simple_config",
+    "create_mooncake_config",
+    "get_connectors_config_for_stage",
+
+    # Manager helpers
+    "initialize_orchestrator_connectors",
+    "get_stage_connector_config",
+    "build_stage_connectors",
+]

vllm_omni/distributed/connectors/adapter.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py 
+# and vllm_omni.entrypoints.omni_llm.py
+
+import time
+import logging
+from typing import Any, Optional, Dict, Callable
+
+logger = logging.getLogger(__name__)
+
+
+def try_send_via_connector(
+    connector: Any,
+    stage_id: int,
+    next_stage_id: int,
+    req_id: str,
+    next_inputs: Any,
+    sampling_params: Any,
+    original_prompt: Any,
+    next_stage_queue_submit_fn: Callable[[Dict[str, Any]], None],
+    metrics: Any
+) -> bool:
+    """
+    Attempts to send data via OmniConnector.
+    Returns True if successful, False otherwise.
+    Encapsulates the logic of preparing payload, sending via connector, 
+    sending notification, and recording metrics.
+    """
+    try:
+        t0 = time.time()
+
+        # Prepare data for connector
+        payload_data = {
+            "engine_inputs": next_inputs,
+            "sampling_params": sampling_params,
+            "metadata": {
+                "original_prompt": original_prompt,
+                "stage_transition": f"{stage_id}->{next_stage_id}",
+                "timestamp": time.time()
+            }
+        }
+
+        # Send data via connector
+        success, serialized_size = connector.put(str(stage_id), str(next_stage_id), str(req_id), payload_data)
+        
+        if success:
+            # Send lightweight notification via queue
+            notify_payload = {
+                "request_id": req_id,
+                "sampling_params": sampling_params,
+                "from_connector": True,
+                "from_stage": str(stage_id),
+                "to_stage": str(next_stage_id),
+                "sent_ts": time.time(),
+            }
+            next_stage_queue_submit_fn(notify_payload)
+
+            t1 = time.time()
+            tx_ms = (t1 - t0) * 1000.0
+            
+            metrics.on_forward(
+                stage_id,
+                next_stage_id,
+                req_id,
+                serialized_size,  # Use size from connector
+                float(tx_ms),
+                True,  # Mark as using connector
+            )
+            return True
+        else:
+            # If put returned False, we let the caller handle fallback
+            return False
+
+    except Exception as e:
+        logger.warning(
+            "[Orchestrator] OmniConnector failed for req %s: %s; falling back to queue",
+            req_id,
+            e,
+        )
+        return False
+
+
+def try_recv_via_connector(
+    task: Dict[str, Any], 
+    connectors: Dict[Any, Any], 
+    stage_id: int,
+) -> tuple[Any, Optional[Dict[str, Any]]]:
+    """
+    Attempts to resolve input data from either connector or IPC.
+    Returns (engine_inputs, rx_metrics) or (None, None) if failed/skipped.
+    """
+    rid = task["request_id"]
+    
+    if task.get("from_connector"):
+        from_stage = task.get("from_stage")
+        to_stage = str(stage_id)
+
+        if not from_stage:
+            logger.error(
+                "[Stage-%s] 'from_connector' is true but 'from_stage' is missing for request %s", stage_id, rid
+            )
+            return None, None
+
+        # Get connector for this edge
+        connector_key = (from_stage, to_stage)
+        connector = connectors.get(connector_key)
+
+        if connector:
+            try:
+                # Get data from connector with timeout
+                _t_start = time.time()
+                payload_data = connector.get(from_stage, to_stage, str(rid))
+                _t_end = time.time()
+                
+                if payload_data and isinstance(payload_data, dict):
+                    ein = payload_data.get("engine_inputs")
+                    # Use connector's serialization size for consistency
+                    serialized_size = len(connector.serialize_obj(payload_data))
+                    decode_ms = (_t_end - _t_start) * 1000.0
+                    
+                    rx_metrics = {
+                        "rx_decode_time_ms": decode_ms,
+                        "rx_transfer_bytes": serialized_size
+                    }
+                    return ein, rx_metrics
+                else:
+                    logger.error(
+                        "[Stage-%s] Failed to get data from connector for request %s or payload is empty", stage_id, rid
+                    )
+                    return None, None
+            except Exception as e:
+                logger.error(
+                    "[Stage-%s] Error retrieving data from connector for request %s: %s", stage_id, rid, e
+                )
+                return None, None
+        else:
+            logger.error(
+                "[Stage-%s] No connector found for edge %s -> %s for request %s", stage_id, from_stage, to_stage, rid
+            )
+            return None, None
+    else:
+        # Data comes from queue as usual
+        #return maybe_load_from_ipc_with_metrics_fn(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+        return None, None
+

vllm_omni/distributed/connectors/base.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class OmniConnectorBase(ABC):
+    """Base class for all OmniConnectors."""
+
+    @abstractmethod
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        """Store Python object, internal serialization handled by connector.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+            data: Python object to store
+
+        Returns:
+            tuple: (success: bool, serialized_size: int)
+        """
+        pass
+
+    @abstractmethod
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        """Retrieve Python object, internal deserialization handled by connector.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+
+        Returns:
+            Python object if found, None otherwise
+        """
+        pass
+
+    @abstractmethod
+    def cleanup(self, request_id: str) -> None:
+        """Clean up resources for a request."""
+        pass
+
+    @abstractmethod
+    def health(self) -> dict[str, Any]:
+        """Return health status and metrics."""
+        pass
+
+    @staticmethod
+    def serialize_obj(obj: Any) -> bytes:
+        """Serialize a Python object to bytes using torch.save (consistent with stage_utils)."""
+        import io
+        import torch
+        buf = io.BytesIO()
+        torch.save(obj, buf)
+        return buf.getvalue()
+
+    @staticmethod
+    def deserialize_obj(data: bytes) -> Any:
+        """Deserialize bytes to Python object using torch.load (consistent with stage_utils)."""
+        import io
+        import torch
+        return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
+
+
+class InMemoryOmniConnector(OmniConnectorBase):
+    """Simple in-memory connector for testing and co-located stages."""
+
+    def __init__(self):
+        self._storage: dict[str, Any] = {}
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "errors": 0,
+        }
+
+    def _make_key(self, from_stage: str, to_stage: str, request_id: str) -> str:
+        return f"{from_stage}:{to_stage}:{request_id}"
+
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        try:
+            key = self._make_key(from_stage, to_stage, request_id)
+            self._storage[key] = data
+            serialized_size = len(self.serialize_obj(data))
+            self._metrics["puts"] += 1
+            logger.debug(f"InMemoryConnector: stored {key}")
+            return True, serialized_size
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"InMemoryConnector put failed: {e}")
+            return False, 0
+
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        try:
+            key = self._make_key(from_stage, to_stage, request_id)
+            data = self._storage.get(key)
+            if data is not None:
+                self._metrics["gets"] += 1
+                return data
+            return None
+        except Exception as e:
+            logger.error(f"InMemoryConnector get failed: {e}")
+            return None
+
+    def cleanup(self, request_id: str) -> None:
+        keys_to_remove = [k for k in self._storage.keys() if k.endswith(f":{request_id}")]
+        for key in keys_to_remove:
+            del self._storage[key]
+        logger.debug(f"InMemoryConnector: cleaned up {len(keys_to_remove)} keys for {request_id}")
+
+    def health(self) -> dict[str, Any]:
+        return {
+            "status": "healthy",
+            "storage_size": len(self._storage),
+            **self._metrics,
+        }

vllm_omni/distributed/connectors/config.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ConnectorSpec:
+    """Specification for a connector instance."""
+    name: str  # e.g., "MooncakeConnector", "InMemoryConnector"
+    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific config
+
+
+@dataclass
+class OmniTransferConfig:
+    """Top-level configuration for OmniConnector system."""
+    # Direct mapping: (from_stage, to_stage) -> connector
+    connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict)
+    default_connector: Optional[ConnectorSpec] = None
+
+    def get_connector_for_edge(self, from_stage: str, to_stage: str) -> Optional[ConnectorSpec]:
+        """Get connector spec for a specific edge."""
+        edge_key = (from_stage, to_stage)
+        return self.connectors.get(edge_key, self.default_connector)
+
+    def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool:
+        """Check if there's a connector configured for the edge."""
+        return self.get_connector_for_edge(from_stage, to_stage) is not None

vllm_omni/distributed/connectors/connector_tests/run_single_prompt.sh

@@ -0,0 +1,9 @@
+export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
+CUDA_VISIBLE_DEVICES=0,1,2,3 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

vllm_omni/distributed/connectors/connector_tests/test_connector_integration.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""
+Test OmniConnector integration with OmniLLM.
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(__file__))
+
+from vllm_omni.distributed.connectors import create_mooncake_config
+
+def test_connector_config():
+    """Test connector configuration creation."""
+    print("Testing connector configuration...")
+
+    # Create Mooncake config
+    config = create_mooncake_config(
+        host="127.0.0.1",
+        metadata_server="http://127.0.0.1:8080/metadata",
+        master="127.0.0.1:50051"
+    )
+
+    print(f"Created config with {len(config.connectors)} connectors")
+    for edge, connector_spec in config.connectors.items():
+        print(f"  {edge}: {connector_spec.name}")
+
+    return config
+
+def test_omni_llm_initialization():
+    """Test OmniLLM with connector config."""
+    print("\nTesting OmniLLM initialization with connectors...")
+
+    try:
+        from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+        # Create connector config
+        connector_config = {
+            "connectors": {
+                "0_to_1": {
+                    "name": "InMemoryConnector",  # Use InMemory for testing
+                    "extra": {}
+                }
+            }
+        }
+
+        # Test connector initialization logic
+        omni = OmniLLM.__new__(OmniLLM)  # Create without calling __init__
+        omni.omni_transfer_config = connector_config
+        omni.connectors = {}
+        omni._initialize_connectors()
+
+        # Check if connector was created
+        if ("0", "1") in omni.connectors:
+            print(f"✅ Connector created: {type(omni.connectors[('0', '1')]).__name__}")
+            return True
+        else:
+            print("❌ Connector not created")
+            return False
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+if __name__ == "__main__":
+    print("OmniConnector Integration Test")
+    print("=" * 40)
+
+    # Test config creation
+    config = test_connector_config()
+
+    # Test OmniLLM integration
+    success = test_omni_llm_initialization()
+
+    if success:
+        print("\n✅ Integration test passed!")
+        print("Connector system is ready for use with OmniLLM")
+    else:
+        print("\n❌ Integration test failed!")

vllm_omni/distributed/connectors/connector_tests/test_epdg_pipeline.py

@@ -0,0 +1,332 @@
+#!/usr/bin/env python3
+"""
+Test EPD (Encode-Prefill-Decode) 3-stage pipeline with connectors.
+Tests data flow between stages using configured connectors.
+"""
+
+import os
+import sys
+from typing import Any, Dict, Optional
+from dataclasses import dataclass
+import uuid
+
+try:
+    import yaml
+except ImportError:
+    print("PyYAML not available, using basic YAML parsing")
+    yaml = None
+
+# Add the connectors package to path
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.insert(0, parent_dir)
+
+# Import from parent package
+try:
+    from connectors.config import ConnectorSpec, OmniTransferConfig
+    from connectors.base import OmniConnectorBase, InMemoryOmniConnector
+    from connectors.factory import OmniConnectorFactory
+    from connectors.utils import load_omni_transfer_config
+except ImportError:
+    # Fallback for direct execution
+    sys.path.insert(0, current_dir)
+    from config import ConnectorSpec, OmniTransferConfig
+    from base import OmniConnectorBase, InMemoryOmniConnector
+    from factory import OmniConnectorFactory
+    from utils import load_omni_transfer_config
+
+
+@dataclass
+class StageResult:
+    """Result from a stage processing."""
+    stage_name: str
+    request_id: str
+    output_data: Any
+    metadata: Dict[str, Any]
+
+
+class EncodeStage:
+    """Simple Encode stage for testing."""
+
+    def __init__(self, stage_id: int, output_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "encode"
+        self.output_connector = output_connector
+
+    def process(self, request_id: str, input_data: Any) -> StageResult:
+        """Process input data and optionally send to next stage."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Simple encoding simulation
+        encoded_data = {
+            "original_text": input_data,
+            "encoded_tokens": [ord(c) for c in str(input_data)][:10],  # Simple tokenization
+            "embeddings": [0.1 * i for i in range(10)],  # Mock embeddings
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=encoded_data,
+            metadata={"tokens_count": len(encoded_data["encoded_tokens"])}
+        )
+
+        # Send to next stage if connector is available
+        if self.output_connector:
+            success = self.output_connector.put(
+                str(self.stage_id), "1", request_id, encoded_data
+            )
+            print(f"[{self.stage_name}] Sent data to stage 1 via connector: {success}")
+
+        return result
+
+
+class PrefillStage:
+    """Simple Prefill stage for testing."""
+
+    def __init__(self, stage_id: int,
+                 input_connector: Optional[OmniConnectorBase] = None,
+                 output_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "prefill"
+        self.input_connector = input_connector
+        self.output_connector = output_connector
+
+    def process(self, request_id: str) -> Optional[StageResult]:
+        """Process data from previous stage."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Try to get data from input connector
+        input_data = None
+        if self.input_connector:
+            input_data = self.input_connector.get("0", str(self.stage_id), request_id)
+
+        if input_data is None:
+            print(f"[{self.stage_name}] No input data received for {request_id}")
+            return None
+
+        print(f"[{self.stage_name}] Received data from encode")
+
+        # Simple prefill simulation
+        prefills_data = {
+            **input_data,
+            "kv_cache": {"layer_0": [0.2 * i for i in range(20)]},  # Mock KV cache
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=prefills_data,
+            metadata={"kv_size": len(prefills_data["kv_cache"])}
+        )
+
+        # Send to next stage if connector is available
+        if self.output_connector:
+            success = self.output_connector.put(
+                str(self.stage_id), "2", request_id, prefills_data
+            )
+            print(f"[{self.stage_name}] Sent data to stage 2 via connector: {success}")
+
+        return result
+
+
+class DecodeStage:
+    """Simple Decode stage for testing (final stage in 3-stage pipeline)."""
+
+    def __init__(self, stage_id: int,
+                 input_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "decode"
+        self.input_connector = input_connector
+
+    def process(self, request_id: str) -> Optional[StageResult]:
+        """Process data from previous stage (final output)."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Try to get data from input connector
+        input_data = None
+        if self.input_connector:
+            input_data = self.input_connector.get("1", str(self.stage_id), request_id)
+
+        if input_data is None:
+            print(f"[{self.stage_name}] No input data received for {request_id}")
+            return None
+
+        print(f"[{self.stage_name}] Received data from prefill")
+
+        # Simple decode simulation (final output)
+        decoded_data = {
+            **input_data,
+            "hidden_states": [0.3 * i for i in range(15)],  # Mock hidden states
+            "final_output": f"Decoded text from: {input_data.get('original_text', 'unknown')}",
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=decoded_data,
+            metadata={"hidden_size": len(decoded_data["hidden_states"]), "final": True}
+        )
+
+        return result
+
+
+class GeneratorStage:
+    """Simple Generator stage for testing (removed for 3-stage pipeline)."""
+    pass  # Not used in current 3-stage setup
+
+
+def load_config_from_yaml(yaml_path: str) -> OmniTransferConfig:
+    """Load connector configuration from YAML file."""
+    if yaml is None:
+        # Fallback: create a simple test config
+        print("Using fallback InMemoryConnector config for testing")
+        return create_test_config()
+
+    with open(yaml_path, 'r', encoding='utf-8') as f:
+        config_dict = yaml.safe_load(f)
+
+    try:
+        config = load_omni_transfer_config(config_dict=config_dict)
+        print(f"Successfully loaded config from YAML with {len(config.connectors)} connectors")
+        return config
+    except Exception as e:
+        print(f"Failed to load YAML config: {e}, using fallback config")
+        return create_test_config()
+
+
+def create_test_config() -> OmniTransferConfig:
+    """Create a test configuration with InMemoryConnectors."""
+    from .config import ConnectorSpec
+
+    connectors = {}
+    # Create connectors for EPD (3-stage) pipeline: 0->1->2
+    pipeline_stages = [("0", "1"), ("1", "2")]
+
+    for from_stage, to_stage in pipeline_stages:
+        connectors[(from_stage, to_stage)] = ConnectorSpec(
+            name="InMemoryConnector",
+            extra={}
+        )
+
+    return OmniTransferConfig(connectors=connectors)
+
+
+def create_stages_with_connectors(config: OmniTransferConfig) -> Dict[str, Any]:
+    """Create EPD (3-stage) pipeline with appropriate connectors."""
+
+    # Create connector instances based on config
+    connectors = {}
+    for edge_key, connector_spec in config.connectors.items():
+        from_stage, to_stage = edge_key
+        connectors[edge_key] = OmniConnectorFactory.create_connector(connector_spec)
+
+    # Create stages with connectors (3-stage pipeline: 0->1->2)
+    stages = {}
+
+    # Encode stage (stage 0)
+    encode_output_connector = connectors.get(("0", "1"))  # to prefill
+    stages["encode"] = EncodeStage(0, encode_output_connector)
+
+    # Prefill stage (stage 1)
+    prefill_input_connector = connectors.get(("0", "1"))   # from encode
+    prefill_output_connector = connectors.get(("1", "2"))  # to decode
+    stages["prefill"] = PrefillStage(1, prefill_input_connector, prefill_output_connector)
+
+    # Decode stage (stage 2) - final stage
+    decode_input_connector = connectors.get(("1", "2"))    # from prefill
+    stages["decode"] = DecodeStage(2, decode_input_connector)
+
+    return stages
+
+
+def test_pipeline_with_config(yaml_path: str):
+    """Test the EPDG pipeline with connector configuration."""
+
+    print("🔧 Loading configuration from YAML...")
+    try:
+        config = load_config_from_yaml(yaml_path)
+        print(f"✓ Loaded config with {len(config.connectors)} connector configurations")
+    except Exception as e:
+        print(f"✗ Failed to load config: {e}")
+        return False
+
+    print("\n🏗️  Creating EPDG stages with connectors...")
+    try:
+        stages = create_stages_with_connectors(config)
+        print("✓ Created all stages with connectors")
+    except Exception as e:
+        print(f"✗ Failed to create stages: {e}")
+        return False
+
+    print("\n🚀 Testing data flow through pipeline...")
+
+    # Test request
+    request_id = f"test_req_{uuid.uuid4()}"
+    input_text = "Hello, EPDG pipeline!"
+
+    try:
+        # Stage 1: Encode
+        print(f"\n--- Stage 1: Encode ---")
+        encode_result = stages["encode"].process(request_id, input_text)
+        print(f"✓ Encode completed: {encode_result.metadata}")
+
+        # Stage 2: Prefill
+        print(f"\n--- Stage 2: Prefill ---")
+        prefill_result = stages["prefill"].process(request_id)
+        if prefill_result:
+            print(f"✓ Prefill completed: {prefill_result.metadata}")
+        else:
+            print("✗ Prefill failed - no input data")
+            return False
+
+        # Stage 3: Decode
+        print(f"\n--- Stage 3: Decode ---")
+        decode_result = stages["decode"].process(request_id)
+        if decode_result:
+            print(f"✓ Decode completed: {decode_result.metadata}")
+        else:
+            print("✗ Decode failed - no input data")
+            return False
+
+        # Decode stage produces final output in 3-stage pipeline
+
+    except Exception as e:
+        print(f"✗ Pipeline test failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+    print("\n🎉 Pipeline test completed successfully!")
+    return True
+
+
+def main():
+    """Main test function."""
+    # Path to the YAML config file
+    yaml_path = "/workspace/omni/vllm-omni/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml"
+
+    if not os.path.exists(yaml_path):
+        print(f"✗ YAML config file not found: {yaml_path}")
+        return False
+
+    print("Testing EPDG Pipeline with Connectors")
+    print("=" * 50)
+    print(f"Config file: {yaml_path}")
+
+    success = test_pipeline_with_config(yaml_path)
+
+    if success:
+        print("\n✅ All tests passed!")
+    else:
+        print("\n❌ Tests failed!")
+
+    return success
+
+
+if __name__ == "__main__":
+    success = main()
+    sys.exit(0 if success else 1)

vllm_omni/distributed/connectors/factory.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+try:
+    from .base import OmniConnectorBase
+    from .config import ConnectorSpec
+except ImportError:
+    # Fallback for direct execution
+    import sys
+    import os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+    from connectors.base import OmniConnectorBase
+    from connectors.config import ConnectorSpec
+
+logger = logging.getLogger(__name__)
+
+
+class OmniConnectorFactory:
+    """Factory for creating OmniConnectors."""
+
+    _registry: dict[str, Callable[[dict[str, Any]], OmniConnectorBase]] = {}
+
+    @classmethod
+    def register_connector(
+        cls,
+        name: str,
+        constructor: Callable[[dict[str, Any]], OmniConnectorBase]
+    ) -> None:
+        """Register a connector constructor."""
+        if name in cls._registry:
+            raise ValueError(f"Connector '{name}' is already registered.")
+        cls._registry[name] = constructor
+        logger.debug(f"Registered connector: {name}")
+
+    @classmethod
+    def create_connector(cls, spec: ConnectorSpec) -> OmniConnectorBase:
+        """Create a connector from specification."""
+        if spec.name not in cls._registry:
+            raise ValueError(f"Unknown connector: {spec.name}. Available: {list(cls._registry.keys())}")
+
+        constructor = cls._registry[spec.name]
+        try:
+            connector = constructor(spec.extra)
+            logger.info(f"Created connector: {spec.name}")
+            return connector
+        except Exception as e:
+            logger.error(f"Failed to create connector {spec.name}: {e}")
+            raise
+
+    @classmethod
+    def list_registered_connectors(cls) -> list[str]:
+        """List all registered connector names."""
+        return list(cls._registry.keys())
+
+
+# Register built-in connectors with lazy imports
+def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .mooncake_connector import MooncakeConnector
+    except ImportError:
+        # Fallback import
+        import sys
+        import os
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from connectors.mooncake_connector import MooncakeConnector
+    return MooncakeConnector(config)
+
+
+def _create_in_memory_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .base import InMemoryOmniConnector
+    except ImportError:
+        # Fallback import
+        import sys
+        import os
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from connectors.base import InMemoryOmniConnector
+    return InMemoryOmniConnector()
+
+
+# Register connectors
+OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector)
+OmniConnectorFactory.register_connector("InMemoryConnector", _create_in_memory_connector)

vllm_omni/distributed/connectors/mooncake_connector.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import io
+import logging
+import time
+from typing import Any, Optional
+
+import torch
+
+from .base import OmniConnectorBase
+
+logger = logging.getLogger(__name__)
+
+try:
+    from mooncake.store import MooncakeDistributedStore, ReplicateConfig
+except ImportError:
+    try:
+        from mooncake import MooncakeDistributedStore, ReplicateConfig
+    except ImportError:
+        logger.warning("Mooncake not available, MooncakeOmniConnector will not work")
+        MooncakeDistributedStore = None
+        ReplicateConfig = None
+
+
+class MooncakeConnector(OmniConnectorBase):
+    """Mooncake-based distributed connector for OmniConnector."""
+
+    def __init__(self, config: dict[str, Any]):
+        if MooncakeDistributedStore is None:
+            raise ImportError("Mooncake not available")
+
+        self.config = config
+        self.host = config.get("host", "127.0.0.1")
+        self.metadata = config.get("metadata_server", "http://127.0.0.1:8080/metadata")
+        self.master = config.get("master", "127.0.0.1:50051")
+        self.segment = config.get("segment", 512 * 1024 * 1024)  # 512MB
+        self.localbuf = config.get("localbuf", 64 * 1024 * 1024)  # 64MB
+        self.proto = config.get("proto", "tcp")
+        self.rdma = config.get("rdma", "")
+
+        self.store: Optional[MooncakeDistributedStore] = None
+        self.pin: Optional[ReplicateConfig] = None
+
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "errors": 0,
+            "timeouts": 0,
+        }
+
+        self._init_store()
+
+    def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str:
+        """Generate store key for request between stages."""
+        return f"{rid}/{from_stage}_to_{to_stage}"
+
+    def _init_store(self):
+        """Initialize Mooncake store."""
+        try:
+            self.store = MooncakeDistributedStore()
+            rc = self.store.setup(
+                self.host, self.metadata, self.segment,
+                self.localbuf, self.proto, self.rdma, self.master
+            )
+            if rc != 0:
+                raise RuntimeError(f"Mooncake setup failed: {rc}")
+
+            self.pin = ReplicateConfig()
+            self.pin.with_soft_pin = True
+            logger.info("MooncakeConnector initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize Mooncake store: {e}")
+            raise
+
+    # Use base class serialization methods for consistency
+
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return False, 0
+
+        try:
+            serialized_data = self.serialize_obj(data)
+            key = self._make_key(request_id, from_stage, to_stage)
+            self.store.put(key, serialized_data, self.pin)
+
+            self._metrics["puts"] += 1
+            self._metrics["bytes_transferred"] += len(serialized_data)
+
+            logger.debug(f"MooncakeConnector: stored {key}, {len(serialized_data)} bytes")
+            print(f"========= mooncake put {from_stage}-{to_stage}-{request_id}, with size {len(serialized_data)}")
+            return True, len(serialized_data)
+
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"MooncakeConnector put failed: {e}")
+            return False, 0
+
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return None
+
+        retries = 20
+        sleep_s = 0.05
+        key = self._make_key(request_id, from_stage, to_stage)
+
+        for attempt in range(retries):
+            try:
+                raw_data = self.store.get(key)
+
+                if raw_data:
+                    data = self.deserialize_obj(raw_data)
+                    self._metrics["gets"] += 1
+                    print(f"========= mooncake get {from_stage}-{to_stage}-{request_id}, with size {len(raw_data)}")
+                    return data
+
+            except Exception as e:
+                logger.debug(f"MooncakeConnector get attempt {attempt} failed: {e}")
+
+            if attempt < retries - 1:
+                time.sleep(sleep_s)
+
+        self._metrics["timeouts"] += 1
+        logger.warning(f"MooncakeConnector: timeout waiting for {key}")
+        return None
+
+    def cleanup(self, request_id: str) -> None:
+        if not self.store:
+            return
+
+        # Note: Mooncake doesn't have explicit delete, data will be garbage collected
+        # We could implement a cleanup mechanism by storing deletion markers
+        logger.debug(f"MooncakeConnector: cleanup requested for {request_id} (no-op)")
+
+    def health(self) -> dict[str, Any]:
+        if not self.store:
+            return {"status": "unhealthy", "error": "Store not initialized"}
+
+        return {
+            "status": "healthy",
+            "host": self.host,
+            "metadata_server": self.metadata,
+            "master": self.master,
+            **self._metrics,
+        }
+
+    def close(self):
+        """Clean shutdown."""
+        if self.store:
+            try:
+                self.store.close()
+                self.store = None
+                logger.info("MooncakeConnector closed")
+            except Exception as e:
+                logger.error(f"Error closing Mooncake store: {e}")

vllm_omni/distributed/connectors/utils.py

@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utilities for OmniConnector configuration and validation."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+from vllm_omni.distributed.connectors.config import OmniTransferConfig, ConnectorSpec
+from vllm_omni.distributed.connectors.factory import OmniConnectorFactory
+from vllm_omni.distributed.connectors.base import OmniConnectorBase
+
+logger = logging.getLogger(__name__)
+
+
+def initialize_connectors_from_config(
+    config_path: Optional[Union[str, Path]] = None
+) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
+    """
+    Initialize connectors from configuration file.
+
+    Returns:
+        tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
+    """
+    try:
+        transfer_config = load_omni_transfer_config(config_path)
+    except Exception as e:
+        logger.warning(f"Failed to load OmniTransferConfig from {config_path}: {e}")
+        return None, {}
+
+    if not transfer_config:
+        logger.info("No OmniTransferConfig provided, using default IPC")
+        return None, {}
+
+    # 使用统一的连接器创建逻辑
+    connectors = create_connectors_from_config(transfer_config.connectors)
+    return transfer_config, connectors
+
+
+def create_connectors_from_config(
+    connectors_config: dict[tuple[str, str], ConnectorSpec]
+) -> dict[tuple[str, str], OmniConnectorBase]:
+    """通用连接器创建逻辑，从配置字典创建连接器实例。"""
+    connectors = {}
+    try:
+        for edge_key, connector_spec in connectors_config.items():
+            connector = OmniConnectorFactory.create_connector(connector_spec)
+            connectors[edge_key] = connector
+            logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
+    except Exception as e:
+        logger.error(f"Failed to initialize connectors: {e}")
+
+    return connectors
+
+
+def get_connectors_config_for_stage(
+    transfer_config: Optional[OmniTransferConfig],
+    stage_id: Union[str, int]
+) -> dict[str, Any]:
+    """
+    Extract connector configurations relevant for a specific stage worker.
+    
+    Returns a dict compatible with worker initialization:
+    {
+        "from_stage_X": {
+            "spec": {
+                "name": "ConnectorName",
+                "extra": {...}
+            }
+        },
+        ...
+    }
+    """
+    if not transfer_config:
+        return {}
+        
+    stage_connectors_config = {}
+    target_stage = str(stage_id)
+    
+    # Iterate through all configured edges
+    for (from_stage, to_stage), spec in transfer_config.connectors.items():
+        # We only care about incoming edges for the worker process
+        # (Worker needs to create connectors to receive data)
+        if to_stage == target_stage:
+            stage_connectors_config[f"from_stage_{from_stage}"] = {
+                "spec": {
+                    "name": spec.name,
+                    "extra": spec.extra
+                }
+            }
+            
+    return stage_connectors_config
+
+
+def load_omni_transfer_config(
+    config_path: Optional[Union[str, Path]] = None,
+    config_dict: Optional[dict[str, Any]] = None
+) -> Optional[OmniTransferConfig]:
+    """Load OmniTransferConfig from file or dict."""
+    if config_path is None and config_dict is None:
+        return None
+
+    if config_path is not None:
+        config_path = Path(config_path)
+        if not config_path.exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        with open(config_path, 'r', encoding='utf-8') as f:
+            if config_path.suffix.lower() == '.json':
+                config_dict = json.load(f)
+            elif config_path.suffix.lower() in ['.yaml', '.yml']:
+                try:
+                    import yaml
+                    config_dict = yaml.safe_load(f)
+                except ImportError:
+                    raise ImportError("PyYAML required for YAML config files")
+            else:
+                raise ValueError(f"Unsupported config file format: {config_path.suffix}")
+
+    if config_dict is None:
+        return None
+
+    # Parse connectors
+    connectors = {}
+    runtime_config = config_dict.get('runtime', {})
+
+    # Parse global connectors (from runtime.connectors)
+    global_connectors = runtime_config.get('connectors', {})
+    for conn_name, conn_config in global_connectors.items():
+        connector = ConnectorSpec(
+            name=conn_config['name'],
+            extra=conn_config.get('extra', {})
+        )
+        # Store globally defined connectors for reference
+        # These will be resolved to specific edges later
+
+    # Parse stage-level connectors
+    for stage_config in config_dict.get('stage_args', []):
+        stage_id = str(stage_config['stage_id'])
+
+        # Input connectors
+        for input_key, conn_ref in stage_config.get('input_connectors', {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(
+                        name=conn_config['name'],
+                        extra=conn_config.get('extra', {})
+                    )
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(
+                    name=conn_ref['name'],
+                    extra=conn_ref.get('extra', {})
+                )
+
+            # Parse from_stage from key (e.g., "from_stage_0" -> "0")
+            from_stage = input_key.replace("from_stage_", "")
+            edge_key = (from_stage, stage_id)
+            connectors[edge_key] = connector
+
+        # Output connectors
+        for output_key, conn_ref in stage_config.get('output_connectors', {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(
+                        name=conn_config['name'],
+                        extra=conn_config.get('extra', {})
+                    )
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(
+                    name=conn_ref['name'],
+                    extra=conn_ref.get('extra', {})
+                )
+
+            # Parse to_stage from key (e.g., "to_stage_1" -> "1")
+            to_stage = output_key.replace("to_stage_", "")
+            edge_key = (stage_id, to_stage)
+            connectors[edge_key] = connector
+
+    config = OmniTransferConfig(connectors=connectors)
+
+    logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")
+    return config
+
+
+def create_simple_config(
+    connector_name: str = "InMemoryConnector",
+    connector_config: Optional[dict[str, Any]] = None
+) -> OmniTransferConfig:
+    """Create a simple OmniTransferConfig for testing."""
+    if connector_config is None:
+        connector_config = {}
+
+    # Create connectors for typical EPDG pipeline
+    connectors = {}
+    pipeline_edges = [
+        ("0", "1"),  # thinker -> talker
+        ("1", "2"),  # talker -> code2wav
+    ]
+
+    for from_stage, to_stage in pipeline_edges:
+        connectors[(from_stage, to_stage)] = ConnectorSpec(
+            name=connector_name,
+            extra=connector_config
+        )
+
+    return OmniTransferConfig(connectors=connectors)
+
+
+def create_mooncake_config(
+    host: str = "127.0.0.1",
+    metadata_server: str = "http://127.0.0.1:8080/metadata",
+    master: str = "127.0.0.1:50051",
+    segment: int = 512 * 1024 * 1024,
+    localbuf: int = 64 * 1024 * 1024,
+) -> OmniTransferConfig:
+    """Create OmniTransferConfig for Mooncake-based setup."""
+    connector_config = {
+        "host": host,
+        "metadata_server": metadata_server,
+        "master": master,
+        "segment": segment,
+        "localbuf": localbuf,
+    }
+
+    return create_simple_config("MooncakeConnector", connector_config)
+
+
+# High-level management functions (moved from manager.py)
+
+def initialize_orchestrator_connectors(
+    config_path: Optional[str],
+) -> Tuple[Optional[OmniTransferConfig], Dict[tuple[str, str], OmniConnectorBase]]:
+    """Initialize connectors shared at orchestrator level."""
+    try:
+        transfer_config, connectors = initialize_connectors_from_config(config_path)
+        return transfer_config, connectors
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.error("Unexpected error initializing connectors: %s", exc)
+        return None, {}
+
+
+def get_stage_connector_config(
+    transfer_config: Optional[OmniTransferConfig],
+    stage_id: int,
+) -> dict[str, Any]:
+    """Return the serialized connector config payload for a specific stage."""
+    if transfer_config is None:
+        return {}
+
+    try:
+        return get_connectors_config_for_stage(transfer_config, stage_id)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.warning(
+            "Failed to build connector config for stage %s: %s. Using IPC fallback.",
+            stage_id,
+            exc,
+        )
+        return {}
+
+
+def build_stage_connectors(
+    stage_id: int,
+    connectors_config: dict[str, Any],
+) -> Optional[dict[tuple[str, str], Any]]:
+    """Instantiate OmniConnectors for a stage based on config."""
+    if not connectors_config:
+        return {}
+
+    logger.info(
+        "[Stage-%s] Initializing OmniConnectors with config keys: %s",
+        stage_id,
+        list(connectors_config.keys()),
+    )
+
+    from .factory import OmniConnectorFactory
+    from .config import ConnectorSpec
+
+    connectors: dict[tuple[str, str], Any] = {}
+    try:
+        # 将字典格式的配置转换为ConnectorSpec对象
+        stage_connector_specs = {}
+        for input_key, config in connectors_config.items():
+            if not input_key.startswith("from_stage_"):
+                continue
+
+            from_stage = input_key.replace("from_stage_", "")
+            spec_dict = config.get("spec", {})
+            if not spec_dict:
+                continue
+
+            connector_spec = ConnectorSpec(
+                name=spec_dict.get("name", "InMemoryConnector"),
+                extra=spec_dict.get("extra", {}),
+            )
+            stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec
+
+        # 使用统一的连接器创建逻辑
+        connectors = create_connectors_from_config(stage_connector_specs)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)
+        return None
+
+    return connectors

vllm_omni/entrypoints/omni_llm.py

@@ -1,6 +1,8 @@
 import multiprocessing as mp
 import os
+import sys
 import time
+import uuid
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Optional, Union
@@ -20,6 +22,11 @@ from vllm.utils import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
 
 # Internal imports (our code)
+from vllm_omni.distributed.connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.connectors import (
+    initialize_orchestrator_connectors,
+    get_stage_connector_config,
+)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
@@ -31,9 +38,14 @@ from vllm_omni.entrypoints.log_utils import (
 )
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
+from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
+from vllm_omni.entrypoints.utils import (
+    load_stage_configs_from_model,
+    load_stage_configs_from_yaml,
+    resolve_model_config_path,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -84,14 +96,23 @@ class OmniLLM:
         init_timeout: int = 300,
         **kwargs: Any,
     ):
+        self.worker_backend = kwargs.get("worker_backend", "process")
+        self.ray_address = kwargs.get("ray_address", None)
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
+        self._ray_pg = None
+
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
         if stage_configs_path is None:
-            self.stage_configs = load_stage_configs_from_model(model)
+            self.config_path = resolve_model_config_path(model)
+            self.stage_configs = load_stage_configs_from_yaml(self.config_path)
         else:
+            self.config_path = stage_configs_path
             self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
 
+        # Initialize connectors
+        self.omni_transfer_config,self.connectors = initialize_orchestrator_connectors(self.config_path)
+
         # Optional file handler for orchestrator
         self._log_file = log_file
 
@@ -124,36 +145,140 @@ class OmniLLM:
         self.stage_list = [st for _, st in results]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
-        self._ctx = mp.get_context("spawn")
+        if self.worker_backend == "ray":
+            try:
+                import ray
+                from ray.util.queue import Queue as RayQueue
+            except ImportError:
+                raise ImportError("ray is required for worker_backend='ray'")
+            self._queue_cls = lambda: RayQueue(maxsize=0)
+        else:
+            self._ctx = mp.get_context("spawn")
+            self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
+
         self._stage_in_queues: list[mp.Queue] = []
         self._stage_out_queues: list[mp.Queue] = []
         self._init_sleep_seconds = max(0, int(init_sleep_seconds))
         self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
-        self._start_stage_processes(model)
+        self._start_stages(model)
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
 
-    def _start_stage_processes(self, model: str) -> None:
+    def _start_stages(self, model: str) -> None:
+        if self.worker_backend == "ray":
+            self._start_stage_actors_ray(model)
+        else:
+            self._start_stage_processes_mp(model)
+
+    def _start_stage_processes_mp(self, model: str) -> None:
         for stage_id, stage in enumerate(self.stage_list):
             # Use unbounded queues to avoid deadlock when seeding many requests
-            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            in_q: mp.Queue = self._queue_cls()
-            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            out_q: mp.Queue = self._queue_cls()
             self._stage_in_queues.append(in_q)
             self._stage_out_queues.append(out_q)
 
             # Attach queues and start Stage-owned worker process
             stage.attach_queues(in_q, out_q)
+
+            # Build connectors config for this stage
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
             stage.init_stage_worker(
                 model,
                 log_file=self._log_file,
                 shm_threshold_bytes=self._shm_threshold_bytes,
                 ctx=self._ctx,
                 batch_timeout=self.batch_timeout,
+                connectors_config=stage_connectors_config,
             )
             logger.debug("[Orchestrator] Stage-%s process started", stage_id)
             time.sleep(self._init_sleep_seconds)
 
+    def _start_stage_actors_ray(self, model: str) -> None:
+        import ray
+        import os
+        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+        # Initialize Ray Cluster
+        if not ray.is_initialized():
+             # Create a dummy parallel config just to reuse vLLM's init logic if needed
+             # or just call ray.init()
+             # Pass current PYTHONPATH to workers to ensure they can find vllm_omni
+             runtime_env = {
+                 "env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}
+             }
+             ray.init(address=self.ray_address, ignore_reinit_error=True, runtime_env=runtime_env)
+
+        # Create Placement Group
+        # We allocate 1 bundle per stage. Assume 1 GPU per stage for now.
+        # Must include CPU=1 because default Actor request includes 1 CPU.
+        bundles = [{"GPU": 1.0, "CPU": 1.0} for _ in self.stage_list]
+        self._ray_pg = ray.util.placement_group(bundles, strategy="SPREAD")
+        ray.get(self._ray_pg.ready())
+        logger.info("[Orchestrator] Ray Placement Group created")
+
+        # Define Actor Wrapper
+        @ray.remote(num_gpus=1)
+        class OmniStageRayWorker:
+            def run(self, func, *args, **kwargs):
+                return func(*args, **kwargs)
+
+        from vllm_omni.entrypoints.omni_stage import _stage_worker
+
+        for stage_id, stage in enumerate(self.stage_list):
+            in_q = self._queue_cls()
+            out_q = self._queue_cls()
+            self._stage_in_queues.append(in_q)
+            self._stage_out_queues.append(out_q)
+            stage.attach_queues(in_q, out_q)
+
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
+            # Prepare payload
+            # Reuse logic from OmniStage.init_stage_worker but do not start process
+            engine_args = _to_dict(stage.engine_args)
+            runtime_cfg = _to_dict(getattr(stage.stage_config, "runtime", {}))
+            stage_payload = {
+                "stage_id": stage.stage_id,
+                "engine_args": engine_args,
+                "runtime": runtime_cfg,
+                "shm_threshold_bytes": sys.maxsize, # Disable SHM for Ray mode by setting huge threshold
+                "connectors_config": stage_connectors_config or {},
+            }
+
+            # Start Actor
+            # Inject PYTHONPATH again for the specific actor to be safe
+            worker_actor = OmniStageRayWorker.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=self._ray_pg,
+                    placement_group_bundle_index=stage_id
+                ),
+                runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}}
+            ).remote()
+
+            # Start worker loop asynchronously
+            worker_actor.run.remote(
+                _stage_worker,
+                model=model,
+                stage_payload=stage_payload,
+                in_q=in_q,
+                out_q=out_q,
+                log_file=self._log_file,
+                batch_timeout=self.batch_timeout,
+            )
+            
+            stage._ray_actor = worker_actor
+            logger.debug("[Orchestrator] Stage-%s Ray Actor started", stage_id)
+            time.sleep(self._init_sleep_seconds)
+
     def close(self) -> None:
         """Close all stage processes and clean up resources.
 
@@ -171,10 +296,21 @@ class OmniLLM:
                 )
         for stage in self.stage_list:
             try:
-                stage.stop_stage_worker()
+                if self.worker_backend == "ray" and hasattr(stage, "_ray_actor"):
+                    import ray
+                    # Ray kill
+                    ray.kill(stage._ray_actor)
+                else:
+                    stage.stop_stage_worker()
             except Exception as e:
                 logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
 
+        if self.worker_backend == "ray" and self._ray_pg:
+            try:
+                ray.util.remove_placement_group(self._ray_pg)
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to remove placement group: %s", e)
+
     def __del__(self) -> None:  # best-effort
         try:
             self.close()
@@ -237,11 +373,14 @@ class OmniLLM:
         # Orchestrator keeps stage objects for input derivation
         num_stages = len(self.stage_list)
 
-        # Map from request_id to original prompt
+        # Generate globally unique request IDs and map them to original prompts
-        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+        request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
+        request_id_to_prompt: dict[str, PromptType] = {
+            rid: p for rid, p in zip(request_ids, request_prompts)
+        }
 
         # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[int, float] = {}
+        _req_start_ts: dict[str, float] = {}
         _wall_start_ts: float = time.time()
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
@@ -387,51 +526,90 @@ class OmniLLM:
                         )
                         continue
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
-                    try:
+
-                        # Measure transfer size and time (encode + enqueue)
+                    # Check if we have a connector for this edge
-                        size_bytes = 0
+                    connector_key = (str(stage_id), str(next_stage_id))
+                    connector = self.connectors.get(connector_key)
+                    print(f"connectors: {self.connectors}")
+
+                    sent_via_connector = False
+                    if connector:
+                        sent_via_connector = try_send_via_connector(
+                            connector=connector,
+                            stage_id=stage_id,
+                            next_stage_id=next_stage_id,
+                            req_id=req_id,
+                            next_inputs=next_inputs,
+                            sampling_params=sp_next,
+                            original_prompt=request_id_to_prompt[req_id],
+                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                            metrics=metrics
+                        )
+
+                    if not sent_via_connector:
+                        # Use original queue mechanism
                         try:
-                            size_bytes = len(_set(next_inputs))
+                            if self.worker_backend == "ray":
-                        except Exception:
+                                # Ray mode: Skip SHM fallback, just send object via Queue (Ray handles serialization)
-                            size_bytes = 0
+                                # This prevents SHM access errors across nodes
-                        t0 = time.time()
+                                ipc_payload = {
-                        ipc_payload = _encode(
+                                    "request_id": req_id,
-                            next_inputs,
+                                    "engine_inputs": next_inputs,
-                            getattr(self, "_shm_threshold_bytes", 65536),
+                                    "sampling_params": sp_next,
-                            obj_key="engine_inputs",
+                                    "sent_ts": time.time(),
-                            shm_key="engine_inputs_shm",
+                                }
-                        )
+                                size_bytes = 0 # Approximation or calculate size if needed
-                        ipc_payload.update(
+                                try:
-                            {
+                                    # Approximate size for stats
-                                "request_id": req_id,
+                                    size_bytes = len(_set(next_inputs))
-                                "sampling_params": sp_next,
+                                except:
-                                "sent_ts": time.time(),
+                                    pass
-                            }
+                                t0 = time.time()
-                        )
+                                # Direct submit to Ray Queue
-                        self.stage_list[next_stage_id].submit(ipc_payload)
+                            else:
-                        t1 = time.time()
+                                size_bytes = 0
-                        tx_ms = (t1 - t0) * 1000.0
+                                try:
-                        metrics.on_forward(
+                                    size_bytes = len(_set(next_inputs))
-                            stage_id,
+                                except Exception:
-                            next_stage_id,
+                                    size_bytes = 0
-                            req_id,
+                                t0 = time.time()
-                            int(size_bytes),
+                                ipc_payload = _encode(
-                            float(tx_ms),
+                                    next_inputs,
-                            bool("engine_inputs_shm" in ipc_payload),
+                                    getattr(self, "_shm_threshold_bytes", 65536),
-                        )
+                                    obj_key="engine_inputs",
-                    except Exception as e:
+                                    shm_key="engine_inputs_shm",
-                        logger.warning(
+                                )
-                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                                ipc_payload.update(
-                            req_id,
+                                    {
-                            e,
+                                        "request_id": req_id,
-                        )
+                                        "sampling_params": sp_next,
-                        self.stage_list[next_stage_id].submit(
+                                        "sent_ts": time.time(),
-                            {
+                                    }
-                                "request_id": req_id,
+                                )
-                                "engine_inputs": next_inputs,
+                            self.stage_list[next_stage_id].submit(ipc_payload)
-                                "sampling_params": sp_next,
+                            t1 = time.time()
-                            }
+                            tx_ms = (t1 - t0) * 1000.0
-                        )
+                            metrics.on_forward(
+                                stage_id,
+                                next_stage_id,
+                                req_id,
+                                int(size_bytes),
+                                float(tx_ms),
+                                bool("engine_inputs_shm" in ipc_payload),
+                            )
+                        except Exception as e:
+                            logger.warning(
+                                "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                                req_id,
+                                e,
+                            )
+                            self.stage_list[next_stage_id].submit(
+                                {
+                                    "request_id": req_id,
+                                    "engine_inputs": next_inputs,
+                                    "sampling_params": sp_next,
+                                }
+                            )
                     logger.debug(
                         "[Orchestrator] Forwarded request %s to stage-%s",
                         req_id,

vllm_omni/entrypoints/omni_stage.py

@@ -163,6 +163,7 @@ class OmniStage:
         shm_threshold_bytes: int = 65536,
         ctx: Optional[mp.context.BaseContext] = None,
         batch_timeout: int = 10,
+        connectors_config: Optional[dict] = None,  # New parameter
     ) -> None:
         """Initialize and start the stage worker process.
 
@@ -192,6 +193,7 @@ class OmniStage:
             "engine_args": engine_args,
             "runtime": runtime_cfg,
             "shm_threshold_bytes": self._shm_threshold_bytes,
+            "connectors_config": connectors_config or {},
         }
         if is_async:
             self._proc = ctx.Process(
@@ -340,6 +342,8 @@ def _stage_worker(
         log_stage_batch_stats,
         log_stage_running_avg,
     )
+    from vllm_omni.distributed.connectors.adapter import try_recv_via_connector
+    from vllm_omni.distributed.connectors.utils import build_stage_connectors
     from vllm_omni.entrypoints.omni_llm import OmniStageLLM
     # no inline JSONL/serialization imports; logging handled by utilities
 
@@ -347,6 +351,7 @@ def _stage_worker(
     engine_args = stage_payload.get("engine_args", {})
     runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+    connectors_config = stage_payload.get("connectors_config", {})
 
     # Per-stage logger: clear inherited handlers to avoid broken parent streams
     try:
@@ -440,6 +445,18 @@ def _stage_worker(
     )
     stage_engine = OmniStageLLM(model=model, **engine_args)
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+
+    # Initialize OmniConnectors if configured
+    connectors = {}
+    if connectors_config:
+        built_connectors = build_stage_connectors(
+            stage_id=stage_id,
+            connectors_config=connectors_config,
+        )
+        if built_connectors is None:
+            return
+        connectors = built_connectors
+
     # Signal readiness to orchestrator
     try:
         out_q.put({"type": "stage_ready", "stage_id": stage_id})
@@ -496,9 +513,21 @@ def _stage_worker(
                     _in_flight_ms_by_rid[rid] = 0.0
             except Exception:
                 _in_flight_ms_by_rid[rid] = 0.0
-            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+
-            _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+            # Resolve input data (via Connector or IPC)
-            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+            ein, _rx_metrics = try_recv_via_connector(
+                task=t,
+                connectors=connectors,
+                stage_id=stage_id
+            )
+
+            if ein is None or _rx_metrics is None:
+                ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+
+            if _rx_metrics:
+                _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+                _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+
             batch_request_ids.append(rid)
             if isinstance(ein, list):
                 batch_engine_inputs.extend(ein)

vllm_omni/entrypoints/utils.py

@@ -10,6 +10,41 @@ from vllm_omni.utils import detect_device_type
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 
 
+def resolve_model_config_path(model: str) -> str:
+    """Resolve the stage config file path from the model name.
+
+    Resolves stage configuration path based on the model type and device type.
+    First tries to find a device-specific YAML file from stage_configs/{device_type}/
+    directory. If not found, falls back to the default config file.
+
+    Args:
+        model: Model name or path (used to determine model_type)
+
+    Returns:
+        String path to the stage configuration file
+
+    Raises:
+        FileNotFoundError: If no stage config file exists for the model type
+    """
+    hf_config = get_config(model, trust_remote_code=True)
+    model_type = hf_config.model_type
+    device_type = detect_device_type()
+
+    # Try device-specific config first
+    if device_type != "cuda":
+        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
+        device_config_path = PROJECT_ROOT / device_config_file
+        if os.path.exists(device_config_path):
+            return str(device_config_path)
+
+    # Fall back to default config
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    stage_config_path = PROJECT_ROOT / stage_config_file
+    if not os.path.exists(stage_config_path):
+        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
+    return str(stage_config_path)
+
+
 def load_stage_configs_from_model(model: str) -> list:
     """Load stage configurations from model's default config file.
 
@@ -26,24 +61,8 @@ def load_stage_configs_from_model(model: str) -> list:
     Raises:
         FileNotFoundError: If no stage config file exists for the model type
     """
-    hf_config = get_config(model, trust_remote_code=True)
+    stage_config_path = resolve_model_config_path(model)
-    model_type = hf_config.model_type
+    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path)
-    device_type = detect_device_type()
-
-    # Try device-specific config first
-    if device_type != "cuda":
-        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
-        device_config_path = PROJECT_ROOT / device_config_file
-        if os.path.exists(device_config_path):
-            stage_configs = load_stage_configs_from_yaml(config_path=str(device_config_path))
-            return stage_configs
-
-    # Fall back to default config
-    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
-    stage_config_path = PROJECT_ROOT / stage_config_file
-    if not os.path.exists(stage_config_path):
-        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
-    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path))
     return stage_configs
 
 

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -26,6 +26,9 @@ stage_args:
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    output_connectors:
+      to_stage_1: connector_thinker2talker # thinker -> talker: use local mooncake
   - stage_id: 1
     runtime:
       process: true
@@ -52,17 +55,22 @@ stage_args:
       detokenize: True
       repetition_penalty: 1.05
       stop_token_ids: [8294]
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_0: connector_thinker2talker # talker <- thinker: use local mooncake
+    output_connectors:
+      to_stage_2: connector_talker2generator # talker -> code2wav: use remote mooncake (example)
   - stage_id: 2
     runtime:
       process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
-      gpu_memory_utilization: 0.15
+      gpu_memory_utilization: 0.3
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
@@ -78,6 +86,9 @@ stage_args:
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_1: connector_talker2generator # code2wav <- talker: use remote mooncake
 
 # Top-level runtime config (concise): default windows and stage edges
 runtime:
@@ -85,6 +96,31 @@ runtime:
   defaults:
     window_size: -1             # Simplified: trigger downstream only after full upstream completion
     max_inflight: 1             # Simplified: process serially within each stage
+
+  # Distributed connectors configuration (optional)
+  connectors:
+    # Local mooncake for same-node communication
+    connector_thinker2talker:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 512000000    # 512MB
+        localbuf: 64000000     # 64MB
+        proto: "tcp"
+
+    # Remote mooncake for cross-node communication (example)
+    connector_talker2generator:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 1073741824    # 1GB
+        localbuf: 134217728     # 128MB
+        proto: "tcp"
+
   edges:
     - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
       to: 1