Merge pull request #51 from tzhouam/feat/multi-request-stream-new

[Feature] Multi Request Stream for Sync Mode

docs/mrs_design.md

@@ -0,0 +1,188 @@
+## Multi-Request Streaming (MRS) on a Single Machine
+
+### 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.
+- 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: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects.
+  - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side.
+- 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
+  - A downstream stage triggers only after the upstream completes the request.
+  - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided.
+
+### 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: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability.
+
+### 8. Scheduling & Cancellation (simplified)
+- 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.
+- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue.
+
+#### 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` (Orchestrator)
+  - Class `OmniLLM` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers.
+  - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow.
+  - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled.
+  - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM.
+  - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions.
+
+### 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 (pseudo-code)
+```python
+def stage_worker(input_q, output_q, runtime, shm_threshold):
+    # Device binding
+    set_stage_gpu_devices(runtime.devices)
+    engine = OmniStageLLM(model, **engine_args)
+    batch_seq = 0
+    while True:
+        first = input_q.get()
+        if first is None:
+            break
+        # Batch intake (up to runtime.max_batch_size)
+        batch = [first]
+        while len(batch) < int(runtime.max_batch_size or 1) and not input_q.empty():
+            nxt = input_q.get()
+            if nxt is None:
+                input_q.put(None)
+                break
+            batch.append(nxt)
+        # Decode IPC payload and sampling params
+        request_ids = []
+        engine_inputs = []
+        rx_bytes = {}
+        rx_decode_ms = {}
+        for t in batch:
+            ein, rx = maybe_load_from_ipc_with_metrics(t, "engine_inputs", "engine_inputs_shm")
+            request_ids.append(t["request_id"])
+            engine_inputs.extend(ein if isinstance(ein, list) else [ein])
+            rx_bytes[t["request_id"]] = rx["rx_transfer_bytes"]
+            rx_decode_ms[t["request_id"]] = rx["rx_decode_time_ms"]
+        # Generate and dispatch (grouped by request_id)
+        batch_seq += 1
+        outputs = list(engine.generate(engine_inputs, batch[0]["sampling_params"], use_tqdm=False))
+        grouped = group_by_request_id(outputs, request_ids)
+        for rid in request_ids:
+            r_out = grouped.get(rid, [])
+            use_shm, payload = maybe_dump_to_shm(r_out, shm_threshold)
+            msg = {
+                "request_id": rid,
+                "metrics": {
+                    "batch_id": batch_seq,
+                    "rx_transfer_bytes": int(rx_bytes.get(rid, 0)),
+                    "rx_decode_time_ms": float(rx_decode_ms.get(rid, 0.0)),
+                    # omitted: generation latency and token stats
+                },
+            }
+            if use_shm:
+                msg["engine_outputs_shm"] = payload
+            else:
+                msg["engine_outputs"] = payload
+            output_q.put(msg)
+```
+
+### 12. Observability
+- Metrics (as implemented):
+  - per-request (emitted by stages): `num_tokens_out`, `stage_gen_time_ms`, `rx_transfer_bytes`, `rx_decode_time_ms`, `rx_in_flight_time_ms`, `batch_id`
+  - orchestrator aggregates: E2E latency, tokens/s (written only when stats are enabled)
+  - optional per-stage JSONL: `{log_file}.stage{stage_id}.stats.jsonl`
+- Logs/tracing:
+  - optional per-stage log files: `{log_file}.stage{stage_id}.log`
+  - the orchestrator can log readiness, forward size/time, and summary information
+
+### 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.
+
+### 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: Orchestrator created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there.
+- After: enhanced `OmniStage` owns its subprocess (`init_stage_worker`), with `_stage_worker` as the entry; device setup, LLM init, batching, and SHM live in `OmniStage`.
+
+Classes & responsibilities
+- OmniLLM (orchestrator)
+  - Build `OmniStage` list in parallel (preserve `process_engine_inputs` wiring)
+  - 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 `runtime`), `in_q/out_q`, subprocess handle, stats
+  - API:
+    - `init_stage_worker(...)`: spawn subprocess
+    - `stop_stage_worker()`: graceful exit
+    - `submit(ipc_payload)`: submit to input queue (may use shared memory)
+    - `try_collect() -> Optional[payload]`: non-blocking get from output queue
+    - `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 `runtime.max_batch_size-1` non-blocking; support `max_batch_size>=1`
+      - `maybe_load_from_ipc_with_metrics` per `engine_inputs` in batch
+      - `stage_engine.generate(batched_inputs, sampling_params)` (window `-1`)
+      - write per-request aggregated outputs to output queue after `maybe_dump_to_shm`
+      - log start/finish; report `error` per request on exceptions
+
+Data flow (sketch)
+```
+OmniLLM
+  ├─ 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)
+- `runtime.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
+- `runtime.max_batch_size`: max requests taken per intake (default 1)
+- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; injected into stages)
+- Stage link and inputs:
+  - `engine_input_source`: list of upstream stage_id(s) from which the downstream stage takes inputs (typically a single source)
+  - `custom_process_input_func`: function to construct downstream engine_inputs from upstream outputs
+- Output flags:
+  - `final_output`, `final_output_type`: marks that a stage produces the final external output and its type
+
+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
+
+Migration (execution order)
+1) Enhanced `OmniStage` with `init_stage_worker/stop_stage_worker/submit/try_collect/_stage_worker`.
+2) `OmniLLM` holds `OmniStage` instances, replacing the older pipelined class; uses Stage APIs.
+3) Reuse `entrypoints/stage_utils`: device/SHM/encode-decode utilities live here.
+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/README.md

@@ -63,13 +63,30 @@ uv pip install -r requirements.txt
 ```
 
 ## Run examples (Qwen2.5-omni)
