[cpu] remove verbose print

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

docs/mrs_design.md

@@ -1,8 +1,4 @@
-## Multi-Request Streaming (MRS) on a Single Machine — Design Doc
+## Multi-Request Streaming (MRS) on a Single Machine
-
-- Author: Team
-- Date: 2025-10-27
-- Version: v0.3 (simplified: copy-based IPC / serial per stage / full-trigger)
 
 ### 1. Background & Scope
 - All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved.
@@ -12,77 +8,30 @@
 ### 2. Key Constraints
 - Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`).
 - Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version.
-- Full-trigger across stages: unify `window_size = -1` for all edges — downstream starts only after upstream completes for that request.
-- Serial within a stage: at most one in-flight request per stage (max_inflight=1). No batching/concurrency inside a stage in this version.
 - Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0).
 
 ### 3. Architecture Overview
 - Processes & IPC queues
   - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue.
-  - Inter-stage connection via IPC: control plane (Unix Domain Socket / Pipe) + data plane (shared memory ring buffer in future; copy-based in this version).
+  - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects.
-  - MPMC queue semantics; control plane passes light-weight pointers/metadata; data plane carries large payloads.
+  - 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 & windows
+- Scheduling
-  - Each edge (stage→stage) defines a window policy; downstream triggers compute by consuming stitched window views.
+  - A downstream stage triggers only after the upstream completes the request.
-  - `-1` mode triggers only after upstream emits `is_last=True` (full aggregation).
+  - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided.
-  - Backpressure propagates upstream via the control plane: once high-watermark is reached, upstream blocks or degrades behavior.
-
-### 4. Stage Breakdown
-- AR (autoregressive)
-  - S0 Intake/Plan: parse request, assign request_id, output normalized tasks.
-  - S1 Prefill: tokenize/encode/KV warmup; may micro-batch across requests.
-  - S2 Decode Loop: step-by-step tokens; flush segments by `max_segment_tokens` and `min_flush_interval_ms`.
-  - S3 Postprocess: detokenize, logprobs/struct wrapping (CPU if necessary).
-  - S4 Consumer: local consumer callback/buffering.
-- DiT/Diffusion
-  - D0 Intake/Plan: scheduler/seed init.
-  - D1 Sampling Loop: preview frames every k steps (downsample allowed).
-  - D2 Decode/Encode: VAE decode and JPEG/WEBP encode (prefer GPU→pinned CPU).
-  - D3 Postprocess/Consumer.
-
-### 5. Unified Events & Segment Structure (in-process)
-- Base event: `{ request_id, sequence, event_type, ts, payload, device_handles }`
-- ARSegmentEvent (segment):
-  - Fields: `tokens` (List[int] or GPU tensor view), optional `logprobs`, `is_last`, optional `hidden_states/aux` refs.
-  - Semantics: strictly ordered by `sequence` for the same `request_id`; `is_last` marks the final segment.
-- ImageProgressEvent: `{ step, image_tensor|jpeg_ref, is_last }` (diffusion preview/completion).
-- Completed/Error: terminal or error events affect only the current request.
-
-### 6. AR Segmentation & Windowing (simplified)
-- Production (upstream stage)
-  - Accumulate tokens during decoding; flush a segment if any of the following holds:
-    - Reaches `max_segment_tokens` length
-    - Exceeds `min_flush_interval_ms` since last flush
-    - Request ended (`is_last=True`)
-  - Flush immediately via `output_queue.put(segment)`; do not wait for entire request completion.
-- Consumption (downstream stage)
-  - Use unified `window_size = -1`: trigger only after receiving `is_last=True`; no window slicing.
-- Backpressure & watermark
-  - Must-deliver events (token.delta, Completed) are not dropped; block upstream briefly on high watermark.
-  - For long sequences under `-1`, allow early segments to spill from GPU to pinned CPU to control VRAM pressure.
-
-#### Example config (YAML, simplified)
-```yaml
-mrs:
-  enabled: true
-  defaults:
-    window_size: -1           # Full-trigger across stages
-    max_inflight: 1           # Serial per stage
-```
 
 ### 7. IPC Implementation (simplified: copy-based)
 - Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data).
 - Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version.
-- Backpressure: blocking reads/writes based on queue capacity, blocking upstream on high watermark.
+- 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)
-- Serial FIFO: each stage `max_inflight=1`; no cross-request micro-batching or concurrency.
-- Cancel/timeout: ongoing work should exit early; remaining queued requests stay pending.
 - Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency.
+- 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
@@ -90,121 +39,118 @@ mrs:
 3) t2: B handles req0 while A handles req1 (parallel across stages)
 
 ### 9. Integration Points (by file)
-- `vllm_omni/entrypoints/omni_llm.py`
+- `vllm_omni/entrypoints/omni_llm.py` (Orchestrator)
-  - Introduce a ProcessOrchestrator (replace threaded version):
+  - Class `OmniLLM` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers.
-    - Spawn stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), create control/data channels, build window state machines.
+  - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow.
-    - Add `enable_mrs` in `generate`; when enabled, do not collect `stage.engine.generate(...)` into lists; stream segments to downstream via IPC instead.
+  - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled.
-  - Register `window_size` per edge; `-1` triggers downstream once `is_last` is observed.
+  - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM.
-  - Manage process lifecycle: start/heartbeat/graceful shutdown/restart on crash; propagate cancel/timeout.
+  - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions.
-- `vllm_omni/outputs.py`
-  - Add in-process events: `ARSegmentEvent`, `ImageProgressEvent`, `CompletedEvent`, `ErrorEvent`.
-  - Keep `OmniRequestOutput` as the external representation of final stage outputs.
-- `vllm_omni/worker/gpu_ar_model_runner.py`
-  - Near `valid_sampled_token_ids` parsing (after updating `req_state.output_token_ids`), aggregate new tokens into segments and emit; send final `is_last=True` and Completed.
-  - Replace current `.to("cpu").contiguous()` sync path: by default keep GPU refs and let downstream copy asynchronously from a pinned pool when needed.
-- `vllm_omni/worker/gpu_diffusion_model_runner.py`
 
 ### 9.1 Process Device Visibility
 - Device binding: set `CUDA_VISIBLE_DEVICES` before process start, or call `torch.cuda.set_device` early in init; configs can be single/multi-GPU.
 - Cross-device transfer (simplified): via CPU copies; zero-copy paths are out of scope for now.
 
-### 10. Worker Template & Window State (pseudo-code)
+### 10. Worker Template (pseudo-code)
 ```python
