Port changes from feat/multi-request-stream (vs main)

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

docs/mrs_design.md

@@ -0,0 +1,238 @@
+## Multi-Request Streaming (MRS) on a Single Machine — Design Doc
+
+- Author: Team
+- Date: 2025-10-27
+- Version: v0.3 (simplified: copy-based IPC / serial per stage / full-trigger)
+
+### 1. Background & Scope
+- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved.
+- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling.
+- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead.
+
+### 2. Key Constraints
+- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`).
+- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version.
+- Full-trigger across stages: unify `window_size = -1` for all edges — downstream starts only after upstream completes for that request.
+- Serial within a stage: at most one in-flight request per stage (max_inflight=1). No batching/concurrency inside a stage in this version.
+- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0).
+
+### 3. Architecture Overview
+- Processes & IPC queues
+  - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue.
+  - Inter-stage connection via IPC: control plane (Unix Domain Socket / Pipe) + data plane (shared memory ring buffer in future; copy-based in this version).
+  - MPMC queue semantics; control plane passes light-weight pointers/metadata; data plane carries large payloads.
+- Device visibility
+  - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets.
+  - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit.
+- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work.
+- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately.
+- Scheduling & windows
+  - Each edge (stage→stage) defines a window policy; downstream triggers compute by consuming stitched window views.
+  - `-1` mode triggers only after upstream emits `is_last=True` (full aggregation).
+  - Backpressure propagates upstream via the control plane: once high-watermark is reached, upstream blocks or degrades behavior.
+
+### 4. Stage Breakdown
+- AR (autoregressive)
+  - S0 Intake/Plan: parse request, assign request_id, output normalized tasks.
+  - S1 Prefill: tokenize/encode/KV warmup; may micro-batch across requests.
+  - S2 Decode Loop: step-by-step tokens; flush segments by `max_segment_tokens` and `min_flush_interval_ms`.
+  - S3 Postprocess: detokenize, logprobs/struct wrapping (CPU if necessary).
+  - S4 Consumer: local consumer callback/buffering.
+- DiT/Diffusion
+  - D0 Intake/Plan: scheduler/seed init.
+  - D1 Sampling Loop: preview frames every k steps (downsample allowed).
+  - D2 Decode/Encode: VAE decode and JPEG/WEBP encode (prefer GPU→pinned CPU).
+  - D3 Postprocess/Consumer.
+
+### 5. Unified Events & Segment Structure (in-process)
+- Base event: `{ request_id, sequence, event_type, ts, payload, device_handles }`
+- ARSegmentEvent (segment):
+  - Fields: `tokens` (List[int] or GPU tensor view), optional `logprobs`, `is_last`, optional `hidden_states/aux` refs.
+  - Semantics: strictly ordered by `sequence` for the same `request_id`; `is_last` marks the final segment.
+- ImageProgressEvent: `{ step, image_tensor|jpeg_ref, is_last }` (diffusion preview/completion).
+- Completed/Error: terminal or error events affect only the current request.
+
+### 6. AR Segmentation & Windowing (simplified)
+- Production (upstream stage)
+  - Accumulate tokens during decoding; flush a segment if any of the following holds:
+    - Reaches `max_segment_tokens` length
+    - Exceeds `min_flush_interval_ms` since last flush
+    - Request ended (`is_last=True`)
+  - Flush immediately via `output_queue.put(segment)`; do not wait for entire request completion.
+- Consumption (downstream stage)
+  - Use unified `window_size = -1`: trigger only after receiving `is_last=True`; no window slicing.
+- Backpressure & watermark
+  - Must-deliver events (token.delta, Completed) are not dropped; block upstream briefly on high watermark.
+  - For long sequences under `-1`, allow early segments to spill from GPU to pinned CPU to control VRAM pressure.
+
+#### Example config (YAML, simplified)
+```yaml
+mrs:
+  enabled: true
+  defaults:
+    window_size: -1           # Full-trigger across stages
+    max_inflight: 1           # Serial per stage
+```
+
+### 7. IPC Implementation (simplified: copy-based)
+- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data).
+- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version.
+- Backpressure: blocking reads/writes based on queue capacity, blocking upstream on high watermark.
+
+### 8. Scheduling & Cancellation (simplified)
+- Serial FIFO: each stage `max_inflight=1`; no cross-request micro-batching or concurrency.
+- Cancel/timeout: ongoing work should exit early; remaining queued requests stay pending.
+- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency.
+
+#### Short sequence example (req0/req1, stage A→B)
+1) t0: stage A handles req0
+2) t1: req0 completes on A → enters B; A immediately starts req1
+3) t2: B handles req0 while A handles req1 (parallel across stages)
+
+### 9. Integration Points (by file)
+- `vllm_omni/entrypoints/omni_llm.py`
+  - Introduce a ProcessOrchestrator (replace threaded version):
+    - Spawn stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), create control/data channels, build window state machines.
+    - Add `enable_mrs` in `generate`; when enabled, do not collect `stage.engine.generate(...)` into lists; stream segments to downstream via IPC instead.
+  - Register `window_size` per edge; `-1` triggers downstream once `is_last` is observed.
+  - Manage process lifecycle: start/heartbeat/graceful shutdown/restart on crash; propagate cancel/timeout.
+- `vllm_omni/outputs.py`
+  - Add in-process events: `ARSegmentEvent`, `ImageProgressEvent`, `CompletedEvent`, `ErrorEvent`.
+  - Keep `OmniRequestOutput` as the external representation of final stage outputs.
+- `vllm_omni/worker/gpu_ar_model_runner.py`
+  - Near `valid_sampled_token_ids` parsing (after updating `req_state.output_token_ids`), aggregate new tokens into segments and emit; send final `is_last=True` and Completed.
+  - Replace current `.to("cpu").contiguous()` sync path: by default keep GPU refs and let downstream copy asynchronously from a pinned pool when needed.
+- `vllm_omni/worker/gpu_diffusion_model_runner.py`
+
+### 9.1 Process Device Visibility
+- Device binding: set `CUDA_VISIBLE_DEVICES` before process start, or call `torch.cuda.set_device` early in init; configs can be single/multi-GPU.
+- Cross-device transfer (simplified): via CPU copies; zero-copy paths are out of scope for now.
+
+### 10. Worker Template & Window State (pseudo-code)
+```python
+def stage_worker(input_q, output_q, ctx):
+    while ctx.alive:
+        item = input_q.get()  # blocking / with timeout
+        if item.is_cancelled():
+            continue
+        with cuda_stream(ctx.stream):
+            out_segments = run_stage_compute(item)  # produce segments
+        for seg in out_segments:
+            output_q.put(seg)
+
+class WindowState:
+    def __init__(self, window_size):
+        self.window_size = window_size
+        self.last_emitted = 0
+        self.buffer = []  # references/paged views
+    def on_segment(self, seg):
+        self.buffer.append(seg)
+        if self.window_size == -1:
+            return self.emit_full_if_last(seg)
+        return self.emit_windows()
+    def emit_windows(self):
+        emits = []
+        while self._pending_len() - self.last_emitted >= self.window_size:
+            emits.append(self._make_view(self.last_emitted, self.window_size))
+            self.last_emitted += self.window_size
+        return emits
+    def emit_full_if_last(self, seg):
+        return [self._make_full_view()] if seg.is_last else []
+```
+
+### 11. Configuration
+- `mrs.enabled`, `mrs.ar.default_window_size`, `mrs.ar.min_flush_interval_ms`, `mrs.ar.max_segment_tokens`, `mrs.ar.scheduler_policy`.
+- Devices/processes:
+  - `mrs.stages`: [{ name, process: true, devices: "0" | "1,2", extra_env: {} }]
+  - `mrs.edges`: [{ from, to, window_size }]
+- Pinned/SHM pool sizing: estimate by concurrency and max sequence; tune watermarks and GC policies; configure CUDA IPC limits and cleanup thresholds if enabled.
+
+### 12. Observability
+- Metrics:
+  - Stage-level: enqueue/dequeue rates, queue depth, blocking time on watermarks, window triggers/counts.
+  - Request-level: TTFT, TPST, completion latency, cancel rate.
+  - Resource-level: GPU utilization, VRAM usage/fragmentation, pinned usage.
+- Logs/tracing: carry `request_id`, stage name, and `sequence`; sample decisions for downsampling/backpressure.
+
+### 13. Test Plan
+- Unit: queue correctness, window state machine, cancel/timeout, memory pool rent/return.
+- Performance: AR QPS/TPST vs. window/segment sizes; diffusion preview frequency vs. throughput.
+- Stability: long-run memory leaks; backpressure trigger/recovery paths.
+- Compatibility: behavior parity when MRS is disabled.
+
+### 14. Milestones
+- M1: Process Orchestrator with simple IPC; serial per stage; full-trigger across stages (window=-1).
+- M2: Hardening and observability; cancel/timeout and error paths.
+- M3: Optional zero-copy and concurrency/micro-batching optimizations.
+
+### 15. Risks & Mitigations
+— End-to-end latency: full-trigger sacrifices interactivity; start simple/correct, add windowed streaming later.
+— CPU copy overhead: copy-based IPC for maintainability; add SHM/zero-copy later as an optimization path.
+— Serial throughput: `max_inflight=1` limits per-stage throughput; enable batching/concurrency later as needed.
+
+### 16. Refactor: Sink LLM init and process into Stage (diagram)
+
+Objective: without changing external call patterns, encapsulate per-stage LLM init, process creation, and worker logic inside `Stage`, so `PipelinedOmniLLM` focuses on orchestration (seeding, polling, forwarding, collecting) while keeping multi-process, shared memory, and device mapping capabilities.
+
+Key changes (vs. current):
+- Before: `PipelinedOmniLLM` created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there.
+- After: enhance `Stage` to start its own subprocess via `start_process()`, with `worker_main()` as the entry; device setup, LLM init, batching, and SHM live in `Stage`.
+
+Classes & responsibilities
+- PipelinedOmniLLM (orchestrator)
+  - Build `Stage` list (preserve `OmniStage` init timing and `process_engine_inputs` flow)
+  - Connect adjacent `Stage` input/output queues; seed requests to stage 0
+  - Poll stage outputs, decode results, call `process_engine_inputs`, then encode and forward to next stage
+  - Termination & cleanup: distribute shutdown signals; join/terminate subprocesses
+- Stage (stage unit)
+  - Members: `stage_config` (with `engine_args` and `mrs`), `in_q/out_q`, subprocess handle, stats
+  - API:
+    - `start_process(in_q, out_q)`: spawn subprocess
+    - `stop_process()`: graceful exit
+    - `submit(ipc_payload)`: submit to `in_q` (may use shared memory)
+    - `try_collect() -> Optional[payload]`: non-blocking get from `out_q`
+    - `process_engine_inputs(stages, prompts)`: reuse/delegate existing logic
+  - Subprocess entry `worker_main()`:
+    - set_stage_gpu_devices (`CUDA_VISIBLE_DEVICES` & logical index mapping)
+    - build `OmniStageLLM(model, **engine_args)`
+    - while-loop:
+      - take first task + up to `mrs.max_batch_size-1` non-blocking; support `max_batch_size>=1`
+      - `maybe_load_from_ipc` per `engine_inputs` in batch
+      - `stage_engine.generate(batched_inputs, sampling_params)` (window `-1`)
+      - write per-request aggregated outputs to `out_q` after `maybe_dump_to_shm`
+      - log start/finish; report `error` per request on exceptions
+
+Data flow (sketch)
+```
+PipelinedOmniLLM
+  ├─ Stage[0].start_process(in0, out0)
+  ├─ Stage[1].start_process(in1, out1)
+  ├─ ...
+  └─ Orchestrator loop:
+       - stage0.submit(encode(engine_inputs))
+       - for each stage:
+           res = stage.try_collect()
+           outputs = decode(res)
+           stage.set_engine_outputs(outputs)
+           if !final: next_inputs = next_stage.process_engine_inputs(stages, original_prompt)
+                    next_stage.submit(encode(next_inputs))
+```
+
+Config (preserve & extend)
+- `mrs.devices`:
+  - Comma-separated visibility list (e.g., "2,5,7"): set visible set; default logical index 0 (first GPU)
+  - Number/string: treat as logical index (0/1-based per implementation) mapping to physical GPUs
+- `mrs.max_batch_size`: max requests taken per intake (default 1)
+- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; can be injected into stages)
+
+Kept & strengthened
+- Multi-process: each stage in its own process; parallel across stages; serial/batching inside a stage controlled by `max_batch_size`
+- Shared memory: large objects via SHM with strict `memoryview -> close -> unlink`
+- Device mapping: unified `set_stage_gpu_devices`, support comma lists and logical indices
+- Windows: keep `window_size=-1` semantics (downstream triggers only after upstream finishes the request)
+
+Migration (execution order)
+1) Implement enhanced `Stage` (e.g., `entrypoints/stage.py`) with `start_process/stop_process/submit/try_collect/worker_main`.
+2) Make `PipelinedOmniLLM` hold `Stage` instances, replacing direct queues and `_stage_worker` with `Stage` APIs.
+3) Reuse `pipeline_utils`: device/SHM/encode-decode utilities stay in utils.
+4) Regression: single/multi-request, various `max_batch_size`, SHM threshold, device mapping, multi-stage parallelism, error paths, and shutdown flows.

examples/offline_inference/qwen_2_5_omni/extract_prompts.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+import argparse
+from typing import Optional
+import os
+import torch
+
+
+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="输入 .lst 文件路径")
+    parser.add_argument("--output", "-o", required=True, help="输出文件路径")
+    parser.add_argument("--topk", "-k", type=int, default=100, help="提取前 K 个 prompt（默认 100）")
+    parser.add_argument("--pt-output", "-p", type=str, default=None, help="可选：输出 .pt 文件路径（默认与 --output 同名 .pt）")
+    args = parser.parse_args()
+
+    prompts = []
+    with open(args.input, "r", 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")
+
+    # 另存为 PyTorch 序列化的 .pt，内容为字符串列表
+    pt_path = args.pt_output if args.pt_output else os.path.splitext(args.output)[0] + ".pt"
+    torch.save(prompts, pt_path)
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+Batch text generation with Qwen 2.5 Omni using Transformers without chat templates.
+
+- Processes 10 fixed prompts in batches of 3
+- Does NOT use tokenizer.apply_chat_template (raw prompts only)
+- Prints only generated continuations for each prompt
+"""
+
+import os
+import argparse
+from typing import Iterable, List
+import time
+import torch
+
+import random
+import numpy as np
+import torch
+from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
+from tqdm import tqdm
+import json
+import soundfile as sf
+
+PROMPTS: List[str] = [
+    "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.",
+    "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.",
+    "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.",
+    "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.",
+    "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.",
+    "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.",
+    "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.",
+    "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.",
+    "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.",
+    "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.",
+]
+def make_text_prompt(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."}
+        ],
+    },
+    {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": prompt},
+        ],
+    }]
+
+def chunked(items: List[str], size: int) -> Iterable[List[str]]:
+    for i in range(0, len(items), size):
+        yield items[i : i + size]
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Qwen 2.5 Omni batch generation without chat templates",
+    )
+    parser.add_argument(
+        "--model-id",
+        type=str,
+        default=os.getenv("QWEN_OMNI_MODEL_ID", "Qwen/Qwen2.5-Omni"),
+        help="HF model repo id (e.g., Qwen/Qwen2.5-Omni)",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=2048,
+        help="Maximum number of new tokens to generate",
+    )
+    parser.add_argument(
+        "--temperature", type=float, default=0.0, help="Sampling temperature"
+    )
+    parser.add_argument("--top-p", type=float, default=1.0, help="Top-p sampling")
+    parser.add_argument(
+        "--repetition-penalty",
+        type=float,
+        default=1.1,
+        help="Repetition penalty (matches pipelined_end2end defaults)",
+    )
+    parser.add_argument(
+        "--stop-token-ids",
+        type=str,
+        default="8294",
+        help="Comma-separated stop token ids (added to eos_token_id)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for reproducibility (sampling)",
+    )
+    parser.add_argument(
+        "--pt-prompts",
+        type=str,
+        default=None,
+        help="Path to .pt file containing List[str] prompts (overrides built-in PROMPTS)",
+    )
+    parser.add_argument(
+        "--tokenize",
+        action="store_true",
+        help="If set, utils.make_text_prompt will return token ids; keep unset for raw text",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="outputs",
+        help="Output directory to save aligned comparison files (same folder as pipelined_end2end)",
+    )
+    args = parser.parse_args()
+
+    if torch.cuda.is_available():
+        torch_dtype = torch.bfloat16
+    else:
+        torch_dtype = torch.float32
+
+    # Match pipelined_end2end.py deterministic settings
+    if args.seed is not None:
+        os.environ["PYTHONHASHSEED"] = str(args.seed)
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+        try:
+            random.seed(args.seed)
+            np.random.seed(args.seed)
+            torch.manual_seed(args.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(args.seed)
+                torch.cuda.manual_seed_all(args.seed)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+        except Exception:
+            pass
+
+    print(f"[Info] Loading model: {args.model_id}")
+    processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id)
+    model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+        args.model_id,
+        trust_remote_code=True,
+        torch_dtype=torch_dtype,
+        device_map="auto",
+        attn_implementation="flash_attention_2",
+    )
+
+    # Load prompts: prefer .pt if provided
+    if args.pt_prompts is not None:
+        loaded = torch.load(args.pt_prompts)
+        if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)):
+            raise ValueError("--pt-prompts must point to a .pt containing List[str]")
+        prompts: List[str] = loaded
+        print(f"[Info] Loaded {len(prompts)} prompts from {args.pt_prompts}")
+    else:
+        prompts = PROMPTS
+
+    # Ensure utils.make_text_prompt can find expected fields
+    # utils expects args.model and args.tokenize
+    setattr(args, "model", args.model_id)
+    if not hasattr(args, "tokenize"):
+        setattr(args, "tokenize", False)
+
+    print(f"[Info] Running single-sample generation (batch size fixed to 1) ...\n")
+    global_index = 0
+    total_requests = len(prompts)
+    sum_request_time_ms = 0.0
+    wall_t0 = time.time()
+    stats_path = os.path.join(args.output_dir, "hf_transformers.stats.jsonl")
+
+    for rid, prompt in tqdm(list(enumerate(prompts)), total=len(prompts), desc="Generating"):
+        _t0 = time.time()
+        # Build single-sample inputs
+        tp = make_text_prompt(prompt)
+        tp = processor.apply_chat_template(tp, add_generation_prompt=True, tokenize=False)
+        inputs = processor(text=tp, return_tensors="pt", padding=True, use_audio_in_video=True)
+        inputs = inputs.to(model.device).to(model.dtype)
+
+        with torch.no_grad():
+            text_ids, audio = model.generate(
+                **inputs,
+                use_audio_in_video=True,
+            )
+
+        # Decode
+        try:
+            texts = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)  # type: ignore[arg-type]
+        except Exception:
+            try:
+                texts = processor.tokenizer.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)  # type: ignore[attr-defined]
+            except Exception:
+                texts = [str(x) for x in text_ids]
+
+        text_str = texts[0] if len(texts) > 0 else ""
+        os.makedirs(args.output_dir, exist_ok=True)
+        out_txt = os.path.join(args.output_dir, f"{rid:05d}.txt")
+
+        content = None
+        if os.path.exists(out_txt):
+            try:
+                with open(out_txt, "r", encoding="utf-8") as f:
+                    content = f.read()
+            except Exception:
+                content = None
+        if content and "transformers_text_output:" in content:
+            try:
+                parts = content.split("transformers_text_output:\n", 1)
+                new_content = parts[0] + "transformers_text_output:\n" + text_str.strip() + "\n"
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.write(new_content)
+            except Exception:
+                pass
+        else:
+            lines = [
+                "Prompt:\n",
+                str(prompt) + "\n",
+                "vllm_text_output:\n\n",
+                "transformers_text_output:\n",
+                text_str.strip() + "\n",
+            ]
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception:
+                pass
+
+        # Save wav per request
+        try:
+            if audio is not None:
+                if isinstance(audio, (list, tuple)):
+                    wav = audio[0]
+                else:
+                    # Tensor or array: take first sample if batched, else as-is
+                    if hasattr(audio, "ndim") and hasattr(audio, "shape") and len(audio.shape) > 1:
+                        wav = audio[0]
+                    else:
+                        wav = audio
+                wav_path = os.path.join(args.output_dir, f"output_{rid}.wav")
+                sf.write(wav_path, torch.as_tensor(wav).reshape(-1).detach().cpu().numpy(), samplerate=24000)
+        except Exception:
+            pass
+
+        # Stats per request
+        req_ms = (time.time() - _t0) * 1000.0
+        sum_request_time_ms += req_ms
+        try:
+            rec = {
+                "type": "transformers_request_stats",
+                "request_id": rid,
+                "batch_size": 1,
+                "per_request_time_ms": req_ms,
+                "num_tokens_out": int(text_ids.shape[-1]) if hasattr(text_ids, "shape") else None,
+            }
+            with open(stats_path, "a", encoding="utf-8") as f:
+                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+        except Exception:
+            pass
+
+        global_index += 1
+
+    print("[Done] Generated", global_index, "responses.")
+    wall_ms = (time.time() - wall_t0) * 1000.0
+    avg_e2e_ms = (sum_request_time_ms / total_requests) if total_requests > 0 else 0.0
+    print(f"[E2E] total_requests={total_requests}, wall_time_ms={wall_ms:.2f}, avg_e2e_per_request_ms={avg_e2e_ms:.2f}")
+    try:
+        summary = {
+            "type": "transformers_summary",
+            "total_requests": total_requests,
+            "wall_time_ms": wall_ms,
+            "sum_request_time_ms": sum_request_time_ms,
+            "avg_e2e_per_request_ms": avg_e2e_ms,
+        }
+        with open(stats_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(summary, ensure_ascii=False) + "\n")
+    except Exception:
+        pass
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/offline_inference/qwen_2_5_omni/pipelined_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.pipelined_omni_llm import PipelinedOmniLLM
+
+_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(
+        "--pt-prompts", type=str, default=None, help="Path to .pt file containing List[str] prompts (overrides --prompts when prompt_type=text)."
+    )
+    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=True,
+        help="Enable writing detailed statistics (default: enabled)",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    # If pt-prompts provided and prompt_type is text, load and override args.prompts
+    try:
+        if getattr(args, "pt_prompts", None) and args.prompt_type == "text":
+            loaded = torch.load(args.pt_prompts)
+            if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)):
+                raise ValueError("--pt-prompts must point to a .pt containing List[str]")
+            args.prompts = loaded
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.pt_prompts}")
+    except Exception as e:
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts <file.pt> (with --prompt_type text)")
+    omni_llm = PipelinedOmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file="omni_llm_pipeline.log",
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
+    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,
+    ]
+
+    prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+    omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+
+    # 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")
+                lines.append("transformers_text_output:\n\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception:
+                    pass
+                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
+                )
+                # Save aligned text file with empty text (no text output in this stage)
+                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\n")
+                lines.append("transformers_text_output:\n\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception:
+                    pass
+                print(f"Request ID: {request_id}, Saved audio to {output_wav} and text to {out_txt}")
+
+
+if __name__ == "__main__":
+    main()

vllm_omni/entrypoints/omni_stage.py

@@ -1,9 +1,22 @@
 """
 Stage manager for orchestrating multiple engines in vLLM-omni.
+
+Enhanced to encapsulate per-stage process lifecycle and worker logic
+(device setup, LLM init, batching, shared-memory IPC), while preserving
+the original input processing utilities for cross-stage data wiring.
 """
 
 import importlib
-from typing import List, Union
+import logging
+import multiprocessing as mp
+from typing import Any, Dict, List, Optional, Union
+
+from vllm_omni.entrypoints.pipeline_utils import (
+    _to_dict,
+    set_stage_gpu_devices,
+    maybe_load_from_ipc,
+    maybe_dump_to_shm,
+)
 
 from vllm.inputs import TextPrompt
 from vllm.v1.engine import EngineCoreOutput
@@ -46,6 +59,14 @@ class OmniStage:
         else:
             self.final_output_type = None
 
+        # Runtime orchestration state (added)
+        self._in_q: Optional[mp.Queue] = None
+        self._out_q: Optional[mp.Queue] = None
+        self._proc: Optional[mp.Process] = None
+        self._log_file: Optional[str] = None
+        self._shm_threshold_bytes: int = 65536
+        self._logger = logging.getLogger(__name__)
+
     def set_engine(self, engine: LLMEngine) -> None:
         """Initialize the engine for the stage."""
         self.engine = engine
@@ -58,6 +79,67 @@ class OmniStage:
         """Set the engine output for the stage."""
         self.engine_outputs = engine_outputs
 
+    # ----------------- New Orchestration APIs -----------------
+    def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None:
+        self._in_q = in_q
+        self._out_q = out_q
+
+    def init_stage_worker(
+        self,
+        model: str,
+        *,
+        log_file: Optional[str] = None,
+        shm_threshold_bytes: int = 65536,
+        ctx: Optional[mp.context.BaseContext] = None,
+        batch_timeout: int = 10,
+    ) -> None:
+        assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
+        self._log_file = log_file
+        self._shm_threshold_bytes = shm_threshold_bytes
+        ctx = ctx or mp.get_context("spawn")
+        # Prepare lightweight dict config for worker
+        engine_args = _to_dict(self.engine_args)
+        mrs_cfg = _to_dict(getattr(self.stage_config, "mrs", {}))
+        stage_payload: Dict[str, Any] = {
+            "stage_id": self.stage_id,
+            "engine_args": engine_args,
+            "mrs": mrs_cfg,
+            "shm_threshold_bytes": self._shm_threshold_bytes,
+        }
+        self._proc = ctx.Process(
+            target=_stage_worker,
+            args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout),
+        )
+        self._proc.start()
+
+    def stop_stage_worker(self) -> None:
+        if self._in_q is not None:
+            try:
+                self._in_q.put_nowait(None)
+            except Exception:
+                pass
+        if self._proc is not None:
+            try:
+                self._proc.join(timeout=5)
+            except Exception:
+                pass
+            if self._proc.is_alive():
+                try:
+                    self._proc.terminate()
+                except Exception:
+                    pass
+
+    def submit(self, payload: Dict[str, Any]) -> None:
+        assert self._in_q is not None
+        self._in_q.put(payload)
+
+    def try_collect(self) -> Optional[Dict[str, Any]]:
+        assert self._out_q is not None
+        try:
+            return self._out_q.get_nowait()
+        except Exception:
+            return None
+
     def process_engine_inputs(
         self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
     ) -> List[Union[OmniTokensPrompt, TextPrompt]]:
@@ -90,3 +172,253 @@ class OmniStage:
             return self.custom_process_input_func(
                 stage_list, engine_input_source, prompt
             )
+
+
+def _stage_worker(
+    model: str,
+    stage_payload: Dict[str, Any],
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    log_file: Optional[str] = None,
+    batch_timeout: int = 10,
+) -> None:
+    """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
+    import logging as _logging
+    from vllm_omni.entrypoints.omni_llm import OmniStageLLM  # noqa: WPS433
+    import queue as _queue
+    import time as _time
+    from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl, serialize_obj as _ser
+
+    stage_id = stage_payload["stage_id"]
+    engine_args = stage_payload.get("engine_args", {})
+    mrs_cfg = stage_payload.get("mrs", {})
+    shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+
+    # Per-stage file logger (optional)
+    try:
+        if log_file:
+            stage_log = _logging.getLogger(__name__)
+            stage_log.setLevel(_logging.DEBUG)
+            fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
+            fh.setLevel(_logging.DEBUG)
+            fh.setFormatter(_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s"))
+            class _StageFilter(_logging.Filter):
+                def filter(self, record: _logging.LogRecord) -> bool:
+                    setattr(record, "stage", stage_id)
+                    return True
+            fh.addFilter(_StageFilter())
+            stage_log.addHandler(fh)
+    except Exception:
+        pass
+
+    # Stage stats JSONL file
+    _stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None
+
+    # Aggregates for running average
+    _agg_total_tokens = 0
+    _agg_total_gen_time_ms = 0.0
+
+    # Device mapping
+    try:
+        set_stage_gpu_devices(stage_id, mrs_cfg.get("devices"))
+    except Exception as e:
+        _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
+
+    # Init LLM
+    _logging.getLogger(__name__).debug("[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys()))
+    stage_engine = OmniStageLLM(model=model, **engine_args)
+    _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+    # Signal readiness to orchestrator
+    try:
+        out_q.put({"type": "stage_ready", "stage_id": stage_id})
+    except Exception:
+        pass
+
+    # Batch processing loop
+    while True:
+        task = in_q.get()
+        _recv_dequeue_ts = _time.time()
+        if task is None:
+            _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
+            break
+
+        max_batch_size = int(mrs_cfg.get("max_batch_size", 1) or 1)
+        batch_tasks: List[Dict[str, Any]] = [task]
+        if max_batch_size > 1:
+            while len(batch_tasks) < max_batch_size:
+                if not in_q.empty():
+                    extra = in_q.get(timeout=batch_timeout)
+                    if extra is None:
+                        in_q.put(None)
+                        break
+                    batch_tasks.append(extra)
+                else:
+                    break
+
+        batch_request_ids: List[Any] = []
+        batch_engine_inputs: List[Any] = []
+        _rx_bytes_by_rid: Dict[Any, int] = {}
+        _rx_decode_ms_by_rid: Dict[Any, float] = {}
+        _in_flight_ms_by_rid: Dict[Any, float] = {}
+        for t in batch_tasks:
+            rid = t["request_id"]
+            try:
+                sent_ts = float(t.get("sent_ts", None)) if isinstance(t, dict) else None
+                if sent_ts is not None:
+                    _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0
+                else:
+                    _in_flight_ms_by_rid[rid] = 0.0
+            except Exception:
+                _in_flight_ms_by_rid[rid] = 0.0
+            _t0 = _time.time()
+            ein = maybe_load_from_ipc(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+            _t1 = _time.time()
+            _rx_decode_ms_by_rid[rid] = (_t1 - _t0) * 1000.0
+            try:
+                if "engine_inputs_shm" in t:
+                    _rx_bytes_by_rid[rid] = int(t["engine_inputs_shm"]["size"])  # type: ignore[index]
+                else:
+                    _rx_bytes_by_rid[rid] = len(_ser(ein))
+            except Exception:
+                _rx_bytes_by_rid[rid] = 0
+            batch_request_ids.append(rid)
+            if isinstance(ein, list):
+                batch_engine_inputs.extend(ein)
+            elif isinstance(ein, dict):
+                batch_engine_inputs.append(ein)
+            else:
+                _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
+        sampling_params = batch_tasks[0]["sampling_params"]
+        _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=%d, request_ids=%s", stage_id, len(batch_tasks), batch_request_ids)
+        print("--------------------------------", flush=True)
+        print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True)
+        print("--------------------------------", flush=True)
+        try:
+            gen_outputs: List[Any] = []
+            _gen_t0 = _time.time()
+            for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False):
+                gen_outputs.append(ro)
+            _gen_t1 = _time.time()
+            _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
+
+            # Group outputs per request id with fallback
+            req_to_outputs: Dict[Any, List[Any]] = {rid: [] for rid in batch_request_ids}
+            unmapped: List[Any] = []
+            for ro in gen_outputs:
+                rid = getattr(ro, "request_id", None)
+                if rid in req_to_outputs:
+                    req_to_outputs[rid].append(ro)
+                else:
+                    unmapped.append(ro)
+            if unmapped:
+                idx = 0
+                for ro in unmapped:
+                    target_rid = batch_request_ids[idx % len(batch_request_ids)]
+                    req_to_outputs[target_rid].append(ro)
+                    idx += 1
+
+            # Per-request stats logging and aggregates
+            def _count_tokens(_ros: List[Any]) -> int:
+                total = 0
+                for _ro in _ros:
+                    try:
+                        outs = getattr(_ro, "outputs", None)
+                        if outs and len(outs) > 0:
+                            tokens = getattr(outs[0], "token_ids", None)
+                            if tokens is not None:
+                                total += len(tokens)
+                    except Exception:
+                        pass
+                return total
+
+            for rid in batch_request_ids:
+                _r_outputs = req_to_outputs.get(rid, [])
+                _num_tokens = _count_tokens(_r_outputs)
+                _agg_total_tokens += _num_tokens
+                _agg_total_gen_time_ms += _gen_ms
+                _tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0
+                if _stats_file:
+                    try:
+                        _append_jsonl(_stats_file, {
+                            "type": "stage_request_stats",
+                            "stage_id": stage_id,
+                            "request_id": rid,
+                            "batch_size": len(batch_tasks),
+                            "num_tokens_out": int(_num_tokens),
+                            "stage_gen_time_ms": _gen_ms,
+                            "tokens_per_s": _tokens_per_s,
+                            "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
+                            "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
+                            "rx_mbps": (float(_rx_bytes_by_rid.get(rid, 0)) * 8.0) / (max(float(_rx_decode_ms_by_rid.get(rid, 0.0)), 1e-6) * 1000.0),
+                        })
+                    except Exception:
+                        pass
+
+            if _stats_file:
+                try:
+                    _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
+                    _append_jsonl(_stats_file, {
+                        "type": "stage_running_avg",
+                        "stage_id": stage_id,
+                        "total_tokens": int(_agg_total_tokens),
+                        "total_gen_time_ms": _agg_total_gen_time_ms,
+                        "avg_tokens_per_s": _avg_tokens_per_s,
+                    })
+                    _append_jsonl(_stats_file, {
+                        "type": "stage_batch_stats",
+                        "stage_id": stage_id,
+                        "batch_size": len(batch_tasks),
+                        "batch_gen_time_ms": _gen_ms,
+                        "request_ids": list(batch_request_ids),
+                    })
+                except Exception:
+                    pass
+
+            # Emit per-request results
+            for rid in batch_request_ids:
+                r_outputs = req_to_outputs.get(rid, [])
+                try:
+                    use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
+                    _metrics = {
+                        "num_tokens_out": int(_count_tokens(r_outputs)),
+                        "stage_gen_time_ms": _gen_ms,
+                        "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
+                        "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
+                        "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
+                    }
+                    if use_shm:
+                        out_q.put({
+                            "request_id": rid,
+                            "stage_id": stage_id,
+                            "engine_outputs_shm": payload,
+                            "metrics": _metrics,
+                        })
+                    else:
+                        out_q.put({
+                            "request_id": rid,
+                            "stage_id": stage_id,
+                            "engine_outputs": payload,
+                            "metrics": _metrics,
+                        })
+                except Exception:
+                    out_q.put({
+                        "request_id": rid,
+                        "stage_id": stage_id,
+                        "engine_outputs": r_outputs,
+                        "metrics": {
+                            "num_tokens_out": int(_count_tokens(r_outputs)),
+                            "stage_gen_time_ms": _gen_ms,
+                            "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
+                            "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
+                            "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
+                        },
+                    })
+                _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
+        except Exception as e:
+            _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
+            for rid in batch_request_ids:
+                out_q.put({
+                    "request_id": rid,
+                    "stage_id": stage_id,
+                    "error": str(e),
+                })

vllm_omni/entrypoints/pipeline_utils.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple, Union
+
+import logging
+import json
+import os
+import pickle
+from multiprocessing import shared_memory as _shm
+
+import cloudpickle
+
+logger = logging.getLogger(__name__)
+
+
+def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> None:
+    """Configure per-stage CUDA visibility and current device.
+
+    Behavior
+    - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly
+      to this list; logical index 0 is used as current device.
+    - Integer or digit-string: treat as logical index (0-based) into the current
+      CUDA_VISIBLE_DEVICES mapping; map to the physical device, and then set
+      CUDA_VISIBLE_DEVICES to this single device.
+    - None/"cpu": keep default visibility.
+    - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
+    """
+    try:
+        selected_physical: Optional[int] = None
+        logical_idx: Optional[int] = None
+
+        if isinstance(devices, str) and "," in devices:
+            os.environ["CUDA_VISIBLE_DEVICES"] = devices
+            toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
+            if toks:
+                try:
+                    selected_physical = int(toks[0])
+                    logger.debug(
+                        "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
+                        stage_id, devices, selected_physical,
+                    )
+                except Exception:
+                    selected_physical = None
+        elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
+            logical_idx = max(0, int(devices))
+            vis = os.environ.get("CUDA_VISIBLE_DEVICES")
+            if vis:
+                try:
+                    mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
+                    if 0 <= logical_idx < len(mapping):
+                        selected_physical = mapping[logical_idx]
+                except Exception:
+                    selected_physical = None
+            if selected_physical is None:
+                selected_physical = int(logical_idx)
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
+            logger.debug(
+                "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device",
+                stage_id, logical_idx + 1, selected_physical,
+            )
+        elif devices in (None, "cpu"):
+            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
+        else:
+            selected_physical = int(str(devices))
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
+            logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical)
+
+        try:
+            import torch  # noqa: WPS433
+            if torch.cuda.is_available():
+                try:
+                    torch.cuda.set_device(0)
+                except Exception:
+                    pass
+                num = torch.cuda.device_count()
+                info = []
+                for i in range(num):
+                    total = torch.cuda.get_device_properties(i).total_memory
+                    free, _ = torch.cuda.mem_get_info(i)
+                    info.append({
+                        "idx": i,
+                        "name": torch.cuda.get_device_name(i),
+                        "total": int(total),
+                        "free": int(free),
+                    })
+                logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
+        except Exception:
+            pass
+    except Exception as e:
+        logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
+
+
+def serialize_obj(obj: Any) -> bytes:
+    """Serialize a Python object to bytes using cloudpickle."""
+    return cloudpickle.dumps(obj)
+
+
+def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
+    """Write bytes into SharedMemory and return meta dict {name,size}.
+
+    Caller should close the segment; the receiver should unlink.
+    """
+    shm = _shm.SharedMemory(create=True, size=len(payload))
+    mv = memoryview(shm.buf)
+    mv[: len(payload)] = payload
+    del mv
+    meta = {"name": shm.name, "size": len(payload)}
+    try:
+        shm.close()
+    except Exception:
+        pass
+    return meta
+
+
+def shm_read_bytes(meta: Dict[str, Any]) -> bytes:
+    """Read bytes from SharedMemory by meta {name,size} and cleanup."""
+    shm = _shm.SharedMemory(name=meta["name"])  # type: ignore[index]
+    mv = memoryview(shm.buf)
+    data = bytes(mv[: meta["size"]])
+    del mv
+    try:
+        shm.close()
+    except Exception:
+        pass
+    try:
+        shm.unlink()
+    except Exception:
+        pass
+    return data
+
+
+def _ensure_parent_dir(path: str) -> None:
+    """Ensure the parent directory for a file path exists (best-effort)."""
+    try:
+        parent = os.path.dirname(path)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+    except Exception:
+        pass
+
+
+def append_jsonl(path: str, record: Dict[str, Any]) -> None:
+    """Append a JSON record as one line to a JSONL file (best-effort).
+
+    This is safe to call from multiple processes when each process writes
+    to a distinct file. For concurrent writes to the same file, OS append
+    semantics typically suffice, but no additional locking is provided.
+    """
+    try:
+        _ensure_parent_dir(path)
+        line = json.dumps(record, ensure_ascii=False)
+        fd = os.open(path, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o644)
+        with os.fdopen(fd, "a", encoding="utf-8") as f:
+            f.write(line + "\n")
+    except Exception:
+        logger.exception("Failed to append JSONL to %s", path)
+
+
+def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]:
+    """Dump object to SHM if serialized size exceeds threshold.
+
+    Returns (True, meta) when dumped; otherwise (False, original_obj).
+    """
+    payload = serialize_obj(obj)
+    if len(payload) > threshold:
+        return True, shm_write_bytes(payload)
+    return False, obj
+
+
+def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -> Any:
+    """Load object from container that may carry SHM or inline object."""
+    if shm_key in container:
+        return pickle.loads(shm_read_bytes(container[shm_key]))
+    return container[obj_key]
+
+
+def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]:
+    """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key).
+
+    When serialized size exceeds threshold, returns {shm_key: {name,size}};
+    otherwise returns {obj_key: obj}.
+    """
+    payload: Dict[str, Any] = {}
+    use_shm, data = maybe_dump_to_shm(obj, threshold)
+    if use_shm:
+        payload[shm_key] = data
+    else:
+        payload[obj_key] = data
+    return payload
+
+
+# Convert OmegaConf/objects to plain dicts
+def _to_dict(x: Any) -> Dict[str, Any]:
+    try:
+        if isinstance(x, dict):
+            return dict(x)
+        return OmegaConf.to_container(x, resolve=True)  # type: ignore[arg-type]
+    except Exception:
+        try:
+            return dict(x)
+        except Exception:
+            return {}

vllm_omni/entrypoints/pipelined_omni_llm.py

@@ -0,0 +1,621 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import os
+import sys
+import queue
+import time
+import multiprocessing as mp
+import logging
+
+from vllm.inputs import PromptType
+from vllm.logger import init_logger
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM, OmniStageLLM
+from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.entrypoints.pipeline_utils import (
+    maybe_load_from_ipc as _load,
+    encode_for_ipc as _encode,
+    serialize_obj as _ser,
+    append_jsonl as _append_jsonl,
+)
+
+logger = init_logger(__name__)
+
+
+class PipelinedOmniLLM(OmniLLM):
+    """Multi-process pipelined OmniLLM.
+
+    - Per-stage process with copy-based IPC (Queues)
+    - window=-1 across stages (downstream starts when upstream finishes)
+    - max_inflight=1 per stage (serial within a stage), but pipeline across stages
+    """
+
+    def __init__(self, model: str,
+                 stage_configs=None,
+                 log_stats: bool = True,
+                 log_file: Optional[str] = None,
+                 init_sleep_seconds: int = 20,
+                 shm_threshold_bytes: int = 65536,
+                 batch_timeout: int = 10,
+                 init_timeout: int = 300,
+                 **kwargs):
+        self.batch_timeout = batch_timeout
+        self._enable_stats: bool = bool(log_stats)
+        # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
+        if stage_configs is None:
+            self.stage_configs = load_stage_configs_from_model(model)
+        else:
+            self.stage_configs = stage_configs
+
+        self.stage_list: List[OmniStage] = []
+        for stage_config in self.stage_configs:
+            # Only construct lightweight OmniStage without setting engine in parent
+            stage = OmniStage(stage_config)
+            self.stage_list.append(stage)
+        logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
+
+        self._ctx = mp.get_context("spawn")
+        self._stage_in_queues: List[mp.Queue] = []
+        self._stage_out_queues: List[mp.Queue] = []
+        # Optional file handler for orchestrator
+        self._log_file = log_file
+        if self._log_file:
+            try:
+                # Avoid duplicate handlers
+                has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
+                if not has_file_handler:
+                    fh = logging.FileHandler(self._log_file)
+                    fh.setLevel(logging.DEBUG)
+                    fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
+                    logger.addHandler(fh)
+                    logger.setLevel(logging.DEBUG)
+            except Exception:
+                pass
+
+        # Orchestrator stats JSONL file
+        self._stats_file: Optional[str] = None
+        if self._enable_stats and self._log_file:
+            try:
+                self._stats_file = f"{self._log_file}.orchestrator.stats.jsonl"
+            except Exception:
+                self._stats_file = None
+        # Overall stats JSONL file (per-request + summary)
+        self._overall_stats_file: Optional[str] = None
+        if self._enable_stats and self._log_file:
+            try:
+                self._overall_stats_file = f"{self._log_file}.overall.stats.jsonl"
+            except Exception:
+                self._overall_stats_file = None
+
+        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)
+        # 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:
+        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)
+            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            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)
+            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,
+            )
+            logger.debug("[Orchestrator] Stage-%s process started", stage_id)
+
+    def close(self) -> None:
+        for q in self._stage_in_queues:
+            try:
+                q.put_nowait(None)
+            except Exception:
+                pass
+        for stage in self.stage_list:
+            try:
+                stage.stop_stage_worker()
+            except Exception:
+                pass
+
+    def __del__(self) -> None:  # best-effort
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    def generate(
+        self,
+        prompts: Union[PromptType, Sequence[PromptType]],
+        sampling_params_list: Optional[
+            Union[SamplingParams, Sequence[SamplingParams]]
+        ] = None,
+    ) -> List[OmniRequestOutput]:
+        logger.debug("[Orchestrator] generate() called")
+        if sampling_params_list is None:
+            raise ValueError("sampling_params_list is required for pipelined generation")
+        if len(sampling_params_list) != len(self.stage_list):
+            raise ValueError(
+                f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}"
+            )
+
+        # Normalize prompts to a list for per-request iteration
+        if not isinstance(prompts, (list, tuple)):
+            request_prompts: List[PromptType] = [prompts]
+        else:
+            request_prompts = list(prompts)
+
+        final_outputs: List[OmniRequestOutput] = []
+
+        # Orchestrator keeps stage objects for input derivation
+        num_stages = len(self.stage_list)
+
+        # Map from request_id to original prompt
+        request_id_to_prompt: Dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+
+        # Track per-request start time for end-to-end timing
+        _req_start_ts: Dict[int, float] = {}
+        _wall_start_ts: float = time.time()
+        _last_finish_ts: float = _wall_start_ts
+
+        # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
+        final_stage_id_for_e2e = -1
+        try:
+            for _sid, _st in enumerate(self.stage_list):
+                if getattr(_st, "final_output", False):
+                    final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid)
+            if final_stage_id_for_e2e < 0:
+                final_stage_id_for_e2e = len(self.stage_list) - 1
+        except Exception:
+            final_stage_id_for_e2e = len(self.stage_list) - 1
+        # In-memory aggregators for this generate() call
+        stage_total_time_ms: List[float] = [0.0 for _ in range(num_stages)]
+        stage_total_tokens: List[int] = [0 for _ in range(num_stages)]
+        stage_req_counts: List[int] = [0 for _ in range(num_stages)]
+        transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {}
+        # Per-edge per-request sender timing to combine with receiver timing later
+        transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {}
+        e2e_total_ms: float = 0.0
+        e2e_total_tokens: int = 0
+        e2e_count: int = 0
+        e2e_done: set[int] = set()
+        # Per-request overall aggregation
+        per_request: Dict[int, Dict[str, Any]] = {}
+        sum_per_request_transfer_ms: float = 0.0
+
+        # Seed stage-0 queue with all requests
+        logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
+
+        for req_id, prompt in request_id_to_prompt.items():
+            sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
+            task = {
+                "request_id": req_id,
+                "engine_inputs": prompt,
+                "sampling_params": sp0,
+            }
+            self.stage_list[0].submit(task)
+            _req_start_ts[req_id] = time.time()
+            logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id)
+
+        # For each stage, forward results to next stage; collect finals at the end
+        # We pipeline by continually polling output queues in stage order
+        remaining_by_stage: List[int] = [len(request_prompts)] + [0] * (num_stages - 1)
+        completed_requests = 0
+        total_requests = len(request_prompts)
+
+        logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages)
+        while completed_requests < total_requests:
+            made_progress = False
+            for stage_id, stage in enumerate(self.stage_list):
+                result = stage.try_collect()
+                if result is None:
+                    continue
+
+                made_progress = True
+                req_id = result.get("request_id")
+                if "error" in result:
+                    logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"])
+                    continue
+                
+                if result.get("type") == "stage_ready":
+                    #Only happens when stage is initialized slower than expected, so we wait for a short time and try again
+                    time.sleep(0.05)
+                    continue
+
+                engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
+                # Aggregate per-request metrics from stage worker if present
+                try:
+                    _m = result.get("metrics")
+                    if _m is not None:
+                        stage_req_counts[stage_id] += 1
+                        stage_total_time_ms[stage_id] += float(_m.get("stage_gen_time_ms", 0.0))
+                        stage_total_tokens[stage_id] += int(_m.get("num_tokens_out", 0))
+                        # record per-request stage metrics
+                        try:
+                            rid_int = int(req_id)
+                            pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+                            pr_stages = pr["stages"]  # type: ignore[index]
+                            pr_stages[stage_id] = {
+                                "stage_gen_time_ms": float(_m.get("stage_gen_time_ms", 0.0)),
+                                "num_tokens_out": int(_m.get("num_tokens_out", 0)),
+                            }
+                        except Exception:
+                            pass
+                        # Also aggregate receiver-side transfer decode time for (stage_id-1 -> stage_id)
+                        try:
+                            if stage_id > 0:
+                                key = (stage_id - 1, stage_id)
+                                agg = transfer_agg.get(key)
+                                if agg is None:
+                                    agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
+                                           "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
+                                           "sum_total_ms": 0.0, "total_count": 0.0}
+                                    transfer_agg[key] = agg
+                                rx_b = float(_m.get("rx_transfer_bytes", 0.0))
+                                rx_ms = float(_m.get("rx_decode_time_ms", 0.0))
+                                in_flight_ms = float(_m.get("rx_in_flight_time_ms", 0.0))
+                                agg["sum_rx_bytes"] += rx_b
+                                agg["sum_rx_ms"] += rx_ms
+                                agg["rx_count"] += 1.0
+                                # If we have sender-side timing stored, compute combined (encode+enqueue + decode)
+                                try:
+                                    edge_req = (stage_id - 1, stage_id, int(req_id))
+                                    s = transfer_edge_req.get(edge_req)
+                                    if s is not None:
+                                        total_ms = float(s.get("tx_ms", 0.0)) + in_flight_ms + rx_ms
+                                        agg["sum_total_ms"] += total_ms
+                                        agg["total_count"] += 1.0
+                                        # accumulate per-request transfer totals
+                                        try:
+                                            rid_int = int(req_id)
+                                            pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+                                            pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms  # type: ignore[index]
+                                            pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_b)  # type: ignore[index]
+                                        except Exception:
+                                            pass
+                                        if getattr(self, "_stats_file", None):
+                                            try:
+                                                size_b = float(s.get("size_bytes", rx_b))
+                                                _append_jsonl(self._stats_file, {  # type: ignore[arg-type]
+                                                    "type": "transfer_total_stats",
+                                                    "from_stage": stage_id - 1,
+                                                    "to_stage": stage_id,
+                                                    "request_id": req_id,
+                                                    "size_bytes": int(size_b),
+                                                    "tx_time_ms": float(s.get("tx_ms", 0.0)),
+                                                    "in_flight_time_ms": in_flight_ms,
+                                                    "rx_decode_time_ms": rx_ms,
+                                                    "total_time_ms": total_ms,
+                                                    "total_time_per_kb_ms": total_ms / max(size_b / 1024.0, 1e-6) if size_b > 0 else 0.0,
+                                                })
+                                            except Exception:
+                                                pass
+                                except Exception:
+                                    pass
+                                # Emit per-request RX stats to JSONL
+                                if getattr(self, "_stats_file", None):
+                                    try:
+                                        _append_jsonl(self._stats_file, {  # type: ignore[arg-type]
+                                            "type": "transfer_rx_stats",
+                                            "from_stage": stage_id - 1,
+                                            "to_stage": stage_id,
+                                            "request_id": req_id,
+                                            "rx_bytes": int(rx_b),
+                                            "rx_decode_time_ms": rx_ms,
+                                            "in_flight_time_ms": in_flight_ms,
+                                            "rx_time_per_kb_ms": rx_ms / max(rx_b / 1024.0, 1e-6) if rx_b > 0 else 0.0,
+                                        })
+                                    except Exception:
+                                        pass
+                        except Exception:
+                            pass
+                except Exception:
+                    pass
+                logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id)
+                stage.set_engine_outputs(engine_outputs)
+
+                if getattr(stage, "final_output", False):
+                    final_outputs.append(
+                        OmniRequestOutput(
+                            stage_id=stage_id,
+                            final_output_type=stage.final_output_type,  # type: ignore[attr-defined]
+                            request_output=engine_outputs,
+                        )
+                    )
+                    logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id)
+
+                    # End-to-end timing and time-per-token for final output (only once per request at the designated final stage)
+                    try:
+                        if stage_id == final_stage_id_for_e2e and req_id not in e2e_done:
+                            _t0 = _req_start_ts.get(req_id)
+                            if _t0 is not None:
+                                _t1 = time.time()
+                                _last_finish_ts = max(_last_finish_ts, _t1)
+                                _e2e_ms = (_t1 - _t0) * 1000.0
+                                # Count tokens from final stage outputs
+                                def _count_tokens(_ros: List[Any]) -> int:  # type: ignore[name-defined]
+                                    total = 0
+                                    for _ro in _ros:
+                                        try:
+                                            outs = getattr(_ro, "outputs", None)
+                                            if outs and len(outs) > 0:
+                                                tokens = getattr(outs[0], "token_ids", None)
+                                                if tokens is not None:
+                                                    total += len(tokens)
+                                        except Exception:
+                                            pass
+                                    return total
+                                _num_tokens = _count_tokens(engine_outputs)
+                                _time_per_token_ms = (_e2e_ms / _num_tokens) if _num_tokens > 0 else 0.0
+                                # Update E2E aggregators
+                                e2e_total_ms += _e2e_ms
+                                e2e_total_tokens += int(_num_tokens)
+                                e2e_count += 1
+                                e2e_done.add(req_id)
+                                # Write per-request overall record
+                                try:
+                                    rid_int = int(req_id)
+                                    pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+                                    per_req_record = {
+                                        "type": "overall_request",
+                                        "request_id": rid_int,
+                                        "e2e_time_ms": _e2e_ms,
+                                        "num_tokens_out": int(_num_tokens),
+                                        "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
+                                        "transfers_total_bytes": int(pr.get("transfers_bytes", 0)),
+                                        "stages": pr.get("stages", {}),
+                                    }
+                                    sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
+                                    if getattr(self, "_overall_stats_file", None):
+                                        try:
+                                            _append_jsonl(self._overall_stats_file, per_req_record)  # type: ignore[arg-type]
+                                        except Exception:
+                                            pass
+                                except Exception:
+                                    pass
+                                if getattr(self, "_stats_file", None):
+                                    try:
+                                        _append_jsonl(self._stats_file, {  # type: ignore[arg-type]
+                                            "type": "orchestrator_request_e2e",
+                                            "request_id": req_id,
+                                            "final_stage_id": stage_id,
+                                            "e2e_time_ms": _e2e_ms,
+                                            "num_tokens_out": int(_num_tokens),
+                                            "e2e_time_per_token_ms": _time_per_token_ms,
+                                        })
+                                    except Exception:
+                                        pass
+                    except Exception:
+                        pass
+
+                next_stage_id = stage_id + 1
+                if next_stage_id < num_stages:
+                    next_stage: OmniStage = self.stage_list[next_stage_id]
+                    next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
+                    sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
+                    try:
+                        # Measure transfer size and time (encode + enqueue)
+                        size_bytes = 0
+                        try:
+                            size_bytes = len(_ser(next_inputs))
+                        except Exception:
+                            size_bytes = 0
+                        t0 = time.time()
+                        ipc_payload = _encode(
+                            next_inputs,
+                            getattr(self, "_shm_threshold_bytes", 65536),
+                            obj_key="engine_inputs",
+                            shm_key="engine_inputs_shm",
+                        )
+                        ipc_payload.update({
+                            "request_id": req_id,
+                            "sampling_params": sp_next,
+                            "sent_ts": time.time(),
+                        })
+                        self.stage_list[next_stage_id].submit(ipc_payload)
+                        t1 = time.time()
+                        tx_ms = (t1 - t0) * 1000.0
+                        if self._enable_stats and getattr(self, "_stats_file", None):
+                            try:
+                                _append_jsonl(self._stats_file, {
+                                    "type": "transfer_stats",
+                                    "from_stage": stage_id,
+                                    "to_stage": next_stage_id,
+                                    "request_id": req_id,
+                                    "size_bytes": int(size_bytes),
+                                    "tx_time_ms": tx_ms,
+                                    "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_ms, 1e-6) * 1000.0),
+                                    "used_shm": bool("engine_inputs_shm" in ipc_payload),
+                                })
+                            except Exception:
+                                pass
+                        # Update in-memory transfer aggregator
+                        try:
+                            key = (stage_id, next_stage_id)
+                            agg = transfer_agg.get(key)
+                            if agg is None:
+                                agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
+                                       "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
+                                       "sum_total_ms": 0.0, "total_count": 0.0}
+                                transfer_agg[key] = agg
+                            agg["sum_bytes"] += float(size_bytes)
+                            agg["sum_ms"] += float(tx_ms)
+                            agg["count"] += 1.0
+                            # Store sender-side timing for per-request combination
+                            transfer_edge_req[(stage_id, next_stage_id, int(req_id))] = {
+                                "tx_ms": float(tx_ms),
+                                "size_bytes": float(size_bytes),
+                            }
+                        except Exception:
+                            pass
+                    except Exception:
+                        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, next_stage_id)
+                    remaining_by_stage[next_stage_id] += 1
+                else:
+                    completed_requests += 1
+                    logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests)
+
+            if not made_progress:
+                time.sleep(0.005)
+        logger.debug("[Orchestrator] All requests completed")
+
+        # Summarize and print stats
+        try:
+            stage_summary: List[Dict[str, Any]] = []
+            for sid in range(num_stages):
+                reqs = stage_req_counts[sid]
+                tokens = stage_total_tokens[sid]
+                total_ms = float(stage_total_time_ms[sid])
+                avg_req = (total_ms / reqs) if reqs > 0 else 0.0
+                avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
+                stage_summary.append({
+                    "stage_id": sid,
+                    "requests": int(reqs),
+                    "tokens": int(tokens),
+                    "total_time_ms": total_ms,
+                    "avg_time_per_request_ms": avg_req,
+                    "avg_tokens_per_s": avg_tok,
+                })
+
+            transfer_summary: List[Dict[str, Any]] = []
+            for (src, dst), agg in transfer_agg.items():
+                sum_bytes = float(agg.get("sum_bytes", 0.0))
+                sum_ms = float(agg.get("sum_ms", 0.0))
+                samples = int(agg.get("count", 0.0))
+                tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
+                sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
+                sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
+                samples_rx = int(agg.get("rx_count", 0.0))
+                rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
+                sum_total_ms = float(agg.get("sum_total_ms", 0.0))
+                samples_total = int(agg.get("total_count", 0.0))
+                total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
+                transfer_summary.append({
+                    "from_stage": src,
+                    "to_stage": dst,
+                    "samples": samples,
+                    "total_bytes": int(sum_bytes),
+                    "total_time_ms": sum_ms,
+                    "tx_mbps": tx_mbps,
+                    "rx_samples": samples_rx,
+                    "rx_total_bytes": int(sum_rx_bytes),
+                    "rx_total_time_ms": sum_rx_ms,
+                    "rx_mbps": rx_mbps,
+                    "total_samples": samples_total,
+                    "total_transfer_time_ms": sum_total_ms,
+                    "total_mbps": total_mbps,
+                })
+
+            e2e_avg_req = (e2e_total_ms / e2e_count) if e2e_count > 0 else 0.0
+            e2e_avg_tok = (e2e_total_tokens * 1000.0 / e2e_total_ms) if e2e_total_ms > 0 else 0.0
+            wall_time_ms = max(0.0, (_last_finish_ts - _wall_start_ts) * 1000.0)
+            summary: Dict[str, Any] = {
+                "e2e_requests": int(e2e_count),
+                # 按你的需求：e2e_total_time_ms 直接使用墙钟时间（非各请求之和）
+                "e2e_total_time_ms": float(wall_time_ms),
+                # 额外保留各请求 E2E 之和，供需要时参考
+                "e2e_sum_time_ms": float(e2e_total_ms),
+                "e2e_total_tokens": int(e2e_total_tokens),
+                "e2e_avg_time_per_request_ms": e2e_avg_req,
+                "e2e_avg_tokens_per_s": e2e_avg_tok,
+                "wall_time_ms": wall_time_ms,
+                "final_stage_id": final_stage_id_for_e2e,
+                "stages": stage_summary,
+                "transfers": transfer_summary,
+            }
+            logger.info("[Summary] %s", summary)
+            if self._enable_stats and getattr(self, "_stats_file", None):
+                try:
+                    _append_jsonl(self._stats_file, {"type": "orchestrator_summary", **summary})  # type: ignore[arg-type]
+                except Exception:
+                    pass
+            if self._enable_stats and getattr(self, "_overall_stats_file", None):
+                try:
+                    _append_jsonl(self._overall_stats_file, {"type": "overall_summary", **summary})  # type: ignore[arg-type]
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+        return final_outputs
+
+    def _wait_for_stages_ready(self, timeout: int = 120) -> None:
+        deadline = time.time() + max(0, int(timeout))
+        num_stages = len(self.stage_list)
+        while len(self._stages_ready) < num_stages and time.time() < deadline:
+            progressed = False
+            for stage_id, stage in enumerate(self.stage_list):
+                if stage_id in self._stages_ready:
+                    continue
+                result = stage.try_collect()
+                if result is None:
+                    continue
+                progressed = True
+                if result.get("type") == "stage_ready":
+                    self._stages_ready.add(stage_id)
+                    logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
+                else:
+                    # No user data should arrive before seeding; ignore other messages
+                    pass
+            if not progressed:
+                time.sleep(0.01)
+        if len(self._stages_ready) < num_stages:
+            not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
+            logger.warning(
+                "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
+                len(self._stages_ready), num_stages, not_ready,
+            )
+            # Provide actionable suggestions before shutdown
+            try:
+                suggestions = [
+                    "Verify GPU/device assignment in config (mrs.devices) is correct.",
+                    "Check GPU/host memory availability; reduce model or batch size if needed.",
+                    "Check model weights path and network reachability (if loading remotely).",
+                    "Increase initialization wait time (init_sleep_seconds or call-site timeout).",
+                ]
+                if getattr(self, "_log_file", None):
+                    suggestions.append(
+                        f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log"
+                    )
+                logger.error(
+                    "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
+                    "\n- ".join(suggestions),
+                )
+            except Exception:
+                # Best-effort logging of suggestions
+                logger.error(
+                    "[Orchestrator] Stage initialization failed and an error occurred while logging suggestions",
+                )
+
+            # Attempt graceful shutdown of all stages before exiting
+            try:
+                self.close()
+            except Exception:
+                pass
+
+            # Terminate the current process with non-zero exit code
+            try:
+                sys.exit(1)
+            except SystemExit:
+                raise
+            except Exception:
+                os._exit(1)
+
+

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -269,12 +269,17 @@ class Qwen2_5OmniForConditionalGeneration(
 
         # 2) Talker (if codec not provided)
         if self.model_stage == "talker":
+            # Mixed-mode support: in一个step内同时支持 Prefill*n 与 Decode*n。
+            # 规则：
+            # - 使用特殊token将 Prefill 片段包裹为 [BOS][PAD...][EOS]
+            # - Decode 片段则为“单一非特殊token”
+            # - 若提供了 additional_information（可为按请求分割的 list 或拼接后的 tensor + 形状列表），
+            #   则对 Prefill 片段按请求重建 thinker→talker 的输入嵌入；
+            # - 对 Decode 片段，若提供了 per-request 的 decode 辅助嵌入（可选），则进行叠加；否则保留原嵌入。
 
             if input_ids is None and additional_information is None:
                 input_ids = torch.zeros(
-                    inputs_embeds.shape[0],
-                    dtype=torch.long,
-                    device=inputs_embeds.device,
+                    inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
                 )
                 additional_information = {}
                 self.thinker_reply_part = torch.zeros_like(inputs_embeds)
@@ -282,63 +287,109 @@ class Qwen2_5OmniForConditionalGeneration(
             else:
                 is_profile = False
 
+            # Ensure we have base embeddings when only ids are provided
+            if inputs_embeds is None and input_ids is not None:
+                inputs_embeds = self.talker.get_input_embeddings(input_ids)
+
+            # ------- Request-scoped additional information (no cross-request concat) -------
+            request_ids: Optional[list[str]] = kwargs.get("request_ids")  # ordered
+            request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans")
+            addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
+            runtime_addi = kwargs.get("runtime_additional_information")
+
+            # Normalize runtime_addi into a mapping by request_id for convenience
+            runtime_addi_by_req: dict[str, dict] = {}
+            if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids):
+                for i, rid in enumerate(request_ids):
+                    if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
+                        runtime_addi_by_req[rid] = runtime_addi[i]
+            elif isinstance(request_ids, list) and isinstance(runtime_addi, dict):
+                for rid in request_ids:
+                    if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict):
+                        runtime_addi_by_req[rid] = runtime_addi[rid]
+
+            # Containers to return per-request updates (e.g., thinker_reply_part_per_request)
+            update_by_req_id: dict[str, dict] = {}
+
+            # ------- Prefill: span_len > 1 -------
             if (
-                input_ids is not None
-                and additional_information is not None
-                and not is_profile
+                not is_profile
+                and isinstance(request_ids, list)
+                and isinstance(request_token_spans, list)
+                and isinstance(addi_by_req, dict)
             ):
-                # read from additional_information dict
-                thinker_result = None
-                if additional_information is not None and isinstance(
-                    additional_information, dict
-                ):
-                    thinker_result = additional_information.get("thinker_result")
-                    prompt_embeds = additional_information.get("prompt_embeds")
-                    prompt_token_ids = additional_information.get("prompt_token_ids")
-                    thinker_output_token_ids = additional_information.get(
-                        "thinker_output_token_ids"
-                    )
-                else:
-                    thinker_result = torch.zeros_like(inputs_embeds)
-                    prompt_embeds = torch.zeros_like(inputs_embeds)
-                    prompt_token_ids = torch.zeros(
-                        inputs_embeds.shape[0],
-                        dtype=torch.int64,
-                        device=inputs_embeds.device,
-                    )
-                    thinker_output_token_ids = torch.zeros(
-                        inputs_embeds.shape[0],
-                        dtype=torch.int64,
-                        device=inputs_embeds.device,
-                    )
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    span_len = int(e) - int(s)
+                    if span_len <= 1:
+                        continue
+                    info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
+                    if not isinstance(info, dict):
+                        info = {}
+                    pe = info.get("prompt_embeds")  # Tensor [P,H]
+                    tr = info.get("thinker_result")  # Tensor [K,H]
+                    ptoks = info.get("prompt_token_ids")  # list[int]
+                    otoks = info.get("thinker_output_token_ids")  # list[int]
 
-                if thinker_result is None:
-                    thinker_result = torch.zeros_like(inputs_embeds)
-                self.thinker_reply_part = thinker_result.squeeze(0)
-                if self.thinker_reply_part.shape[1] > 1:
-                    self.thinker_reply_part = self.thinker_reply_part[1:, :]
-                input_ids, inputs_embeds = self._thinker_to_talker_prefill(
-                    voice_type=voice_type,
-                    output_prompt_embeds=thinker_result,
-                    output_token_ids=thinker_output_token_ids,
-                    thinker_prompt_embeds=prompt_embeds,
-                    prompt_token_ids=prompt_token_ids,
-                )
-            elif not is_profile:
-                input_ids, inputs_embeds = self._thinker_to_talker_decode_one_step(
-                    output_prompt_embeds=(
-                        self.thinker_reply_part[:1]
-                        if self.thinker_reply_part.shape[0] >= 1
-                        else torch.zeros(1, self.thinker_reply_part.shape[1])
-                        .to(self._module_device(self.model))
-                        .to(torch.bfloat16)
-                        + (-1.25 * 2 ** (-123))
-                    ),
-                    output_token_ids=input_ids,
-                )
+                    if not isinstance(pe, torch.Tensor):
+                        pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                    if not isinstance(tr, torch.Tensor):
+                        tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                    if not isinstance(ptoks, (list, torch.Tensor)):
+                        ptoks = []
+                    if not isinstance(otoks, (list, torch.Tensor)):
+                        otoks = []
 
-                if self.thinker_reply_part.shape[0] >= 1:
-                    self.thinker_reply_part = self.thinker_reply_part[1:, :]
+                    req_input_ids, req_embeds = self._thinker_to_talker_prefill(
+                        voice_type=voice_type,
+                        output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        output_token_ids=otoks,
+                        thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        prompt_token_ids=ptoks,
+                    )
+                    seg_len = min(span_len, req_embeds.shape[0])
+                    inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
+                    if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
+                        input_ids[s : s + seg_len] = req_input_ids
+
+                    # Prepare per-request reply queue for subsequent decode: drop first row
+                    if tr.ndim == 2 and tr.shape[0] > 0:
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous()
+
+            # ------- Decode: span_len == 1 -------
+            if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    if (int(e) - int(s)) != 1:
+                        continue
+                    # choose step vector in priority order
+                    step_vec = None
+                    # A) runtime queue
+                    q = None
+                    if isinstance(rid, str):
+                        q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
+                    if isinstance(q, torch.Tensor) and q.numel() > 0:
+                        step_vec = q[0:1]
+                        new_q = q[1:].detach().to("cpu").contiguous()
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
+                    else:
+                        # B) per-request provided decode vector (optional)
+                        info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
+                        dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
+                        if isinstance(dv, torch.Tensor) and dv.numel() > 0:
+                            step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
+                        elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
+                            # C) fallback shared pool
+                            step_vec = self.thinker_reply_part[0:1]
+                            self.thinker_reply_part = self.thinker_reply_part[1:]
+
+                    if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
+                        one_id = input_ids[s : s + 1]
+                        _, one_embed = self._thinker_to_talker_decode_one_step(
+                            output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                            output_token_ids=one_id,
+                        )
+                        inputs_embeds[s] = one_embed[0]
 
             with torch.inference_mode():
                 talker_hidden = self.talker(
@@ -346,8 +397,19 @@ class Qwen2_5OmniForConditionalGeneration(
                     positions=positions[0],
                     inputs_embeds=inputs_embeds,
                 )
+            multimodal_outputs: dict = None
+            # Return updates if any
+            if update_by_req_id:
+                multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
 
-            return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=None)
+            if sampling_metadata is not None:
+                # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size
+                sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448
+
+            return OmniOutput(
+                text_hidden_states=talker_hidden,
+                multimodal_outputs=multimodal_outputs,
+            )
 
         if self.model_stage == "code2wav":
             code = (
@@ -359,8 +421,12 @@ class Qwen2_5OmniForConditionalGeneration(
                     device=inputs_embeds.device,
                 )
             )
+            
+            code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code
+            code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code
+            
             audio_tensor = self.generate_audio(
-                code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code, voice_type
+                code, voice_type
             )
             return OmniOutput(
                 text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -1,36 +1,50 @@
 # stage config for running qwen2.5-omni with architecture of OmniLLM.
 stage_args:
   - stage_id: 0
+    mrs:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-      gpu_memory_utilization: 0.32
+      gpu_memory_utilization: 0.8
       enforce_eager: true  # need to discuss
       trust_remote_code: true
       engine_output_type: latent  # change the param name,such as pooling_output
+      enable_prefix_caching: false
     final_output: true
     final_output_type: text
   - stage_id: 1
+    mrs:
+      process: true
+      devices: "1"
+      max_batch_size: 3
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-      gpu_memory_utilization: 0.32
+      gpu_memory_utilization: 0.8
       enforce_eager: true
       trust_remote_code: true
+      enable_prefix_caching: false
       engine_output_type: latent
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
   - stage_id: 2
+    mrs:
+      process: true
+      devices: "0"            # 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.3
+      gpu_memory_utilization: 0.15
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
@@ -38,3 +52,17 @@ stage_args:
     engine_input_source: [1]
     final_output: true
     final_output_type: audio
+
+# Top-level MRS config (concise): default windows and stage edges
+mrs:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+  edges:
+    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
+      to: 1
+      window_size: -1
+    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
+      to: 2
+      window_size: -1

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -4,8 +4,9 @@ import torch
 
 from vllm.inputs import TextPrompt
 from vllm_omni.inputs.data import OmniTokensPrompt
-
-
+TALKER_CODEC_PAD_TOKEN_ID = 8292
+TALKER_CODEC_START_TOKEN_ID = 8293
+TALKER_CODEC_END_TOKEN_ID = 8294
 def thinker2talker(
     stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None
 ):
@@ -33,8 +34,7 @@ def thinker2talker(
         )
         talker_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=[0]
-                * (len(prompt_token_ids) + 2),  # add 2 for codec pad and start token
+                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) + [TALKER_CODEC_END_TOKEN_ID],
                 additional_information={
                     "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(
                         torch.float32
@@ -44,6 +44,8 @@ def thinker2talker(
                     ),
                     "prompt_token_ids": prompt_token_ids,
                     "thinker_output_token_ids": thinker_output_ids,
+                    "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
+                    "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
                 },
                 multi_modal_data=(
                     multi_modal_data[thinker_output.request_id]

vllm_omni/worker/gpu_ar_model_runner.py

@@ -99,9 +99,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
             # TODO(woosuk): Avoid the copy. Optimize.
             self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
 
-            # Reset per-step additional information collector
+            # Reset per-step additional information collector (deprecated concat path)
             if hasattr(self, "_forward_additional_information"):
                 self._forward_additional_information = None
+            # New: per-request additional information for this step
+            per_req_additional_information: dict[str, dict] = {}
 
             # Overlay custom prompt_embeds per request for the prompt portion;
             # collect additional_information (tensor/list) for prefill portion only
@@ -126,50 +128,27 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
                         src
                     )
-                # For additional_information: handle arbitrary keys
+                # Build per-request additional information (no cross-request concat)
                 if addi_cpu is not None and isinstance(addi_cpu, dict):
-                    # Lazy init collector dict
-                    if (
-                        not hasattr(self, "_forward_additional_information")
-                        or self._forward_additional_information is None
-                    ):
-                        self._forward_additional_information = {}
-                    # Process tensors (slice by scheduled prompt range) and
-                    # lists (append per-request)
+                    req_info: dict[str, object] = {}
                     for k, v in addi_cpu.items():
                         if isinstance(v, torch.Tensor):
-                            # Slice along token dimension for prefill part
-                            try:
-                                seg = v[
-                                    num_computed_tokens : num_computed_tokens
-                                    + overlay_len
-                                ].to(
-                                    dtype=self.dtype,
-                                    device=self.device,
-                                    non_blocking=True,
-                                )
-                            except Exception:
-                                # Fallback: move whole tensor if slicing fails
-                                seg = v.to(
-                                    dtype=self.dtype,
-                                    device=self.device,
-                                    non_blocking=True,
-                                )
-                            prev_val = self._forward_additional_information.get(k)
-                            self._forward_additional_information[k] = (
-                                torch.cat([prev_val, seg], dim=0)
-                                if isinstance(prev_val, torch.Tensor)
-                                else seg.clone()
-                            )
-                        elif isinstance(v, list):
-                            prev_val = self._forward_additional_information.get(k)
-                            if prev_val is None:
-                                self._forward_additional_information[k] = [v]
-                            elif isinstance(prev_val, list):
-                                self._forward_additional_information[k].append(v)
+                            # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor
+                            if overlay_len > 0:
+                                try:
+                                    seg = v[
+                                        num_computed_tokens : num_computed_tokens + overlay_len
+                                    ].detach().to("cpu").contiguous()
+                                except Exception:
+                                    seg = v.detach().to("cpu").contiguous()
+                                req_info[k] = seg
                             else:
-                                # Mixed types: wrap existing into list
-                                self._forward_additional_information[k] = [prev_val, v]
+                                req_info[k] = v.detach().to("cpu").contiguous()
+                        elif isinstance(v, list):
+                            req_info[k] = v
+                        else:
+                            req_info[k] = v
+                    per_req_additional_information[req_id] = req_info
 
             input_ids = self.input_ids.gpu[:num_input_tokens]
             inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
@@ -395,15 +374,33 @@ class GPUARModelRunner(OmniGPUModelRunner):
         ):
 
             model_kwargs_extra = {}
-            # Only pass additional_information for the prefill part
-            if (
-                hasattr(self, "_forward_additional_information")
-                and self._forward_additional_information is not None
-                and isinstance(self._forward_additional_information, dict)
-            ):
-                model_kwargs_extra["additional_information"] = (
-                    self._forward_additional_information
+            # Pass per-request additional information map for this step (no concat)
+            if per_req_additional_information:
+                model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
+            # 始终传递每请求的运行期 additional_information（持久化在 request state）
+            try:
+                per_req_runtime_info = []
+                for req_id in self.input_batch.req_ids:
+                    req_state = self.requests.get(req_id)
+                    info = (
+                        getattr(req_state, "additional_information_cpu", None)
+                        if req_state is not None
+                        else None
+                    )
+                    per_req_runtime_info.append(info if isinstance(info, dict) else {})
+                model_kwargs_extra["runtime_additional_information"] = (
+                    per_req_runtime_info
                 )
+                model_kwargs_extra["request_ids"] = self.input_batch.req_ids
+                # 传递每个请求在本步扁平序列中的 token 区间，便于模型按请求映射 decode/prefill
+                req_token_spans = []
+                for req_index in range(len(self.input_batch.req_ids)):
+                    start_offset = int(self.query_start_loc_cpu[req_index])
+                    sched_tokens = int(num_scheduled_tokens_np[req_index])
+                    req_token_spans.append((start_offset, start_offset + sched_tokens))
+                model_kwargs_extra["request_token_spans"] = req_token_spans
+            except Exception:
+                pass
             model_output = self.model(
                 input_ids=input_ids,
                 positions=positions,
@@ -428,8 +425,42 @@ class GPUARModelRunner(OmniGPUModelRunner):
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
                 hidden_states
             )
-
+            # 模型侧可能返回 per-request 的 additional_information 更新（与模型无关的通用通道）。
+            # 约定：multimodal_outputs["additional_information_update"] 为按 batch 顺序的 list[dict]，
+            # runner 将其合并到对应 request 的 additional_information_cpu，供后续 decode 使用。
+            try:
+                if (
+                    isinstance(multimodal_outputs, dict)
+                    and (
+                        "additional_information_update" in multimodal_outputs
+                        or "additional_information_update_by_req_id" in multimodal_outputs
+                    )
+                ):
+                    # 方案 A：按 batch 顺序的 list[dict]
+                    updates_list = multimodal_outputs.get(
+                        "additional_information_update"
+                    )
+                    if isinstance(updates_list, list):
+                        for idx, upd in enumerate(updates_list):
+                            if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
+                                continue
+                            req_id = self.input_batch.req_ids[idx]
+                            self._merge_additional_information_update(req_id, upd)
+                    # 方案 B：按 req_id 的 dict[str, dict]
+                    updates_map = multimodal_outputs.get(
+                        "additional_information_update_by_req_id"
+                    )
+                    if isinstance(updates_map, dict):
+                        for req_id, upd in updates_map.items():
+                            if not isinstance(upd, dict):
+                                continue
+                            if req_id not in self.requests:
+                                continue
+                            self._merge_additional_information_update(req_id, upd)
+            except Exception as e:
+                logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}")
             if not self.broadcast_pp_output:
+
                 # Common case.
                 if not get_pp_group().is_last_rank:
                     # Return the intermediate tensors.
@@ -589,3 +620,26 @@ class GPUARModelRunner(OmniGPUModelRunner):
             invalid_req_indices=invalid_req_indices,
             async_output_copy_stream=self.async_output_copy_stream,
         )
+        
+    def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
+        req_state = self.requests.get(req_id)
+        if req_state is None:
+            return
+        existing = getattr(req_state, "additional_information_cpu", {})
+        if not isinstance(existing, dict):
+            existing = {}
+        merged = dict(existing)
+        for k, v in upd.items():
+            if isinstance(v, torch.Tensor):
+                merged[k] = v.detach().to("cpu").contiguous()
+            elif isinstance(v, list):
+                new_list = []
+                for item in v:
+                    if isinstance(item, torch.Tensor):
+                        new_list.append(item.detach().to("cpu").contiguous())
+                    else:
+                        new_list.append(item)
+                merged[k] = new_list
+            else:
+                merged[k] = v
+        setattr(req_state, "additional_information_cpu", merged)