-
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
+```bash
+tar -xf <Your Download Path>/seedtts_testset.tar
+cp seedtts_testset/en/meta.lst examples/offline_inference/qwen_2_5_omni/meta.lst
+python3 examples/offline_inference/qwen_2_5_omni/extract_prompts.py \
+  --input examples/offline_inference/qwen_2_5_omni/meta.lst \
+  --output examples/offline_inference/qwen_2_5_omni/top100.txt \
+  --topk 100
+```
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen_2_5_omni
 ```
-Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run.
+Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run.
 ```bash
-bash run.sh
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen_2_5_omni
+```
+Modify PYTHONPATH in run_single_prompt.sh as your path of vllm_omni. Then run.
+```bash
+bash run_single_prompt.sh
 ```
-The output audio is saved in ./output_audio

examples/offline_inference/qwen_2_5_omni/end2end.py

@@ -7,8 +7,8 @@ import numpy as np
 import soundfile as sf
 import torch
 from utils import make_omni_prompt
-
 from vllm.sampling_params import SamplingParams
+
 from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 _os_env_toggle.environ["VLLM_USE_V1"] = "1"
@@ -49,7 +49,7 @@ def parse_args():
         "--hf-revision", default=None, help="Optional HF revision (branch/tag/commit)."
     )
     parser.add_argument(
-        "--prompts", required=True, nargs="+", help="Input text prompts."
+        "--prompts", nargs="+", default=None, help="Input text prompts."
     )
     parser.add_argument(
         "--voice-type", default="default", help="Voice type, e.g., m02, f030, default."
@@ -69,6 +69,12 @@ def parse_args():
         "--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")
@@ -103,13 +109,46 @@ def parse_args():
     parser.add_argument("--use-torchvision", action="store_true")
     parser.add_argument("--tokenize", action="store_true")
     parser.add_argument(
-        "--output-wav", default="output.wav", help="Output wav file path."
+        "--output-wav", default="output.wav", help="[Deprecated] Output wav directory (use --output-dir)."
+    )
+    parser.add_argument(
+        "--output-dir", default="outputs", help="Output directory to save text and wav files together."
     )
     parser.add_argument(
         "--thinker-hidden-states-dir",
         default="thinker_hidden_states",
         help="Path to thinker hidden states directory.",
     )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
     args = parser.parse_args()
     return args
 
@@ -117,7 +156,28 @@ def parse_args():
 def main():
     args = parse_args()
     model_name = args.model
-    omni_llm = OmniLLM(model=model_name)
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, "r", encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            args.prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
+    omni_llm = OmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
     thinker_sampling_params = SamplingParams(
         temperature=0.0,  # Deterministic - no randomness
         top_p=1.0,  # Disable nucleus sampling
@@ -128,13 +188,13 @@ def main():
         repetition_penalty=1.1,
     )
     talker_sampling_params = SamplingParams(
-        temperature=0.0,  # Deterministic - no randomness
-        top_p=1.0,  # Disable nucleus sampling
-        top_k=-1,  # Disable top-k sampling
+        temperature=0.9,
+        top_p=0.8,
+        top_k=40,
         max_tokens=2048,
         seed=SEED,  # Fixed seed for sampling
         detokenize=True,
-        repetition_penalty=1.1,
+        repetition_penalty=1.05,
         stop_token_ids=[8294],
     )
     code2wav_sampling_params = SamplingParams(
@@ -156,20 +216,33 @@ def main():
     prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
     omni_outputs = omni_llm.generate(prompt, sampling_params_list)
 
-    os.makedirs(args.output_wav, exist_ok=True)
+    # 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 = output.request_id
+                request_id = int(output.request_id)
                 text_output = output.outputs[0].text
-                print(f"Request ID: {request_id}, Text Output: {text_output}")
+                # Save aligned text file per request
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
-                request_id = output.request_id
+                request_id = int(output.request_id)
                 audio_tensor = output.multimodal_output["audio"]
-                output_wav = os.path.join(
-                    args.output_wav, f"output_{output.request_id}.wav"
-                )
+                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
                 sf.write(
                     output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000
                 )

examples/offline_inference/qwen_2_5_omni/extract_prompts.py

@@ -0,0 +1,43 @@
+#!/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="Input .lst file path")
+    parser.add_argument("--output", "-o", required=True, help="Output file path")
+    parser.add_argument("--topk", "-k", type=int, default=100, help="Extract the top K prompts (default: 100)")
+    args = parser.parse_args()
+
+    prompts = []
+    with open(args.input, "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")
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh

@@ -0,0 +1,9 @@
+export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --txt-prompts top100.txt

examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh

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

examples/offline_inference/qwen_2_5_omni/utils.py

@@ -13,6 +13,10 @@ from processing_omni import fetch_image, fetch_video
 from vllm.inputs import TextPrompt
 from vllm_omni.inputs.data import OmniTokensPrompt
 
+# Simple caches to avoid repeated heavy HF loads per prompt
+_PROCESSOR_CACHE: Dict[str, "AutoProcessor"] = {}
+_CONFIG_CACHE: Dict[str, "AutoConfig"] = {}
+
 
 def get_system_prompt():
 
@@ -113,16 +117,25 @@ def make_inputs_qwen2_omni(
 
     from transformers import AutoConfig, AutoProcessor
 
-    processor = AutoProcessor.from_pretrained(args.model)
+    # Cached processor/config to prevent per-prompt reloading and repeated warnings
+    if args.model not in _PROCESSOR_CACHE:
+        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
+    processor = _PROCESSOR_CACHE[args.model]
 
-    try:
-        config = AutoConfig.from_pretrained(args.model)
-        if "Qwen2_5OmniModel" in config.architectures:
-            args.legacy_omni_video = False
+    config = _CONFIG_CACHE.get(args.model)
+    if config is None:
+        try:
+            config = AutoConfig.from_pretrained(args.model)
+        except Exception:
+            config = None
+        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
+
+    # Decide legacy flag only once based on config (default True if unknown)
+    if getattr(args, "legacy_omni_video", None) is None:
+        if config is not None and hasattr(config, "architectures"):
+            args.legacy_omni_video = not ("Qwen2_5OmniModel" in config.architectures)
         else:
             args.legacy_omni_video = True
-    except Exception:
-        args.legacy_omni_video = True
 
     audios, images, videos = [], [], []
     for message in messages:

vllm_omni/entrypoints/log_utils.py

@@ -0,0 +1,570 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import logging
+import os
+import time
+
+from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl
+
+
+def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
+    try:
+        if not log_file:
+            return
+        try:
+            if os.path.exists(log_file):
+                os.remove(log_file)
+        except Exception:
+            pass
+        # Per-stage logs and stats
+        for sid in range(num_stages):
+            try:
+                p = f"{log_file}.stage{sid}.log"
+                if os.path.exists(p):
+                    os.remove(p)
+            except Exception:
+                pass
+            try:
+                p = f"{log_file}.stage{sid}.stats.jsonl"
+                if os.path.exists(p):
+                    os.remove(p)
+            except Exception:
+                pass
+        # Orchestrator stats files
+        try:
+            p = f"{log_file}.orchestrator.stats.jsonl"
+            if os.path.exists(p):
+                os.remove(p)
+        except Exception:
+            pass
+        try:
+            p = f"{log_file}.overall.stats.jsonl"
+            if os.path.exists(p):
+                os.remove(p)
+        except Exception:
+            pass
+    except Exception:
+        pass
+
+
+def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None:
+    try:
+        if not log_file:
+            return
+        has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
+        if not has_file_handler:
+            fh = logging.FileHandler(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
+
+
+def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optional[str], Optional[str]]:
+    stats_file: Optional[str] = None
+    overall_file: Optional[str] = None
+    try:
+        if enable_stats and log_file:
+            stats_file = f"{log_file}.orchestrator.stats.jsonl"
+            overall_file = f"{log_file}.overall.stats.jsonl"
+    except Exception:
+        stats_file = None
+        overall_file = None
+    return stats_file, overall_file
+
+
+def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
+    if not path:
+        return
+    try:
+        _append_jsonl(path, record)  # type: ignore[arg-type]
+    except Exception:
+        pass
+
+
+def log_transfer_tx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
+                    size_bytes: int, tx_time_ms: float, used_shm: bool) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "transfer_stats",
+        "from_stage": from_stage,
+        "to_stage": to_stage,
+        "request_id": request_id,
+        "size_bytes": int(size_bytes),
+        "tx_time_ms": float(tx_time_ms),
+        "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
+        "used_shm": bool(used_shm),
+    })
+
+
+def log_transfer_rx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
+                    rx_bytes: int, rx_decode_time_ms: float, in_flight_time_ms: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "transfer_rx_stats",
+        "from_stage": from_stage,
+        "to_stage": to_stage,
+        "request_id": request_id,
+        "rx_bytes": int(rx_bytes),
+        "rx_decode_time_ms": float(rx_decode_time_ms),
+        "in_flight_time_ms": float(in_flight_time_ms),
+        "rx_time_per_kb_ms": (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0,
+    })
+
+
+def log_transfer_total(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
+                       size_bytes: int, tx_time_ms: float, in_flight_time_ms: float,
+                       rx_decode_time_ms: float, total_time_ms: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "transfer_total_stats",
+        "from_stage": from_stage,
+        "to_stage": to_stage,
+        "request_id": request_id,
+        "size_bytes": int(size_bytes),
+        "tx_time_ms": float(tx_time_ms),
+        "in_flight_time_ms": float(in_flight_time_ms),
+        "rx_decode_time_ms": float(rx_decode_time_ms),
+        "total_time_ms": float(total_time_ms),
+        "total_time_per_kb_ms": float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0,
+    })
+
+
+def log_orchestrator_e2e(stats_file: Optional[str], request_id: Any, final_stage_id: int,
+                         e2e_time_ms: float, num_tokens_out: int, e2e_time_per_token_ms: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "orchestrator_request_e2e",
+        "request_id": request_id,
+        "final_stage_id": final_stage_id,
+        "e2e_time_ms": float(e2e_time_ms),
+        "num_tokens_out": int(num_tokens_out),
+        "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
+    })
+
+
+def log_orchestrator_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+    _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
+
+
+def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+    _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
+
+
+def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None:
+    _safe_append_jsonl(overall_stats_file, record)
+
+
+def log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, batch_size: int,
+                            num_tokens_out: int, stage_gen_time_ms: float, tokens_per_s: float,
+                            rx_transfer_bytes: int, rx_decode_time_ms: float, rx_mbps: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "stage_request_stats",
+        "stage_id": stage_id,
+        "request_id": request_id,
+        "batch_size": int(batch_size),
+        "num_tokens_out": int(num_tokens_out),
+        "stage_gen_time_ms": float(stage_gen_time_ms),
+        "tokens_per_s": float(tokens_per_s),
+        "rx_transfer_bytes": int(rx_transfer_bytes),
+        "rx_decode_time_ms": float(rx_decode_time_ms),
+        "rx_mbps": float(rx_mbps),
+    })
+
+
+def log_stage_running_avg(stats_file: Optional[str], stage_id: int, total_tokens: int,
+                          total_gen_time_ms: float, avg_tokens_per_s: float) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "stage_running_avg",
+        "stage_id": stage_id,
+        "total_tokens": int(total_tokens),
+        "total_gen_time_ms": float(total_gen_time_ms),
+        "avg_tokens_per_s": float(avg_tokens_per_s),
+    })
+
+
+def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: int,
+                          batch_gen_time_ms: float, request_ids: List[Any]) -> None:
+    _safe_append_jsonl(stats_file, {
+        "type": "stage_batch_stats",
+        "stage_id": stage_id,
+        "batch_size": int(batch_size),
+        "batch_gen_time_ms": float(batch_gen_time_ms),
+        "request_ids": list(request_ids),
+    })
+
+
+def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any,
+                                       batch_size: int, engine_outputs: List[Any],
+                                       stage_gen_time_ms: float, rx_transfer_bytes: int,
+                                       rx_decode_time_ms: float) -> None:
+    """Compute per-request metrics and log them in one call."""
+    num_tokens = count_tokens_from_outputs(engine_outputs)
+    tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
+    rx_mbps = (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) if rx_transfer_bytes > 0 else 0.0
+    log_stage_request_stats(
+        stats_file,
+        stage_id,
+        request_id,
+        int(batch_size),
+        int(num_tokens),
+        float(stage_gen_time_ms),
+        float(tokens_per_s),
+        int(rx_transfer_bytes),
+        float(rx_decode_time_ms),
+        float(rx_mbps),
+    )
+
+
+# ----------------- Aggregation helpers for orchestrator -----------------
+
+def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
+                         stage_req_counts: List[int],
+                         stage_total_time_ms: List[float],
+                         stage_total_tokens: List[int],
+                         stage_id: int,
+                         req_id: Any,
+                         metrics: Dict[str, Any]) -> None:
+    try:
+        stage_req_counts[stage_id] += 1
+        stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
+        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(metrics.get("stage_gen_time_ms", 0.0)),
+            "num_tokens_out": int(metrics.get("num_tokens_out", 0)),
+        }
+    except Exception:
+        pass
+
+
+def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
+                                 transfer_agg: Dict[tuple[int, int], Dict[str, float]],
+                                 per_request: Dict[int, Dict[str, Any]],
+                                 stage_id: int,
+                                 req_id: Any,
+                                 rx_bytes: float,
+                                 rx_ms: float,
+                                 in_flight_ms: float) -> Optional[tuple[int, float, float]]:
+    try:
+        # Update RX aggregates for (stage_id-1 -> stage_id)
+        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
+            agg["sum_rx_bytes"] += float(rx_bytes)
+            agg["sum_rx_ms"] += float(rx_ms)
+            agg["rx_count"] += 1.0
+
+            # Try combine with sender-side timing if present
+            try:
+                rid_int = int(req_id)
+            except Exception:
+                rid_int = req_id  # best effort
+            s = transfer_edge_req.get((stage_id - 1, stage_id, rid_int))
+            if s is None:
+                return None
+            tx_ms = float(s.get("tx_ms", 0.0))
+            size_b = float(s.get("size_bytes", rx_bytes))
+            total_ms = tx_ms + float(in_flight_ms) + float(rx_ms)
+            agg["sum_total_ms"] += total_ms
+            agg["total_count"] += 1.0
+            # accumulate per-request transfer totals
+            try:
+                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_bytes)  # type: ignore[index]
+            except Exception:
+                pass
+            return int(size_b), float(tx_ms), float(total_ms)
+        return None
+    except Exception:
+        return None
+
+
+def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, float]],
+                               transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
+                               from_stage: int,
+                               to_stage: int,
+                               req_id: Any,
+                               size_bytes: int,
+                               tx_ms: float) -> None:
+    try:
+        key = (from_stage, to_stage)
+        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
+        try:
+            rid_int = int(req_id)
+        except Exception:
+            rid_int = req_id
+        transfer_edge_req[(from_stage, to_stage, rid_int)] = {
+            "tx_ms": float(tx_ms),
+            "size_bytes": float(size_bytes),
+        }
+    except Exception:
+        pass
+
+
+def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
+    total = 0
+    for _ro in engine_outputs:
+        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
+
+
+def build_stage_summary(stage_req_counts: List[int],
+                        stage_total_tokens: List[int],
+                        stage_total_time_ms: List[float]) -> List[Dict[str, Any]]:
+    summary: List[Dict[str, Any]] = []
+    for sid in range(len(stage_req_counts)):
+        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
+        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,
+        })
+    return summary
+
+
+def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]) -> List[Dict[str, Any]]:
+    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
+        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,
+        })
+    return summary
+
+
+class OrchestratorMetrics:
+    def __init__(self, num_stages: int, enable_stats: bool,
+                 stats_file: Optional[str], overall_stats_file: Optional[str],
+                 wall_start_ts: float) -> None:
+        self.num_stages = int(num_stages)
+        self.enable_stats = bool(enable_stats)
+        self.stats_file = stats_file
+        self.overall_stats_file = overall_stats_file
+        self.stage_total_time_ms: List[float] = [0.0 for _ in range(self.num_stages)]
+        self.stage_total_tokens: List[int] = [0 for _ in range(self.num_stages)]
+        self.stage_req_counts: List[int] = [0 for _ in range(self.num_stages)]
+        self.transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {}
+        self.transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {}
+        self.e2e_total_ms: float = 0.0
+        self.e2e_total_tokens: int = 0
+        self.e2e_count: int = 0
+        self.e2e_done: set[int] = set()
+        self.per_request: Dict[int, Dict[str, Any]] = {}
+        self.sum_per_request_transfer_ms: float = 0.0
+        self.wall_start_ts: float = float(wall_start_ts)
+        self.last_finish_ts: float = float(wall_start_ts)
+        self.stage_seen_batches: Dict[int, set] = {sid: set() for sid in range(self.num_stages)}
+        self.stage_first_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
+        self.stage_last_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
+
+    def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: Dict[str, Any]) -> None:
+        record_stage_metrics(
+            self.per_request,
+            self.stage_req_counts,
+            self.stage_total_time_ms,
+            self.stage_total_tokens,
+            stage_id,
+            req_id,
+            metrics,
+        )
+        try:
+            batch_id_raw = metrics.get("batch_id", None)
+            if batch_id_raw is not None:
+                batch_id = int(batch_id_raw)
+                if batch_id not in self.stage_seen_batches[stage_id]:
+                    self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0))
+                    self.stage_seen_batches[stage_id].add(batch_id)
+        except Exception:
+            pass
+        rx_b = float(metrics.get("rx_transfer_bytes", 0.0))
+        rx_ms = float(metrics.get("rx_decode_time_ms", 0.0))
+        in_flight_ms = float(metrics.get("rx_in_flight_time_ms", 0.0))
+        combined = aggregate_rx_and_maybe_total(
+            self.transfer_edge_req,
+            self.transfer_agg,
+            self.per_request,
+            stage_id,
+            req_id,
+            rx_b,
+            rx_ms,
+            in_flight_ms,
+        )
+        if self.enable_stats and self.stats_file and stage_id > 0:
+            log_transfer_rx(
+                self.stats_file,
+                stage_id - 1,
+                stage_id,
+                req_id,
+                int(rx_b),
+                float(rx_ms),
+                float(in_flight_ms),
+            )
+            if combined is not None:
+                size_b_c, tx_ms_c, total_ms_c = combined
+                log_transfer_total(
+                    self.stats_file,
+                    stage_id - 1,
+                    stage_id,
+                    req_id,
+                    int(size_b_c),
+                    float(tx_ms_c),
+                    float(in_flight_ms),
+                    float(rx_ms),
+                    float(total_ms_c),
+                )
+
+    def on_forward(self, from_stage: int, to_stage: int, req_id: Any,
+                   size_bytes: int, tx_ms: float, used_shm: bool) -> None:
+        # Mark first input time for the destination stage if not set
+        if self.stage_first_ts[to_stage] is None:
+            self.stage_first_ts[to_stage] = time.time()
+        if self.enable_stats and self.stats_file:
+            log_transfer_tx(
+                self.stats_file,
+                from_stage,
+                to_stage,
+                req_id,
+                int(size_bytes),
+                float(tx_ms),
+                bool(used_shm),
+            )
+        record_sender_transfer_agg(
+            self.transfer_agg,
+            self.transfer_edge_req,
+            from_stage,
+            to_stage,
+            req_id,
+            int(size_bytes),
+            float(tx_ms),
+        )
+
+    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: List[Any],
+                            req_start_ts: float) -> None:
+        try:
+            rid_int = int(req_id)
+        except Exception:
+            rid_int = req_id
+        _t0 = float(req_start_ts)
+        _t1 = time.time()
+        # Update last output time for this stage
+        prev_last = self.stage_last_ts[stage_id]
+        self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1)
+        self.last_finish_ts = max(self.last_finish_ts, _t1)
+        e2e_ms = (_t1 - _t0) * 1000.0
+        num_tokens = count_tokens_from_outputs(engine_outputs)
+        self.e2e_total_ms += e2e_ms
+        self.e2e_total_tokens += int(num_tokens)
+        self.e2e_count += 1
+        self.e2e_done.add(rid_int)
+        pr = self.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", {}),
+        }
+        self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
+        if self.enable_stats and self.overall_stats_file:
+            log_overall_record(self.overall_stats_file, per_req_record)
+        if self.enable_stats and self.stats_file:
+            e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
+            log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
+
+    def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]:
+        # Compute stage summary using wall time between first input and last output per stage
+        stage_summary: List[Dict[str, Any]] = []
+        for sid in range(self.num_stages):
+            first_ts = self.stage_first_ts[sid]
+            last_ts = self.stage_last_ts[sid]
+            total_ms = (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
+            reqs = self.stage_req_counts[sid]
+            tokens = self.stage_total_tokens[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": float(total_ms),
+                "avg_time_per_request_ms": float(avg_req),
+                "avg_tokens_per_s": float(avg_tok),
+            })
+        transfer_summary = build_transfer_summary(self.transfer_agg)
+        e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
+        e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
+        wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
+        summary: Dict[str, Any] = {
+            "e2e_requests": int(self.e2e_count),
+            "e2e_total_time_ms": float(wall_time_ms),
+            "e2e_sum_time_ms": float(self.e2e_total_ms),
+            "e2e_total_tokens": int(self.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,
+        }
+        if self.enable_stats and self.stats_file:
+            log_orchestrator_summary(self.stats_file, summary)
+        if self.enable_stats and self.overall_stats_file:
+            log_overall_summary(self.overall_stats_file, summary)
+        return summary
+
+

vllm_omni/entrypoints/omni_llm.py

@@ -1,8 +1,16 @@
 from typing import Any, Optional, Sequence, Union
+import logging
+import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
+import queue
+import sys
+import time
 
 import cloudpickle
 from pydantic import ValidationError
 
+# External library imports (vLLM)
 from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
 from vllm.engine.arg_utils import HfOverrides
 from vllm.entrypoints.llm import LLM
@@ -13,37 +21,128 @@ from vllm.sampling_params import SamplingParams
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
+
+# Internal imports (our code)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.utils import load_stage_configs_from_model
+from vllm_omni.entrypoints.log_utils import (
+    remove_old_logs,
+    configure_orchestrator_logger,
+    init_stats_paths,
+    OrchestratorMetrics,
+)
+from vllm_omni.entrypoints.stage_utils import (
+    maybe_load_from_ipc as _load,
+    encode_for_ipc as _encode,
+    serialize_obj as _ser,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
 
 class OmniLLM:
-    def __init__(
-        self, model: str, stage_configs=None, log_stats: bool = False, **kwargs
-    ):
+    """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 = False,
+                 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.initialize_stage_configs(model)
+            self.stage_configs = load_stage_configs_from_model(model)
         else:
             self.stage_configs = stage_configs
 
-        self.stage_list = []
-        self.initialize_stages(model)
 
-    def initialize_stage_configs(self, model: str):
-        self.stage_configs = load_stage_configs_from_model(model)
+        # Optional file handler for orchestrator
+        self._log_file = log_file
+        if self._log_file:
+            remove_old_logs(self._log_file, len(self.stage_list))
+            configure_orchestrator_logger(logger, self._log_file)
 
-    def initialize_stages(self, model: str):
-        for stage_config in self.stage_configs:
-            stage = OmniStage(stage_config)
-            stage_llm = OmniStageLLM(model=model, **stage_config.engine_args)
-            stage.set_engine(stage_llm)
-            self.stage_list.append(stage)
+        self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
+        self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
+
+    def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None:
+        self.stage_list: list[OmniStage] = []
+        # Build OmniStage instances in parallel, preserve original order
+        def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
+            idx, cfg = idx_cfg
+            return idx, OmniStage(cfg)
+
+        with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
+            futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
+            results: list[tuple[int, OmniStage]] = []
+            for fut in as_completed(futures):
+                results.append(fut.result())
+        results.sort(key=lambda x: x[0])
+        self.stage_list = [st for _, st in results]
+        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] = []
+        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)
+            time.sleep(self._init_sleep_seconds)
+
+    def close(self) -> None:
+        for q in self._stage_in_queues:
+            try:
+                q.put_nowait(None)
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e)
+        for stage in self.stage_list:
+            try:
+                stage.stop_stage_worker()
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
+
+    def __del__(self) -> None:  # best-effort
+        try:
+            self.close()
+        except Exception as e:
+            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
 
     def generate(
         self,
@@ -52,42 +151,233 @@ class OmniLLM:
             Union[SamplingParams, Sequence[SamplingParams]]
         ] = None,
     ) -> list[OmniRequestOutput]:
-        """Generate text outputs for the given prompts."""
-        final_outputs: 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, "
-                f"got {len(sampling_params_list)}"
+                f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}"
             )
-        for stage_id, stage in enumerate(self.stage_list):
-            if stage_id > 0:
-                engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)
-            else:
-                engine_inputs = prompts
-            engine_outputs = self._run_generation(
-                stage, sampling_params_list[stage_id], engine_inputs
-            )
-            stage.set_engine_outputs(engine_outputs)
-            if hasattr(stage, "final_output") and stage.final_output:
-                final_outputs.append(
-                    OmniRequestOutput(
-                        stage_id=stage_id,
-                        final_output_type=stage.final_output_type,
-                        request_output=engine_outputs,
+
+        # 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 as e:
+            logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True)
+            final_stage_id_for_e2e = len(self.stage_list) - 1
+        # Metrics/aggregation helper
+        metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts)
+
+        # Seed stage-0 queue with all requests
+        logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
+        # Mark first input time for stage-0
+        metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
+
+        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")
+                # Mark last output time for this stage whenever we receive outputs
+                metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
+                try:
+                    _m = result.get("metrics")
+                    if _m is not None:
+                        metrics.on_stage_metrics(stage_id, req_id, _m)
+                except Exception as e:
+                    logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e)
+                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:
+                        rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
+                        if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
+                            metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts))
+                    except Exception as e:
+                        logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e)
+
+                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
+                        metrics.on_forward(stage_id, next_stage_id, req_id, int(size_bytes), float(tx_ms), bool("engine_inputs_shm" in ipc_payload))
+                    except Exception as e:
+                        logger.warning("[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", req_id, e)
+                        self.stage_list[next_stage_id].submit({
+                            "request_id": req_id,
+                            "engine_inputs": next_inputs,
+                            "sampling_params": sp_next,
+                        })
+                    logger.debug("[Orchestrator] Forwarded request %s to stage-%s", req_id, 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:
+            summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
+            logger.info("[Summary] %s", summary)
+        except Exception as e:
+            logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
+
         return final_outputs
 
-    def _run_generation(
-        self,
-        stage: OmniStage,
-        sampling_params: SamplingParams,
-        prompts: Union[PromptType, Sequence[PromptType]],
-    ):
-        engine_outputs = []
-        for ro in stage.engine.generate(prompts, sampling_params):
-            engine_outputs.append(ro)
-        return engine_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 (runtime.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)
 
 
 class OmniStageLLM(LLM):

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.stage_utils import (
+    _to_dict,
+    set_stage_gpu_devices,
+    maybe_load_from_ipc_with_metrics,
+    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)
+        runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
+        stage_payload: Dict[str, Any] = {
+            "stage_id": self.stage_id,
+            "engine_args": engine_args,
+            "runtime": runtime_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 as e:
+                self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
+        if self._proc is not None:
+            try:
+                self._proc.join(timeout=5)
+            except Exception as e:
+                self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True)
+            if self._proc.is_alive():
+                try:
+                    self._proc.terminate()
+                except Exception as e:
+                    self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
+
+    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,224 @@ 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
+    from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
+        log_stage_running_avg,
+        log_stage_batch_stats,
+        compute_and_log_stage_request_stats,
+        count_tokens_from_outputs,
+    )
+    import queue as _queue
+    import time as _time
+    # no inline JSONL/serialization imports; logging handled by utilities
+
+    stage_id = stage_payload["stage_id"]
+    engine_args = stage_payload.get("engine_args", {})
+    runtime_cfg = stage_payload.get("runtime", {})
+    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
+    # Monotonic batch id per stage process for orchestrator dedup on time aggregation
+    _batch_seq = 0
+
+    # Device mapping
+    try:
+        set_stage_gpu_devices(stage_id, runtime_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(runtime_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
+            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(
+                t, obj_key="engine_inputs", shm_key="engine_inputs_shm"
+            )
+            _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+            batch_request_ids.append(rid)
+            if isinstance(ein, list):
+                batch_engine_inputs.extend(ein)
+            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:
+            _batch_seq += 1
+            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
+            for rid in batch_request_ids:
+                _r_outputs = req_to_outputs.get(rid, [])
+                _num_tokens = count_tokens_from_outputs(_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:
+                _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
+                log_stage_running_avg(_stats_file, stage_id, int(_agg_total_tokens), float(_agg_total_gen_time_ms), float(_avg_tokens_per_s))
+                log_stage_batch_stats(_stats_file, stage_id, len(batch_tasks), float(_gen_ms), list(batch_request_ids))
+
+            # 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_from_outputs(r_outputs)),
+                        "stage_gen_time_ms": _gen_ms,
+                        "batch_id": int(_batch_seq),
+                        "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 _stats_file:
+                        compute_and_log_stage_request_stats(
+                            _stats_file,
+                            stage_id,
+                            rid,
+                            len(batch_tasks),
+                            r_outputs,
+                            float(_gen_ms),
+                            int(_metrics["rx_transfer_bytes"]),   # type: ignore[index]
+                            float(_metrics["rx_decode_time_ms"]), # type: ignore[index]
+                        )
+                    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_from_outputs(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/stage_utils.py

@@ -0,0 +1,239 @@
+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 as e:
+                    logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
+                    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 as e:
+                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
+                    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 as e:
+                    logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True)
+                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 as e:
+            logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True)
+    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 as e:
+        logger.debug("Ensure parent dir failed for %s: %s", path, e)
+    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.
+
+    Deprecated: prefer `maybe_load_from_ipc_with_metrics` to also obtain
+    decode-time and size metrics.
+    """
+    if shm_key in container:
+        return pickle.loads(shm_read_bytes(container[shm_key]))
+    return container[obj_key]
+
+
+def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, shm_key: str) -> tuple[Any, Dict[str, float]]:
+    """Load object and return (object, metrics) with RX bytes and decode time.
+
+    Metrics keys:
+      - rx_transfer_bytes: int
+      - rx_decode_time_ms: float
+    """
+    import time as _time  # local import to avoid overhead at module import
+    t0 = _time.time()
+    if shm_key in container:
+        meta = container[shm_key]  # type: ignore[index]
+        payload = shm_read_bytes(meta)
+        obj = pickle.loads(payload)
+        try:
+            rx_bytes = int(meta.get("size", len(payload)))  # type: ignore[call-arg]
+        except Exception:
+            rx_bytes = len(payload)
+    else:
+        obj = container[obj_key]
+        try:
+            rx_bytes = len(serialize_obj(obj))
+        except Exception:
+            rx_bytes = 0
+    t1 = _time.time()
+    rx_decode_ms = (t1 - t0) * 1000.0
+    return obj, {
+        "rx_transfer_bytes": int(rx_bytes),
+        "rx_decode_time_ms": float(rx_decode_ms),
+    }
+
+
+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/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 a single step, both Prefill*n and Decode*n are supported.
+            # Rules:
+            # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
+            # - Decode segments consist of a single non-special token.
+            # - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes),
+            #   then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
+            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding.
 
             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
+    runtime:
+      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
+    runtime:
+      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
+    runtime:
+      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 runtime config (concise): default windows and stage edges
+runtime:
+  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

@@ -65,6 +65,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
         torch.Tensor,
         Optional[IntermediateTensors],
         dict[str, Any],
+        Optional[dict[str, dict]],
     ]:
 
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
@@ -79,6 +80,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         # _prepare_inputs may reorder the batch, so we must gather multi
         # modal outputs after that to ensure the correct order
+        per_req_additional_information: Optional[dict[str, dict]] = None
         if (
             self.supports_mm_inputs
             and get_pp_group().is_first_rank
@@ -99,9 +101,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 = {}
 
             # Overlay custom prompt_embeds per request for the prompt portion;
             # collect additional_information (tensor/list) for prefill portion only
@@ -122,54 +126,31 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     src = pe_cpu[
                         num_computed_tokens : num_computed_tokens + overlay_len
                     ].to(dtype=self.dtype, device=self.device, non_blocking=True)
-                    start_offset = int(self.query_start_loc_cpu[req_index])
+                    start_offset = int(self.query_start_loc.cpu[req_index])
                     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]
@@ -240,6 +221,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             positions,
             intermediate_tensors,
             model_kwargs,
+            per_req_additional_information,
         )
 
     @torch.inference_mode()
@@ -355,6 +337,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 positions,
                 intermediate_tensors,
                 model_kwargs,
+                per_req_additional_information,
             ) = self._preprocess(
                 scheduler_output,
                 num_scheduled_tokens_np,
@@ -395,15 +378,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
+            # Always pass per-request runtime additional_information (persisted in 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
+                # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
+                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 +429,42 @@ class GPUARModelRunner(OmniGPUModelRunner):
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
                 hidden_states
             )
-
+            # The model side may return per-request additional_information updates (model-agnostic channel).
+            # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
+            # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
+            try:
+                if (
+                    isinstance(multimodal_outputs, dict)
+                    and (
+                        "additional_information_update" in multimodal_outputs
+                        or "additional_information_update_by_req_id" in multimodal_outputs
+                    )
+                ):
+                    # Option A: list[dict] in batch order
+                    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)
+                    # Option B: dict[str, dict] keyed by req_id
+                    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 +624,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)