-def stage_worker(input_q, output_q, ctx):
+def stage_worker(input_q, output_q, runtime, shm_threshold):
-    while ctx.alive:
+    # Device binding
-        item = input_q.get()  # blocking / with timeout
+    set_stage_gpu_devices(runtime.devices)
-        if item.is_cancelled():
+    engine = OmniStageLLM(model, **engine_args)
-            continue
+    batch_seq = 0
-        with cuda_stream(ctx.stream):
+    while True:
-            out_segments = run_stage_compute(item)  # produce segments
+        first = input_q.get()
-        for seg in out_segments:
+        if first is None:
-            output_q.put(seg)
+            break
-
+        # Batch intake (up to runtime.max_batch_size)
-class WindowState:
+        batch = [first]
-    def __init__(self, window_size):
+        while len(batch) < int(runtime.max_batch_size or 1) and not input_q.empty():
-        self.window_size = window_size
+            nxt = input_q.get()
-        self.last_emitted = 0
+            if nxt is None:
-        self.buffer = []  # references/paged views
+                input_q.put(None)
-    def on_segment(self, seg):
+                break
-        self.buffer.append(seg)
+            batch.append(nxt)
-        if self.window_size == -1:
+        # Decode IPC payload and sampling params
-            return self.emit_full_if_last(seg)
+        request_ids = []
-        return self.emit_windows()
+        engine_inputs = []
-    def emit_windows(self):
+        rx_bytes = {}
-        emits = []
+        rx_decode_ms = {}
-        while self._pending_len() - self.last_emitted >= self.window_size:
+        for t in batch:
-            emits.append(self._make_view(self.last_emitted, self.window_size))
+            ein, rx = maybe_load_from_ipc_with_metrics(t, "engine_inputs", "engine_inputs_shm")
-            self.last_emitted += self.window_size
+            request_ids.append(t["request_id"])
-        return emits
+            engine_inputs.extend(ein if isinstance(ein, list) else [ein])
-    def emit_full_if_last(self, seg):
+            rx_bytes[t["request_id"]] = rx["rx_transfer_bytes"]
-        return [self._make_full_view()] if seg.is_last else []
+            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)
 ```
 
-### 11. Configuration
-- `mrs.enabled`, `mrs.ar.default_window_size`, `mrs.ar.min_flush_interval_ms`, `mrs.ar.max_segment_tokens`, `mrs.ar.scheduler_policy`.
-- Devices/processes:
-  - `mrs.stages`: [{ name, process: true, devices: "0" | "1,2", extra_env: {} }]
-  - `mrs.edges`: [{ from, to, window_size }]
-- Pinned/SHM pool sizing: estimate by concurrency and max sequence; tune watermarks and GC policies; configure CUDA IPC limits and cleanup thresholds if enabled.
-
 ### 12. Observability
-- Metrics:
+- Metrics (as implemented):
-  - Stage-level: enqueue/dequeue rates, queue depth, blocking time on watermarks, window triggers/counts.
+  - 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`
-  - Request-level: TTFT, TPST, completion latency, cancel rate.
+  - orchestrator aggregates: E2E latency, tokens/s (written only when stats are enabled)
-  - Resource-level: GPU utilization, VRAM usage/fragmentation, pinned usage.
+  - optional per-stage JSONL: `{log_file}.stage{stage_id}.stats.jsonl`
-- Logs/tracing: carry `request_id`, stage name, and `sequence`; sample decisions for downsampling/backpressure.
+- Logs/tracing:
-
+  - optional per-stage log files: `{log_file}.stage{stage_id}.log`
-### 13. Test Plan
+  - the orchestrator can log readiness, forward size/time, and summary information
-- Unit: queue correctness, window state machine, cancel/timeout, memory pool rent/return.
-- Performance: AR QPS/TPST vs. window/segment sizes; diffusion preview frequency vs. throughput.
-- Stability: long-run memory leaks; backpressure trigger/recovery paths.
-- Compatibility: behavior parity when MRS is disabled.
-
-### 14. Milestones
-- M1: Process Orchestrator with simple IPC; serial per stage; full-trigger across stages (window=-1).
-- M2: Hardening and observability; cancel/timeout and error paths.
-- M3: Optional zero-copy and concurrency/micro-batching optimizations.
 
 ### 15. Risks & Mitigations
 — End-to-end latency: full-trigger sacrifices interactivity; start simple/correct, add windowed streaming later.
 — CPU copy overhead: copy-based IPC for maintainability; add SHM/zero-copy later as an optimization path.
-— Serial throughput: `max_inflight=1` limits per-stage throughput; enable batching/concurrency later as needed.
 
 ### 16. Refactor: Sink LLM init and process into Stage (diagram)
 
 Objective: without changing external call patterns, encapsulate per-stage LLM init, process creation, and worker logic inside `Stage`, so `PipelinedOmniLLM` focuses on orchestration (seeding, polling, forwarding, collecting) while keeping multi-process, shared memory, and device mapping capabilities.
 
 Key changes (vs. current):
-- Before: `PipelinedOmniLLM` created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there.
+- Before: Orchestrator created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there.
-- After: enhance `Stage` to start its own subprocess via `start_process()`, with `worker_main()` as the entry; device setup, LLM init, batching, and SHM live in `Stage`.
+- 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
-- PipelinedOmniLLM (orchestrator)
+- OmniLLM (orchestrator)
-  - Build `Stage` list (preserve `OmniStage` init timing and `process_engine_inputs` flow)
+  - 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 `mrs`), `in_q/out_q`, subprocess handle, stats
+  - Members: `stage_config` (with `engine_args` and `runtime`), `in_q/out_q`, subprocess handle, stats
   - API:
-    - `start_process(in_q, out_q)`: spawn subprocess
+    - `init_stage_worker(...)`: spawn subprocess
-    - `stop_process()`: graceful exit
+    - `stop_stage_worker()`: graceful exit
-    - `submit(ipc_payload)`: submit to `in_q` (may use shared memory)
+    - `submit(ipc_payload)`: submit to input queue (may use shared memory)
-    - `try_collect() -> Optional[payload]`: non-blocking get from `out_q`
+    - `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)
+    - `set_stage_gpu_devices` (`CUDA_VISIBLE_DEVICES` & logical index mapping)
     - build `OmniStageLLM(model, **engine_args)`
     - while-loop:
-      - take first task + up to `mrs.max_batch_size-1` non-blocking; support `max_batch_size>=1`
+      - take first task + up to `runtime.max_batch_size-1` non-blocking; support `max_batch_size>=1`
-      - `maybe_load_from_ipc` per `engine_inputs` in batch
+      - `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 `out_q` after `maybe_dump_to_shm`
+      - 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)
 ```
-PipelinedOmniLLM
+OmniLLM
   ├─ Stage[0].start_process(in0, out0)
   ├─ Stage[1].start_process(in1, out1)
   ├─ ...
@@ -219,20 +165,24 @@ PipelinedOmniLLM
 ```
 
 Config (preserve & extend)
-- `mrs.devices`:
+- `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
-- `mrs.max_batch_size`: max requests taken per intake (default 1)
+- `runtime.max_batch_size`: max requests taken per intake (default 1)
-- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; can be injected into stages)
+- `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
-- Windows: keep `window_size=-1` semantics (downstream triggers only after upstream finishes the request)
 
 Migration (execution order)
-1) Implement enhanced `Stage` (e.g., `entrypoints/stage.py`) with `start_process/stop_process/submit/try_collect/worker_main`.
+1) Enhanced `OmniStage` with `init_stage_worker/stop_stage_worker/submit/try_collect/_stage_worker`.
-2) Make `PipelinedOmniLLM` hold `Stage` instances, replacing direct queues and `_stage_worker` with `Stage` APIs.
+2) `OmniLLM` holds `OmniStage` instances, replacing the older pipelined class; uses Stage APIs.
-3) Reuse `pipeline_utils`: device/SHM/encode-decode utilities stay in utils.
+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,22 +63,30 @@ uv pip install -r requirements.txt
 ```
 
 ## Run examples (Qwen2.5-omni)
-Download dataset from seed_tts[https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit]. To get the prompt, you can:
+### 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 \
+  --topk 100
-  --pt-output examples/offline_inference/qwen_2_5_omni/top100.pt
 ```
 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

@@ -51,9 +51,6 @@ def parse_args():
     parser.add_argument(
         "--prompts", nargs="+", default=None, help="Input text prompts."
     )
-    parser.add_argument(
-        "--pt-prompts", type=str, default=None, help="Path to .pt file containing List[str] prompts (overrides --prompts when prompt_type=text)."
-    )
     parser.add_argument(
         "--voice-type", default="default", help="Voice type, e.g., m02, f030, default."
     )
@@ -146,6 +143,12 @@ def parse_args():
         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
 
@@ -153,20 +156,19 @@ def parse_args():
 def main():
     args = parse_args()
     model_name = args.model
-    # If pt-prompts provided and prompt_type is text, load and override args.prompts
     try:
-        if getattr(args, "pt_prompts", None) and args.prompt_type == "text":
+        # Preferred: load from txt file (one prompt per line)
-            loaded = torch.load(args.pt_prompts)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
-            if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)):
+            with open(args.txt_prompts, "r", encoding="utf-8") as f:
-                raise ValueError("--pt-prompts must point to a .pt containing List[str]")
+                lines = [ln.strip() for ln in f.readlines()]
-            args.prompts = loaded
+            args.prompts = [ln for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.pt_prompts}")
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
     except Exception as e:
-        print(f"[Error] Failed to load pt-prompts: {e}")
+        print(f"[Error] Failed to load prompts: {e}")
         raise
 
     if args.prompts is None:
-        raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts <file.pt> (with --prompt_type text)")
+        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,

examples/offline_inference/qwen_2_5_omni/extract_prompts.py

@@ -18,10 +18,9 @@ def extract_prompt(line: str) -> Optional[str]:
 
 def main() -> None:
     parser = argparse.ArgumentParser()
-    parser.add_argument("--input", "-i", required=True, help="输入 .lst 文件路径")
+    parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
-    parser.add_argument("--output", "-o", required=True, help="输出文件路径")
+    parser.add_argument("--output", "-o", required=True, help="Output file path")
-    parser.add_argument("--topk", "-k", type=int, default=100, help="提取前 K 个 prompt（默认 100）")
+    parser.add_argument("--topk", "-k", type=int, default=100, help="Extract the top K prompts (default: 100)")
-    parser.add_argument("--pt-output", "-p", type=str, default=None, help="可选：输出 .pt 文件路径（默认与 --output 同名 .pt）")
     args = parser.parse_args()
 
     prompts = []
@@ -37,10 +36,6 @@ def main() -> None:
         for p in prompts:
             f.write(p + "\n")
 
-    # 另存为 PyTorch 序列化的 .pt，内容为字符串列表
-    pt_path = args.pt_output if args.pt_output else os.path.splitext(args.output)[0] + ".pt"
-    torch.save(prompts, pt_path)
-
 
 if __name__ == "__main__":
     main()

examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh

@@ -6,4 +6,4 @@ python end2end.py --model Qwen/Qwen2.5-Omni-7B \
                                  --output-wav output_audio \
                                  --prompt_type text \
                                  --init-sleep-seconds 0 \
-                                 --pt-prompts <Your_Prompt_File>.pt
+                                 --txt-prompts top100.txt

examples/offline_inference/qwen_2_5_omni/run_single_prompt.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 \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

vllm_omni/entrypoints/log_utils.py

@@ -6,7 +6,7 @@ import logging
 import os
 import time
 
-from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl
+from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl
 
 
 def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:

vllm_omni/entrypoints/omni_llm.py

@@ -34,7 +34,7 @@ from vllm_omni.entrypoints.log_utils import (
     init_stats_paths,
     OrchestratorMetrics,
 )
-from vllm_omni.entrypoints.pipeline_utils import (
+from vllm_omni.entrypoints.stage_utils import (
     maybe_load_from_ipc as _load,
     encode_for_ipc as _encode,
     serialize_obj as _ser,

vllm_omni/entrypoints/omni_stage.py

@@ -11,7 +11,7 @@ import logging
 import multiprocessing as mp
 from typing import Any, Dict, List, Optional, Union
 
-from vllm_omni.entrypoints.pipeline_utils import (
+from vllm_omni.entrypoints.stage_utils import (
     _to_dict,
     set_stage_gpu_devices,
     maybe_load_from_ipc_with_metrics,

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -8,7 +8,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # need to discuss
@@ -25,7 +25,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -42,7 +42,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

vllm_omni/worker/cpu/cpu_ar_model_runner.py

@@ -30,13 +30,20 @@ from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.cpu_model_runner import OmniCPUModelRunner
+from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner
 
 logger = init_logger(__name__)
 
 
 class CPUARModelRunner(OmniCPUModelRunner):
+    """Autoregressive GPU model runner that returns hidden states per request.
 
+    This runner follows the same preparation and forward path as GPUModelRunner
+    (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
+    and additionally performs lightweight sampling so that sampled tokens are
+    available in outputs. Hidden representations are taken at the same indices
+    that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
+    """
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -58,6 +65,7 @@ class CPUARModelRunner(OmniCPUModelRunner):
         torch.Tensor,
         Optional[IntermediateTensors],
         dict[str, Any],
+        Optional[dict[str, dict]],
     ]:
 
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
@@ -72,6 +80,7 @@ class CPUARModelRunner(OmniCPUModelRunner):
 
         # _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
@@ -92,9 +101,11 @@ class CPUARModelRunner(OmniCPUModelRunner):
             # 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
@@ -115,54 +126,31 @@ class CPUARModelRunner(OmniCPUModelRunner):
                     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
+                    req_info: dict[str, object] = {}
-                    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)
                     for k, v in addi_cpu.items():
                         if isinstance(v, torch.Tensor):
-                            # Slice along token dimension for prefill part
+                            # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor
-                            try:
+                            if overlay_len > 0:
-                                seg = v[
+                                try:
-                                    num_computed_tokens : num_computed_tokens
+                                    seg = v[
-                                    + overlay_len
+                                        num_computed_tokens : num_computed_tokens + overlay_len
-                                ].to(
+                                    ].detach().to("cpu").contiguous()
-                                    dtype=self.dtype,
+                                except Exception:
-                                    device=self.device,
+                                    seg = v.detach().to("cpu").contiguous()
-                                    non_blocking=True,
+                                req_info[k] = seg
-                                )
-                            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)
                             else:
-                                # Mixed types: wrap existing into list
+                                req_info[k] = v.detach().to("cpu").contiguous()
-                                self._forward_additional_information[k] = [prev_val, v]
+                        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]
@@ -233,6 +221,7 @@ class CPUARModelRunner(OmniCPUModelRunner):
             positions,
             intermediate_tensors,
             model_kwargs,
+            per_req_additional_information,
         )
 
     @torch.inference_mode()
@@ -348,6 +337,7 @@ class CPUARModelRunner(OmniCPUModelRunner):
                 positions,
                 intermediate_tensors,
                 model_kwargs,
+                per_req_additional_information,
             ) = self._preprocess(
                 scheduler_output,
                 num_scheduled_tokens_np,
@@ -388,15 +378,33 @@ class CPUARModelRunner(OmniCPUModelRunner):
         ):
 
             model_kwargs_extra = {}
-            # Only pass additional_information for the prefill part
+            # Pass per-request additional information map for this step (no concat)
-            if (
+            if per_req_additional_information:
-                hasattr(self, "_forward_additional_information")
+                model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
-                and self._forward_additional_information is not None
+            # Always pass per-request runtime additional_information (persisted in request state)
-                and isinstance(self._forward_additional_information, dict)
+            try:
-            ):
+                per_req_runtime_info = []
-                model_kwargs_extra["additional_information"] = (
+                for req_id in self.input_batch.req_ids:
-                    self._forward_additional_information
+                    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,
@@ -421,8 +429,42 @@ class CPUARModelRunner(OmniCPUModelRunner):
             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.
@@ -582,3 +624,26 @@ class CPUARModelRunner(OmniCPUModelRunner):
             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)

vllm_omni/worker/cpu/cpu_ar_worker.py

@@ -11,7 +11,7 @@ from vllm.platforms import current_platform, CpuArchEnum
 from vllm.v1.worker.cpu_worker import CPUWorker
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
 
-from vllm_omni.worker.cpu_ar_model_runner import CPUARModelRunner
+from vllm_omni.worker.cpu.cpu_ar_model_runner import CPUARModelRunner
 
 logger = init_logger(__name__)
 

vllm_omni/worker/cpu/cpu_diffusion_model_runner.py

@@ -10,11 +10,7 @@ import torch
 from vllm.config import CUDAGraphMode
 from vllm.forward_context import BatchDescriptor
 from vllm.multimodal.inputs import MultiModalKwargs
-from vllm.utils import cdiv
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
-    split_attn_metadata,
-)
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.utils import record_function_or_nullcontext
@@ -29,7 +25,7 @@ from vllm.v1.worker.gpu_model_runner import (
 from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.cpu_model_runner import OmniCPUModelRunner
+from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner
 
 logger = logging.getLogger(__name__)
 
@@ -621,4 +617,4 @@ class CPUDiffusionModelRunner(OmniCPUModelRunner):
         self._sync_device()
         del hidden_states
         self.encoder_cache.clear()
-        gc.collect()
+        gc.collect()

vllm_omni/worker/cpu/cpu_diffusion_worker.py

@@ -13,7 +13,7 @@ from vllm.v1.worker.cpu_worker import CPUWorker
 
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
 
-from vllm_omni.worker.cpu_diffusion_model_runner import CPUDiffusionModelRunner
+from vllm_omni.worker.cpu.cpu_diffusion_model_runner import CPUDiffusionModelRunner
 
 
 logger = init_logger(__name__)

vllm_omni/worker/cpu/cpu_model_runner.py

@@ -266,7 +266,6 @@ class OmniCPUModelRunner(CPUModelRunner):
     def extract_multimodal_outputs(
             self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]
     ) -> dict:
-        print('hidden_states',hidden_states)
         if (
                 hasattr(self.model, "have_multimodal_outputs")
                 and self.model.have_multimodal_outputs
@@ -588,7 +587,7 @@ class OmniCPUModelRunner(CPUModelRunner):
         hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
             hidden_states
         )
-        if len(hidden_states.shape)==3:
+        if len(hidden_states.shape) == 3:
             return hidden_states[0], hidden_states[0][logit_indices]
         else:
-            return hidden_states, hidden_states[logit_indices]
+            return hidden_states, hidden_states[logit_indices]

vllm_omni/worker/gpu/gpu_ar_model_runner.py

@@ -30,7 +30,7 @@ from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner
 
 logger = init_logger(__name__)
 

vllm_omni/worker/gpu/gpu_ar_worker.py

@@ -9,7 +9,7 @@ from vllm.utils import GiB_bytes, MemorySnapshot
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.gpu.gpu_ar_model_runner import GPUARModelRunner
 
 
 class GPUARWorker(GPUWorker):

vllm_omni/worker/gpu/gpu_diffusion_model_runner.py

@@ -25,7 +25,7 @@ from vllm.v1.worker.gpu_model_runner import (
 from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner
 
 logger = logging.getLogger(__name__)
 

vllm_omni/worker/gpu/gpu_diffusion_worker.py

@@ -9,7 +9,7 @@ from vllm.utils import GiB_bytes, MemorySnapshot
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
+from vllm_omni.worker.gpu.gpu_diffusion_model_runner import GPUDiffusionModelRunner
 
 
 class GPUDiffusionWorker(GPUWorker):