diff --git a/docs/mrs_design.md b/docs/mrs_design.md new file mode 100644 index 000000000..2c463fa64 --- /dev/null +++ b/docs/mrs_design.md @@ -0,0 +1,238 @@ +## Multi-Request Streaming (MRS) on a Single Machine — Design Doc + +- Author: Team +- Date: 2025-10-27 +- Version: v0.3 (simplified: copy-based IPC / serial per stage / full-trigger) + +### 1. Background & Scope +- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved. +- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling. +- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead. + +### 2. Key Constraints +- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`). +- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version. +- Full-trigger across stages: unify `window_size = -1` for all edges — downstream starts only after upstream completes for that request. +- Serial within a stage: at most one in-flight request per stage (max_inflight=1). No batching/concurrency inside a stage in this version. +- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0). + +### 3. Architecture Overview +- Processes & IPC queues + - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue. + - Inter-stage connection via IPC: control plane (Unix Domain Socket / Pipe) + data plane (shared memory ring buffer in future; copy-based in this version). + - MPMC queue semantics; control plane passes light-weight pointers/metadata; data plane carries large payloads. +- Device visibility + - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets. + - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit. +- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work. +- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately. +- Scheduling & windows + - Each edge (stage→stage) defines a window policy; downstream triggers compute by consuming stitched window views. + - `-1` mode triggers only after upstream emits `is_last=True` (full aggregation). + - Backpressure propagates upstream via the control plane: once high-watermark is reached, upstream blocks or degrades behavior. + +### 4. Stage Breakdown +- AR (autoregressive) + - S0 Intake/Plan: parse request, assign request_id, output normalized tasks. + - S1 Prefill: tokenize/encode/KV warmup; may micro-batch across requests. + - S2 Decode Loop: step-by-step tokens; flush segments by `max_segment_tokens` and `min_flush_interval_ms`. + - S3 Postprocess: detokenize, logprobs/struct wrapping (CPU if necessary). + - S4 Consumer: local consumer callback/buffering. +- DiT/Diffusion + - D0 Intake/Plan: scheduler/seed init. + - D1 Sampling Loop: preview frames every k steps (downsample allowed). + - D2 Decode/Encode: VAE decode and JPEG/WEBP encode (prefer GPU→pinned CPU). + - D3 Postprocess/Consumer. + +### 5. Unified Events & Segment Structure (in-process) +- Base event: `{ request_id, sequence, event_type, ts, payload, device_handles }` +- ARSegmentEvent (segment): + - Fields: `tokens` (List[int] or GPU tensor view), optional `logprobs`, `is_last`, optional `hidden_states/aux` refs. + - Semantics: strictly ordered by `sequence` for the same `request_id`; `is_last` marks the final segment. +- ImageProgressEvent: `{ step, image_tensor|jpeg_ref, is_last }` (diffusion preview/completion). +- Completed/Error: terminal or error events affect only the current request. + +### 6. AR Segmentation & Windowing (simplified) +- Production (upstream stage) + - Accumulate tokens during decoding; flush a segment if any of the following holds: + - Reaches `max_segment_tokens` length + - Exceeds `min_flush_interval_ms` since last flush + - Request ended (`is_last=True`) + - Flush immediately via `output_queue.put(segment)`; do not wait for entire request completion. +- Consumption (downstream stage) + - Use unified `window_size = -1`: trigger only after receiving `is_last=True`; no window slicing. +- Backpressure & watermark + - Must-deliver events (token.delta, Completed) are not dropped; block upstream briefly on high watermark. + - For long sequences under `-1`, allow early segments to spill from GPU to pinned CPU to control VRAM pressure. + +#### Example config (YAML, simplified) +```yaml +mrs: + enabled: true + defaults: + window_size: -1 # Full-trigger across stages + max_inflight: 1 # Serial per stage +``` + +### 7. IPC Implementation (simplified: copy-based) +- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data). +- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version. +- Backpressure: blocking reads/writes based on queue capacity, blocking upstream on high watermark. + +### 8. Scheduling & Cancellation (simplified) +- Serial FIFO: each stage `max_inflight=1`; no cross-request micro-batching or concurrency. +- Cancel/timeout: ongoing work should exit early; remaining queued requests stay pending. +- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency. + +#### Short sequence example (req0/req1, stage A→B) +1) t0: stage A handles req0 +2) t1: req0 completes on A → enters B; A immediately starts req1 +3) t2: B handles req0 while A handles req1 (parallel across stages) + +### 9. Integration Points (by file) +- `vllm_omni/entrypoints/omni_llm.py` + - Introduce a ProcessOrchestrator (replace threaded version): + - Spawn stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), create control/data channels, build window state machines. + - Add `enable_mrs` in `generate`; when enabled, do not collect `stage.engine.generate(...)` into lists; stream segments to downstream via IPC instead. + - Register `window_size` per edge; `-1` triggers downstream once `is_last` is observed. + - Manage process lifecycle: start/heartbeat/graceful shutdown/restart on crash; propagate cancel/timeout. +- `vllm_omni/outputs.py` + - Add in-process events: `ARSegmentEvent`, `ImageProgressEvent`, `CompletedEvent`, `ErrorEvent`. + - Keep `OmniRequestOutput` as the external representation of final stage outputs. +- `vllm_omni/worker/gpu_ar_model_runner.py` + - Near `valid_sampled_token_ids` parsing (after updating `req_state.output_token_ids`), aggregate new tokens into segments and emit; send final `is_last=True` and Completed. + - Replace current `.to("cpu").contiguous()` sync path: by default keep GPU refs and let downstream copy asynchronously from a pinned pool when needed. +- `vllm_omni/worker/gpu_diffusion_model_runner.py` + +### 9.1 Process Device Visibility +- Device binding: set `CUDA_VISIBLE_DEVICES` before process start, or call `torch.cuda.set_device` early in init; configs can be single/multi-GPU. +- Cross-device transfer (simplified): via CPU copies; zero-copy paths are out of scope for now. + +### 10. Worker Template & Window State (pseudo-code) +```python +def stage_worker(input_q, output_q, ctx): + while ctx.alive: + item = input_q.get() # blocking / with timeout + if item.is_cancelled(): + continue + with cuda_stream(ctx.stream): + out_segments = run_stage_compute(item) # produce segments + for seg in out_segments: + output_q.put(seg) + +class WindowState: + def __init__(self, window_size): + self.window_size = window_size + self.last_emitted = 0 + self.buffer = [] # references/paged views + def on_segment(self, seg): + self.buffer.append(seg) + if self.window_size == -1: + return self.emit_full_if_last(seg) + return self.emit_windows() + def emit_windows(self): + emits = [] + while self._pending_len() - self.last_emitted >= self.window_size: + emits.append(self._make_view(self.last_emitted, self.window_size)) + self.last_emitted += self.window_size + return emits + def emit_full_if_last(self, seg): + return [self._make_full_view()] if seg.is_last else [] +``` + +### 11. Configuration +- `mrs.enabled`, `mrs.ar.default_window_size`, `mrs.ar.min_flush_interval_ms`, `mrs.ar.max_segment_tokens`, `mrs.ar.scheduler_policy`. +- Devices/processes: + - `mrs.stages`: [{ name, process: true, devices: "0" | "1,2", extra_env: {} }] + - `mrs.edges`: [{ from, to, window_size }] +- Pinned/SHM pool sizing: estimate by concurrency and max sequence; tune watermarks and GC policies; configure CUDA IPC limits and cleanup thresholds if enabled. + +### 12. Observability +- Metrics: + - Stage-level: enqueue/dequeue rates, queue depth, blocking time on watermarks, window triggers/counts. + - Request-level: TTFT, TPST, completion latency, cancel rate. + - Resource-level: GPU utilization, VRAM usage/fragmentation, pinned usage. +- Logs/tracing: carry `request_id`, stage name, and `sequence`; sample decisions for downsampling/backpressure. + +### 13. Test Plan +- Unit: queue correctness, window state machine, cancel/timeout, memory pool rent/return. +- Performance: AR QPS/TPST vs. window/segment sizes; diffusion preview frequency vs. throughput. +- Stability: long-run memory leaks; backpressure trigger/recovery paths. +- Compatibility: behavior parity when MRS is disabled. + +### 14. Milestones +- M1: Process Orchestrator with simple IPC; serial per stage; full-trigger across stages (window=-1). +- M2: Hardening and observability; cancel/timeout and error paths. +- M3: Optional zero-copy and concurrency/micro-batching optimizations. + +### 15. Risks & Mitigations +— End-to-end latency: full-trigger sacrifices interactivity; start simple/correct, add windowed streaming later. +— CPU copy overhead: copy-based IPC for maintainability; add SHM/zero-copy later as an optimization path. +— Serial throughput: `max_inflight=1` limits per-stage throughput; enable batching/concurrency later as needed. + +### 16. Refactor: Sink LLM init and process into Stage (diagram) + +Objective: without changing external call patterns, encapsulate per-stage LLM init, process creation, and worker logic inside `Stage`, so `PipelinedOmniLLM` focuses on orchestration (seeding, polling, forwarding, collecting) while keeping multi-process, shared memory, and device mapping capabilities. + +Key changes (vs. current): +- Before: `PipelinedOmniLLM` created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there. +- After: enhance `Stage` to start its own subprocess via `start_process()`, with `worker_main()` as the entry; device setup, LLM init, batching, and SHM live in `Stage`. + +Classes & responsibilities +- PipelinedOmniLLM (orchestrator) + - Build `Stage` list (preserve `OmniStage` init timing and `process_engine_inputs` flow) + - Connect adjacent `Stage` input/output queues; seed requests to stage 0 + - Poll stage outputs, decode results, call `process_engine_inputs`, then encode and forward to next stage + - Termination & cleanup: distribute shutdown signals; join/terminate subprocesses +- Stage (stage unit) + - Members: `stage_config` (with `engine_args` and `mrs`), `in_q/out_q`, subprocess handle, stats + - API: + - `start_process(in_q, out_q)`: spawn subprocess + - `stop_process()`: graceful exit + - `submit(ipc_payload)`: submit to `in_q` (may use shared memory) + - `try_collect() -> Optional[payload]`: non-blocking get from `out_q` + - `process_engine_inputs(stages, prompts)`: reuse/delegate existing logic + - Subprocess entry `worker_main()`: + - set_stage_gpu_devices (`CUDA_VISIBLE_DEVICES` & logical index mapping) + - build `OmniStageLLM(model, **engine_args)` + - while-loop: + - take first task + up to `mrs.max_batch_size-1` non-blocking; support `max_batch_size>=1` + - `maybe_load_from_ipc` per `engine_inputs` in batch + - `stage_engine.generate(batched_inputs, sampling_params)` (window `-1`) + - write per-request aggregated outputs to `out_q` after `maybe_dump_to_shm` + - log start/finish; report `error` per request on exceptions + +Data flow (sketch) +``` +PipelinedOmniLLM + ├─ Stage[0].start_process(in0, out0) + ├─ Stage[1].start_process(in1, out1) + ├─ ... + └─ Orchestrator loop: + - stage0.submit(encode(engine_inputs)) + - for each stage: + res = stage.try_collect() + outputs = decode(res) + stage.set_engine_outputs(outputs) + if !final: next_inputs = next_stage.process_engine_inputs(stages, original_prompt) + next_stage.submit(encode(next_inputs)) +``` + +Config (preserve & extend) +- `mrs.devices`: + - Comma-separated visibility list (e.g., "2,5,7"): set visible set; default logical index 0 (first GPU) + - Number/string: treat as logical index (0/1-based per implementation) mapping to physical GPUs +- `mrs.max_batch_size`: max requests taken per intake (default 1) +- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; can be injected into stages) + +Kept & strengthened +- Multi-process: each stage in its own process; parallel across stages; serial/batching inside a stage controlled by `max_batch_size` +- Shared memory: large objects via SHM with strict `memoryview -> close -> unlink` +- Device mapping: unified `set_stage_gpu_devices`, support comma lists and logical indices +- Windows: keep `window_size=-1` semantics (downstream triggers only after upstream finishes the request) + +Migration (execution order) +1) Implement enhanced `Stage` (e.g., `entrypoints/stage.py`) with `start_process/stop_process/submit/try_collect/worker_main`. +2) Make `PipelinedOmniLLM` hold `Stage` instances, replacing direct queues and `_stage_worker` with `Stage` APIs. +3) Reuse `pipeline_utils`: device/SHM/encode-decode utilities stay in utils. +4) Regression: single/multi-request, various `max_batch_size`, SHM threshold, device mapping, multi-stage parallelism, error paths, and shutdown flows. diff --git a/examples/offline_inference/qwen_2_5_omni/extract_prompts.py b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py new file mode 100644 index 000000000..a1bc7a78f --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/extract_prompts.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +import argparse +from typing import Optional +import os +import torch + + +def extract_prompt(line: str) -> Optional[str]: + # 提取第一个 '|' 与第二个 '|' 之间的内容 + i = line.find('|') + if i == -1: + return None + j = line.find('|', i + 1) + if j == -1: + return None + return line[i + 1:j].strip() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input", "-i", required=True, help="输入 .lst 文件路径") + parser.add_argument("--output", "-o", required=True, help="输出文件路径") + parser.add_argument("--topk", "-k", type=int, default=100, help="提取前 K 个 prompt（默认 100）") + parser.add_argument("--pt-output", "-p", type=str, default=None, help="可选：输出 .pt 文件路径（默认与 --output 同名 .pt）") + args = parser.parse_args() + + prompts = [] + with open(args.input, "r", encoding="utf-8", errors="ignore") as f: + for line in f: + if len(prompts) >= args.topk: + break + p = extract_prompt(line.rstrip("\n")) + if p: + prompts.append(p) + + with open(args.output, "w", encoding="utf-8") as f: + for p in prompts: + f.write(p + "\n") + + # 另存为 PyTorch 序列化的 .pt，内容为字符串列表 + pt_path = args.pt_output if args.pt_output else os.path.splitext(args.output)[0] + ".pt" + torch.save(prompts, pt_path) + + +if __name__ == "__main__": + main() + + diff --git a/examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py b/examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py new file mode 100644 index 000000000..510c60f82 --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/hf_qwen25_omni_batch_no_template.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +Batch text generation with Qwen 2.5 Omni using Transformers without chat templates. + +- Processes 10 fixed prompts in batches of 3 +- Does NOT use tokenizer.apply_chat_template (raw prompts only) +- Prints only generated continuations for each prompt +""" + +import os +import argparse +from typing import Iterable, List +import time +import torch + +import random +import numpy as np +import torch +from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor +from tqdm import tqdm +import json +import soundfile as sf + +PROMPTS: List[str] = [ + "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.", + "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.", + "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.", + "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.", + "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.", + "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.", + "Generate a friendly greeting message suitable for a voice assistant. Answer in 15 words.", + "Summarize the benefits of on-device inference for edge deployments. Answer in 15 words.", + "Offer troubleshooting steps for degraded audio quality in a TTS pipeline. Answer in 15 words.", + "List best practices for optimizing latency in real-time voice synthesis. Answer in 15 words.", +] +def make_text_prompt(prompt): + return [{ + "role": "system", + "content": [ + {"type": "text", "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."} + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + ], + }] + +def chunked(items: List[str], size: int) -> Iterable[List[str]]: + for i in range(0, len(items), size): + yield items[i : i + size] + + +def main(): + parser = argparse.ArgumentParser( + description="Qwen 2.5 Omni batch generation without chat templates", + ) + parser.add_argument( + "--model-id", + type=str, + default=os.getenv("QWEN_OMNI_MODEL_ID", "Qwen/Qwen2.5-Omni"), + help="HF model repo id (e.g., Qwen/Qwen2.5-Omni)", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=2048, + help="Maximum number of new tokens to generate", + ) + parser.add_argument( + "--temperature", type=float, default=0.0, help="Sampling temperature" + ) + parser.add_argument("--top-p", type=float, default=1.0, help="Top-p sampling") + parser.add_argument( + "--repetition-penalty", + type=float, + default=1.1, + help="Repetition penalty (matches pipelined_end2end defaults)", + ) + parser.add_argument( + "--stop-token-ids", + type=str, + default="8294", + help="Comma-separated stop token ids (added to eos_token_id)", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for reproducibility (sampling)", + ) + parser.add_argument( + "--pt-prompts", + type=str, + default=None, + help="Path to .pt file containing List[str] prompts (overrides built-in PROMPTS)", + ) + parser.add_argument( + "--tokenize", + action="store_true", + help="If set, utils.make_text_prompt will return token ids; keep unset for raw text", + ) + parser.add_argument( + "--output-dir", + type=str, + default="outputs", + help="Output directory to save aligned comparison files (same folder as pipelined_end2end)", + ) + args = parser.parse_args() + + if torch.cuda.is_available(): + torch_dtype = torch.bfloat16 + else: + torch_dtype = torch.float32 + + # Match pipelined_end2end.py deterministic settings + if args.seed is not None: + os.environ["PYTHONHASHSEED"] = str(args.seed) + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + try: + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + except Exception: + pass + + print(f"[Info] Loading model: {args.model_id}") + processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id) + model = Qwen2_5OmniForConditionalGeneration.from_pretrained( + args.model_id, + trust_remote_code=True, + torch_dtype=torch_dtype, + device_map="auto", + attn_implementation="flash_attention_2", + ) + + # Load prompts: prefer .pt if provided + if args.pt_prompts is not None: + loaded = torch.load(args.pt_prompts) + if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)): + raise ValueError("--pt-prompts must point to a .pt containing List[str]") + prompts: List[str] = loaded + print(f"[Info] Loaded {len(prompts)} prompts from {args.pt_prompts}") + else: + prompts = PROMPTS + + # Ensure utils.make_text_prompt can find expected fields + # utils expects args.model and args.tokenize + setattr(args, "model", args.model_id) + if not hasattr(args, "tokenize"): + setattr(args, "tokenize", False) + + print(f"[Info] Running single-sample generation (batch size fixed to 1) ...\n") + global_index = 0 + total_requests = len(prompts) + sum_request_time_ms = 0.0 + wall_t0 = time.time() + stats_path = os.path.join(args.output_dir, "hf_transformers.stats.jsonl") + + for rid, prompt in tqdm(list(enumerate(prompts)), total=len(prompts), desc="Generating"): + _t0 = time.time() + # Build single-sample inputs + tp = make_text_prompt(prompt) + tp = processor.apply_chat_template(tp, add_generation_prompt=True, tokenize=False) + inputs = processor(text=tp, return_tensors="pt", padding=True, use_audio_in_video=True) + inputs = inputs.to(model.device).to(model.dtype) + + with torch.no_grad(): + text_ids, audio = model.generate( + **inputs, + use_audio_in_video=True, + ) + + # Decode + try: + texts = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # type: ignore[arg-type] + except Exception: + try: + texts = processor.tokenizer.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # type: ignore[attr-defined] + except Exception: + texts = [str(x) for x in text_ids] + + text_str = texts[0] if len(texts) > 0 else "" + os.makedirs(args.output_dir, exist_ok=True) + out_txt = os.path.join(args.output_dir, f"{rid:05d}.txt") + + content = None + if os.path.exists(out_txt): + try: + with open(out_txt, "r", encoding="utf-8") as f: + content = f.read() + except Exception: + content = None + if content and "transformers_text_output:" in content: + try: + parts = content.split("transformers_text_output:\n", 1) + new_content = parts[0] + "transformers_text_output:\n" + text_str.strip() + "\n" + with open(out_txt, "w", encoding="utf-8") as f: + f.write(new_content) + except Exception: + pass + else: + lines = [ + "Prompt:\n", + str(prompt) + "\n", + "vllm_text_output:\n\n", + "transformers_text_output:\n", + text_str.strip() + "\n", + ] + try: + with open(out_txt, "w", encoding="utf-8") as f: + f.writelines(lines) + except Exception: + pass + + # Save wav per request + try: + if audio is not None: + if isinstance(audio, (list, tuple)): + wav = audio[0] + else: + # Tensor or array: take first sample if batched, else as-is + if hasattr(audio, "ndim") and hasattr(audio, "shape") and len(audio.shape) > 1: + wav = audio[0] + else: + wav = audio + wav_path = os.path.join(args.output_dir, f"output_{rid}.wav") + sf.write(wav_path, torch.as_tensor(wav).reshape(-1).detach().cpu().numpy(), samplerate=24000) + except Exception: + pass + + # Stats per request + req_ms = (time.time() - _t0) * 1000.0 + sum_request_time_ms += req_ms + try: + rec = { + "type": "transformers_request_stats", + "request_id": rid, + "batch_size": 1, + "per_request_time_ms": req_ms, + "num_tokens_out": int(text_ids.shape[-1]) if hasattr(text_ids, "shape") else None, + } + with open(stats_path, "a", encoding="utf-8") as f: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + except Exception: + pass + + global_index += 1 + + print("[Done] Generated", global_index, "responses.") + wall_ms = (time.time() - wall_t0) * 1000.0 + avg_e2e_ms = (sum_request_time_ms / total_requests) if total_requests > 0 else 0.0 + print(f"[E2E] total_requests={total_requests}, wall_time_ms={wall_ms:.2f}, avg_e2e_per_request_ms={avg_e2e_ms:.2f}") + try: + summary = { + "type": "transformers_summary", + "total_requests": total_requests, + "wall_time_ms": wall_ms, + "sum_request_time_ms": sum_request_time_ms, + "avg_e2e_per_request_ms": avg_e2e_ms, + } + with open(stats_path, "a", encoding="utf-8") as f: + f.write(json.dumps(summary, ensure_ascii=False) + "\n") + except Exception: + pass + + +if __name__ == "__main__": + main() + + diff --git a/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py b/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py new file mode 100644 index 000000000..9b034521e --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/pipelined_end2end.py @@ -0,0 +1,264 @@ +import argparse +import os +import os as _os_env_toggle +import random + +import numpy as np +import soundfile as sf +import torch +from utils import make_omni_prompt +from vllm.sampling_params import SamplingParams + +from vllm_omni.entrypoints.pipelined_omni_llm import PipelinedOmniLLM + +_os_env_toggle.environ["VLLM_USE_V1"] = "1" + +SEED = 42 +# Set all random seeds +random.seed(SEED) +np.random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) + +# Make PyTorch deterministic +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False + +# Set environment variables for deterministic behavior +os.environ["PYTHONHASHSEED"] = str(SEED) +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + required=True, + help="Path to merged model directory (will be created if downloading).", + ) + parser.add_argument("--thinker-model", type=str, default=None) + parser.add_argument("--talker-model", type=str, default=None) + parser.add_argument("--code2wav-model", type=str, default=None) + parser.add_argument( + "--hf-hub-id", + default="Qwen/Qwen2.5-Omni-7B", + help="Hugging Face repo id to download if needed.", + ) + parser.add_argument( + "--hf-revision", default=None, help="Optional HF revision (branch/tag/commit)." + ) + parser.add_argument( + "--prompts", nargs="+", default=None, help="Input text prompts." + ) + parser.add_argument( + "--pt-prompts", type=str, default=None, help="Path to .pt file containing List[str] prompts (overrides --prompts when prompt_type=text)." + ) + parser.add_argument( + "--voice-type", default="default", help="Voice type, e.g., m02, f030, default." + ) + parser.add_argument( + "--code2wav-dir", + default=None, + help="Path to code2wav folder (contains spk_dict.pt).", + ) + parser.add_argument( + "--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt)." + ) + parser.add_argument( + "--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file." + ) + parser.add_argument( + "--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"] + ) + parser.add_argument("--max-model-len", type=int, default=32768) + parser.add_argument( + "--init-sleep-seconds", + type=int, + default=20, + help="Sleep seconds after starting each stage process to allow initialization (default: 20)", + ) + + parser.add_argument("--thinker-only", action="store_true") + parser.add_argument("--text-only", action="store_true") + parser.add_argument("--do-wave", action="store_true") + parser.add_argument( + "--prompt_type", + choices=[ + "text", + "audio", + "audio-long", + "audio-long-chunks", + "audio-long-expand-chunks", + "image", + "video", + "video-frames", + "audio-in-video", + "audio-in-video-v2", + "audio-multi-round", + "badcase-vl", + "badcase-text", + "badcase-image-early-stop", + "badcase-two-audios", + "badcase-two-videos", + "badcase-multi-round", + "badcase-voice-type", + "badcase-voice-type-v2", + "badcase-audio-tower-1", + "badcase-audio-only", + ], + default="text", + ) + parser.add_argument("--use-torchvision", action="store_true") + parser.add_argument("--tokenize", action="store_true") + parser.add_argument( + "--output-wav", default="output.wav", help="[Deprecated] Output wav directory (use --output-dir)." + ) + parser.add_argument( + "--output-dir", default="outputs", help="Output directory to save text and wav files together." + ) + parser.add_argument( + "--thinker-hidden-states-dir", + default="thinker_hidden_states", + help="Path to thinker hidden states directory.", + ) + parser.add_argument( + "--batch-timeout", + type=int, + default=5, + help="Timeout for batching in seconds (default: 5)", + ) + parser.add_argument( + "--init-timeout", + type=int, + default=300, + help="Timeout for initializing stages in seconds (default: 300)", + ) + parser.add_argument( + "--shm-threshold-bytes", + type=int, + default=65536, + help="Threshold for using shared memory in bytes (default: 65536)", + ) + parser.add_argument( + "--enable-stats", + action="store_true", + default=True, + help="Enable writing detailed statistics (default: enabled)", + ) + args = parser.parse_args() + return args + + +def main(): + args = parse_args() + model_name = args.model + # If pt-prompts provided and prompt_type is text, load and override args.prompts + try: + if getattr(args, "pt_prompts", None) and args.prompt_type == "text": + loaded = torch.load(args.pt_prompts) + if not isinstance(loaded, list) or (len(loaded) > 0 and not isinstance(loaded[0], str)): + raise ValueError("--pt-prompts must point to a .pt containing List[str]") + args.prompts = loaded + print(f"[Info] Loaded {len(args.prompts)} prompts from {args.pt_prompts}") + except Exception as e: + raise + + if args.prompts is None: + raise ValueError("No prompts provided. Use --prompts ... or --pt-prompts (with --prompt_type text)") + omni_llm = PipelinedOmniLLM( + model=model_name, + log_stats=args.enable_stats, + log_file="omni_llm_pipeline.log", + init_sleep_seconds=args.init_sleep_seconds, + batch_timeout=args.batch_timeout, + init_timeout=args.init_timeout, + shm_threshold_bytes=args.shm_threshold_bytes, + ) + thinker_sampling_params = SamplingParams( + temperature=0.0, # Deterministic - no randomness + top_p=1.0, # Disable nucleus sampling + top_k=-1, # Disable top-k sampling + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.1, + ) + talker_sampling_params = SamplingParams( + temperature=0.9, + top_p=0.8, + top_k=40, + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.05, + stop_token_ids=[8294], + ) + code2wav_sampling_params = SamplingParams( + temperature=0.0, # Deterministic - no randomness + top_p=1.0, # Disable nucleus sampling + top_k=-1, # Disable top-k sampling + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.1, + ) + + sampling_params_list = [ + thinker_sampling_params, + talker_sampling_params, + code2wav_sampling_params, + ] + + prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts] + omni_outputs = omni_llm.generate(prompt, sampling_params_list) + + # Determine output directory: prefer --output-dir; fallback to --output-wav + output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav + os.makedirs(output_dir, exist_ok=True) + for stage_outputs in omni_outputs: + if stage_outputs.final_output_type == "text": + for output in stage_outputs.request_output: + request_id = int(output.request_id) + text_output = output.outputs[0].text + # Save aligned text file per request + prompt_text = args.prompts[request_id] + out_txt = os.path.join(output_dir, f"{request_id:05d}.txt") + lines = [] + lines.append("Prompt:\n") + lines.append(str(prompt_text) + "\n") + lines.append("vllm_text_output:\n") + lines.append(str(text_output).strip() + "\n") + lines.append("transformers_text_output:\n\n") + try: + with open(out_txt, "w", encoding="utf-8") as f: + f.writelines(lines) + except Exception: + pass + print(f"Request ID: {request_id}, Text saved to {out_txt}") + elif stage_outputs.final_output_type == "audio": + for output in stage_outputs.request_output: + request_id = int(output.request_id) + audio_tensor = output.multimodal_output["audio"] + output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav") + sf.write( + output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000 + ) + # Save aligned text file with empty text (no text output in this stage) + prompt_text = args.prompts[request_id] + out_txt = os.path.join(output_dir, f"{request_id:05d}.txt") + lines = [] + lines.append("Prompt:\n") + lines.append(str(prompt_text) + "\n") + lines.append("vllm_text_output:\n\n") + lines.append("transformers_text_output:\n\n") + try: + with open(out_txt, "w", encoding="utf-8") as f: + f.writelines(lines) + except Exception: + pass + print(f"Request ID: {request_id}, Saved audio to {output_wav} and text to {out_txt}") + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 06b1f1b4c..a76c1126f 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -1,9 +1,22 @@ """ Stage manager for orchestrating multiple engines in vLLM-omni. + +Enhanced to encapsulate per-stage process lifecycle and worker logic +(device setup, LLM init, batching, shared-memory IPC), while preserving +the original input processing utilities for cross-stage data wiring. """ import importlib -from typing import List, Union +import logging +import multiprocessing as mp +from typing import Any, Dict, List, Optional, Union + +from vllm_omni.entrypoints.pipeline_utils import ( + _to_dict, + set_stage_gpu_devices, + maybe_load_from_ipc, + maybe_dump_to_shm, +) from vllm.inputs import TextPrompt from vllm.v1.engine import EngineCoreOutput @@ -46,6 +59,14 @@ class OmniStage: else: self.final_output_type = None + # Runtime orchestration state (added) + self._in_q: Optional[mp.Queue] = None + self._out_q: Optional[mp.Queue] = None + self._proc: Optional[mp.Process] = None + self._log_file: Optional[str] = None + self._shm_threshold_bytes: int = 65536 + self._logger = logging.getLogger(__name__) + def set_engine(self, engine: LLMEngine) -> None: """Initialize the engine for the stage.""" self.engine = engine @@ -58,6 +79,67 @@ class OmniStage: """Set the engine output for the stage.""" self.engine_outputs = engine_outputs + # ----------------- New Orchestration APIs ----------------- + def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None: + self._in_q = in_q + self._out_q = out_q + + def init_stage_worker( + self, + model: str, + *, + log_file: Optional[str] = None, + shm_threshold_bytes: int = 65536, + ctx: Optional[mp.context.BaseContext] = None, + batch_timeout: int = 10, + ) -> None: + assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process" + self._log_file = log_file + self._shm_threshold_bytes = shm_threshold_bytes + ctx = ctx or mp.get_context("spawn") + # Prepare lightweight dict config for worker + engine_args = _to_dict(self.engine_args) + mrs_cfg = _to_dict(getattr(self.stage_config, "mrs", {})) + stage_payload: Dict[str, Any] = { + "stage_id": self.stage_id, + "engine_args": engine_args, + "mrs": mrs_cfg, + "shm_threshold_bytes": self._shm_threshold_bytes, + } + self._proc = ctx.Process( + target=_stage_worker, + args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout), + ) + self._proc.start() + + def stop_stage_worker(self) -> None: + if self._in_q is not None: + try: + self._in_q.put_nowait(None) + except Exception: + pass + if self._proc is not None: + try: + self._proc.join(timeout=5) + except Exception: + pass + if self._proc.is_alive(): + try: + self._proc.terminate() + except Exception: + pass + + def submit(self, payload: Dict[str, Any]) -> None: + assert self._in_q is not None + self._in_q.put(payload) + + def try_collect(self) -> Optional[Dict[str, Any]]: + assert self._out_q is not None + try: + return self._out_q.get_nowait() + except Exception: + return None + def process_engine_inputs( self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None ) -> List[Union[OmniTokensPrompt, TextPrompt]]: @@ -90,3 +172,253 @@ class OmniStage: return self.custom_process_input_func( stage_list, engine_input_source, prompt ) + + +def _stage_worker( + model: str, + stage_payload: Dict[str, Any], + in_q: mp.Queue, + out_q: mp.Queue, + log_file: Optional[str] = None, + batch_timeout: int = 10, +) -> None: + """Stage worker entry: device setup, LLM init, batching, SHM IPC.""" + import logging as _logging + from vllm_omni.entrypoints.omni_llm import OmniStageLLM # noqa: WPS433 + import queue as _queue + import time as _time + from vllm_omni.entrypoints.pipeline_utils import append_jsonl as _append_jsonl, serialize_obj as _ser + + stage_id = stage_payload["stage_id"] + engine_args = stage_payload.get("engine_args", {}) + mrs_cfg = stage_payload.get("mrs", {}) + shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536)) + + # Per-stage file logger (optional) + try: + if log_file: + stage_log = _logging.getLogger(__name__) + stage_log.setLevel(_logging.DEBUG) + fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log") + fh.setLevel(_logging.DEBUG) + fh.setFormatter(_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")) + class _StageFilter(_logging.Filter): + def filter(self, record: _logging.LogRecord) -> bool: + setattr(record, "stage", stage_id) + return True + fh.addFilter(_StageFilter()) + stage_log.addHandler(fh) + except Exception: + pass + + # Stage stats JSONL file + _stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None + + # Aggregates for running average + _agg_total_tokens = 0 + _agg_total_gen_time_ms = 0.0 + + # Device mapping + try: + set_stage_gpu_devices(stage_id, mrs_cfg.get("devices")) + except Exception as e: + _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e) + + # Init LLM + _logging.getLogger(__name__).debug("[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys())) + stage_engine = OmniStageLLM(model=model, **engine_args) + _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id) + # Signal readiness to orchestrator + try: + out_q.put({"type": "stage_ready", "stage_id": stage_id}) + except Exception: + pass + + # Batch processing loop + while True: + task = in_q.get() + _recv_dequeue_ts = _time.time() + if task is None: + _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id) + break + + max_batch_size = int(mrs_cfg.get("max_batch_size", 1) or 1) + batch_tasks: List[Dict[str, Any]] = [task] + if max_batch_size > 1: + while len(batch_tasks) < max_batch_size: + if not in_q.empty(): + extra = in_q.get(timeout=batch_timeout) + if extra is None: + in_q.put(None) + break + batch_tasks.append(extra) + else: + break + + batch_request_ids: List[Any] = [] + batch_engine_inputs: List[Any] = [] + _rx_bytes_by_rid: Dict[Any, int] = {} + _rx_decode_ms_by_rid: Dict[Any, float] = {} + _in_flight_ms_by_rid: Dict[Any, float] = {} + for t in batch_tasks: + rid = t["request_id"] + try: + sent_ts = float(t.get("sent_ts", None)) if isinstance(t, dict) else None + if sent_ts is not None: + _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0 + else: + _in_flight_ms_by_rid[rid] = 0.0 + except Exception: + _in_flight_ms_by_rid[rid] = 0.0 + _t0 = _time.time() + ein = maybe_load_from_ipc(t, obj_key="engine_inputs", shm_key="engine_inputs_shm") + _t1 = _time.time() + _rx_decode_ms_by_rid[rid] = (_t1 - _t0) * 1000.0 + try: + if "engine_inputs_shm" in t: + _rx_bytes_by_rid[rid] = int(t["engine_inputs_shm"]["size"]) # type: ignore[index] + else: + _rx_bytes_by_rid[rid] = len(_ser(ein)) + except Exception: + _rx_bytes_by_rid[rid] = 0 + batch_request_ids.append(rid) + if isinstance(ein, list): + batch_engine_inputs.extend(ein) + elif isinstance(ein, dict): + batch_engine_inputs.append(ein) + else: + _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein)) + sampling_params = batch_tasks[0]["sampling_params"] + _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=%d, request_ids=%s", stage_id, len(batch_tasks), batch_request_ids) + print("--------------------------------", flush=True) + print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True) + print("--------------------------------", flush=True) + try: + gen_outputs: List[Any] = [] + _gen_t0 = _time.time() + for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False): + gen_outputs.append(ro) + _gen_t1 = _time.time() + _gen_ms = (_gen_t1 - _gen_t0) * 1000.0 + + # Group outputs per request id with fallback + req_to_outputs: Dict[Any, List[Any]] = {rid: [] for rid in batch_request_ids} + unmapped: List[Any] = [] + for ro in gen_outputs: + rid = getattr(ro, "request_id", None) + if rid in req_to_outputs: + req_to_outputs[rid].append(ro) + else: + unmapped.append(ro) + if unmapped: + idx = 0 + for ro in unmapped: + target_rid = batch_request_ids[idx % len(batch_request_ids)] + req_to_outputs[target_rid].append(ro) + idx += 1 + + # Per-request stats logging and aggregates + def _count_tokens(_ros: List[Any]) -> int: + total = 0 + for _ro in _ros: + try: + outs = getattr(_ro, "outputs", None) + if outs and len(outs) > 0: + tokens = getattr(outs[0], "token_ids", None) + if tokens is not None: + total += len(tokens) + except Exception: + pass + return total + + for rid in batch_request_ids: + _r_outputs = req_to_outputs.get(rid, []) + _num_tokens = _count_tokens(_r_outputs) + _agg_total_tokens += _num_tokens + _agg_total_gen_time_ms += _gen_ms + _tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0 + if _stats_file: + try: + _append_jsonl(_stats_file, { + "type": "stage_request_stats", + "stage_id": stage_id, + "request_id": rid, + "batch_size": len(batch_tasks), + "num_tokens_out": int(_num_tokens), + "stage_gen_time_ms": _gen_ms, + "tokens_per_s": _tokens_per_s, + "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), + "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), + "rx_mbps": (float(_rx_bytes_by_rid.get(rid, 0)) * 8.0) / (max(float(_rx_decode_ms_by_rid.get(rid, 0.0)), 1e-6) * 1000.0), + }) + except Exception: + pass + + if _stats_file: + try: + _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0 + _append_jsonl(_stats_file, { + "type": "stage_running_avg", + "stage_id": stage_id, + "total_tokens": int(_agg_total_tokens), + "total_gen_time_ms": _agg_total_gen_time_ms, + "avg_tokens_per_s": _avg_tokens_per_s, + }) + _append_jsonl(_stats_file, { + "type": "stage_batch_stats", + "stage_id": stage_id, + "batch_size": len(batch_tasks), + "batch_gen_time_ms": _gen_ms, + "request_ids": list(batch_request_ids), + }) + except Exception: + pass + + # Emit per-request results + for rid in batch_request_ids: + r_outputs = req_to_outputs.get(rid, []) + try: + use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes) + _metrics = { + "num_tokens_out": int(_count_tokens(r_outputs)), + "stage_gen_time_ms": _gen_ms, + "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), + "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), + "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), + } + if use_shm: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "engine_outputs_shm": payload, + "metrics": _metrics, + }) + else: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "engine_outputs": payload, + "metrics": _metrics, + }) + except Exception: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "engine_outputs": r_outputs, + "metrics": { + "num_tokens_out": int(_count_tokens(r_outputs)), + "stage_gen_time_ms": _gen_ms, + "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)), + "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)), + "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)), + }, + }) + _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid) + except Exception as e: + _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e) + for rid in batch_request_ids: + out_q.put({ + "request_id": rid, + "stage_id": stage_id, + "error": str(e), + }) diff --git a/vllm_omni/entrypoints/pipeline_utils.py b/vllm_omni/entrypoints/pipeline_utils.py new file mode 100644 index 000000000..cbf036332 --- /dev/null +++ b/vllm_omni/entrypoints/pipeline_utils.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from typing import Any, Dict, Optional, Tuple, Union + +import logging +import json +import os +import pickle +from multiprocessing import shared_memory as _shm + +import cloudpickle + +logger = logging.getLogger(__name__) + + +def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> None: + """Configure per-stage CUDA visibility and current device. + + Behavior + - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly + to this list; logical index 0 is used as current device. + - Integer or digit-string: treat as logical index (0-based) into the current + CUDA_VISIBLE_DEVICES mapping; map to the physical device, and then set + CUDA_VISIBLE_DEVICES to this single device. + - None/"cpu": keep default visibility. + - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string. + """ + try: + selected_physical: Optional[int] = None + logical_idx: Optional[int] = None + + if isinstance(devices, str) and "," in devices: + os.environ["CUDA_VISIBLE_DEVICES"] = devices + toks = [t.strip() for t in devices.split(",") if t.strip() != ""] + if toks: + try: + selected_physical = int(toks[0]) + logger.debug( + "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s", + stage_id, devices, selected_physical, + ) + except Exception: + selected_physical = None + elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()): + logical_idx = max(0, int(devices)) + vis = os.environ.get("CUDA_VISIBLE_DEVICES") + if vis: + try: + mapping = [int(x) for x in vis.split(",") if x.strip() != ""] + if 0 <= logical_idx < len(mapping): + selected_physical = mapping[logical_idx] + except Exception: + selected_physical = None + if selected_physical is None: + selected_physical = int(logical_idx) + os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical) + logger.debug( + "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device", + stage_id, logical_idx + 1, selected_physical, + ) + elif devices in (None, "cpu"): + logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices) + else: + selected_physical = int(str(devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical) + logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical) + + try: + import torch # noqa: WPS433 + if torch.cuda.is_available(): + try: + torch.cuda.set_device(0) + except Exception: + pass + num = torch.cuda.device_count() + info = [] + for i in range(num): + total = torch.cuda.get_device_properties(i).total_memory + free, _ = torch.cuda.mem_get_info(i) + info.append({ + "idx": i, + "name": torch.cuda.get_device_name(i), + "total": int(total), + "free": int(free), + }) + logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info) + except Exception: + pass + except Exception as e: + logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e) + + +def serialize_obj(obj: Any) -> bytes: + """Serialize a Python object to bytes using cloudpickle.""" + return cloudpickle.dumps(obj) + + +def shm_write_bytes(payload: bytes) -> Dict[str, Any]: + """Write bytes into SharedMemory and return meta dict {name,size}. + + Caller should close the segment; the receiver should unlink. + """ + shm = _shm.SharedMemory(create=True, size=len(payload)) + mv = memoryview(shm.buf) + mv[: len(payload)] = payload + del mv + meta = {"name": shm.name, "size": len(payload)} + try: + shm.close() + except Exception: + pass + return meta + + +def shm_read_bytes(meta: Dict[str, Any]) -> bytes: + """Read bytes from SharedMemory by meta {name,size} and cleanup.""" + shm = _shm.SharedMemory(name=meta["name"]) # type: ignore[index] + mv = memoryview(shm.buf) + data = bytes(mv[: meta["size"]]) + del mv + try: + shm.close() + except Exception: + pass + try: + shm.unlink() + except Exception: + pass + return data + + +def _ensure_parent_dir(path: str) -> None: + """Ensure the parent directory for a file path exists (best-effort).""" + try: + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + except Exception: + pass + + +def append_jsonl(path: str, record: Dict[str, Any]) -> None: + """Append a JSON record as one line to a JSONL file (best-effort). + + This is safe to call from multiple processes when each process writes + to a distinct file. For concurrent writes to the same file, OS append + semantics typically suffice, but no additional locking is provided. + """ + try: + _ensure_parent_dir(path) + line = json.dumps(record, ensure_ascii=False) + fd = os.open(path, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o644) + with os.fdopen(fd, "a", encoding="utf-8") as f: + f.write(line + "\n") + except Exception: + logger.exception("Failed to append JSONL to %s", path) + + +def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]: + """Dump object to SHM if serialized size exceeds threshold. + + Returns (True, meta) when dumped; otherwise (False, original_obj). + """ + payload = serialize_obj(obj) + if len(payload) > threshold: + return True, shm_write_bytes(payload) + return False, obj + + +def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -> Any: + """Load object from container that may carry SHM or inline object.""" + if shm_key in container: + return pickle.loads(shm_read_bytes(container[shm_key])) + return container[obj_key] + + +def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]: + """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key). + + When serialized size exceeds threshold, returns {shm_key: {name,size}}; + otherwise returns {obj_key: obj}. + """ + payload: Dict[str, Any] = {} + use_shm, data = maybe_dump_to_shm(obj, threshold) + if use_shm: + payload[shm_key] = data + else: + payload[obj_key] = data + return payload + + +# Convert OmegaConf/objects to plain dicts +def _to_dict(x: Any) -> Dict[str, Any]: + try: + if isinstance(x, dict): + return dict(x) + return OmegaConf.to_container(x, resolve=True) # type: ignore[arg-type] + except Exception: + try: + return dict(x) + except Exception: + return {} diff --git a/vllm_omni/entrypoints/pipelined_omni_llm.py b/vllm_omni/entrypoints/pipelined_omni_llm.py new file mode 100644 index 000000000..b874c2ee1 --- /dev/null +++ b/vllm_omni/entrypoints/pipelined_omni_llm.py @@ -0,0 +1,621 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Sequence, Union + +import os +import sys +import queue +import time +import multiprocessing as mp +import logging + +from vllm.inputs import PromptType +from vllm.logger import init_logger +from vllm.sampling_params import SamplingParams + +from vllm_omni.entrypoints.omni_llm import OmniLLM, OmniStageLLM +from vllm_omni.entrypoints.omni_stage import OmniStage +from vllm_omni.entrypoints.utils import load_stage_configs_from_model +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.entrypoints.pipeline_utils import ( + maybe_load_from_ipc as _load, + encode_for_ipc as _encode, + serialize_obj as _ser, + append_jsonl as _append_jsonl, +) + +logger = init_logger(__name__) + + +class PipelinedOmniLLM(OmniLLM): + """Multi-process pipelined OmniLLM. + + - Per-stage process with copy-based IPC (Queues) + - window=-1 across stages (downstream starts when upstream finishes) + - max_inflight=1 per stage (serial within a stage), but pipeline across stages + """ + + def __init__(self, model: str, + stage_configs=None, + log_stats: bool = True, + log_file: Optional[str] = None, + init_sleep_seconds: int = 20, + shm_threshold_bytes: int = 65536, + batch_timeout: int = 10, + init_timeout: int = 300, + **kwargs): + self.batch_timeout = batch_timeout + self._enable_stats: bool = bool(log_stats) + # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent. + if stage_configs is None: + self.stage_configs = load_stage_configs_from_model(model) + else: + self.stage_configs = stage_configs + + self.stage_list: List[OmniStage] = [] + for stage_config in self.stage_configs: + # Only construct lightweight OmniStage without setting engine in parent + stage = OmniStage(stage_config) + self.stage_list.append(stage) + logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) + + self._ctx = mp.get_context("spawn") + self._stage_in_queues: List[mp.Queue] = [] + self._stage_out_queues: List[mp.Queue] = [] + # Optional file handler for orchestrator + self._log_file = log_file + if self._log_file: + try: + # Avoid duplicate handlers + has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers) + if not has_file_handler: + fh = logging.FileHandler(self._log_file) + fh.setLevel(logging.DEBUG) + fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s")) + logger.addHandler(fh) + logger.setLevel(logging.DEBUG) + except Exception: + pass + + # Orchestrator stats JSONL file + self._stats_file: Optional[str] = None + if self._enable_stats and self._log_file: + try: + self._stats_file = f"{self._log_file}.orchestrator.stats.jsonl" + except Exception: + self._stats_file = None + # Overall stats JSONL file (per-request + summary) + self._overall_stats_file: Optional[str] = None + if self._enable_stats and self._log_file: + try: + self._overall_stats_file = f"{self._log_file}.overall.stats.jsonl" + except Exception: + self._overall_stats_file = None + + self._init_sleep_seconds = max(0, int(init_sleep_seconds)) + self._shm_threshold_bytes = max(0, int(shm_threshold_bytes)) + self._start_stage_processes(model) + # Wait for all stages to report readiness before seeding + self._stages_ready: set[int] = set() + self._wait_for_stages_ready(timeout=init_timeout) + + def _start_stage_processes(self, model: str) -> None: + for stage_id, stage in enumerate(self.stage_list): + # Use unbounded queues to avoid deadlock when seeding many requests + in_q: mp.Queue = self._ctx.Queue(maxsize=0) + out_q: mp.Queue = self._ctx.Queue(maxsize=0) + self._stage_in_queues.append(in_q) + self._stage_out_queues.append(out_q) + + # Attach queues and start Stage-owned worker process + stage.attach_queues(in_q, out_q) + stage.init_stage_worker( + model, + log_file=self._log_file, + shm_threshold_bytes=self._shm_threshold_bytes, + ctx=self._ctx, + batch_timeout=self.batch_timeout, + ) + logger.debug("[Orchestrator] Stage-%s process started", stage_id) + + def close(self) -> None: + for q in self._stage_in_queues: + try: + q.put_nowait(None) + except Exception: + pass + for stage in self.stage_list: + try: + stage.stop_stage_worker() + except Exception: + pass + + def __del__(self) -> None: # best-effort + try: + self.close() + except Exception: + pass + + def generate( + self, + prompts: Union[PromptType, Sequence[PromptType]], + sampling_params_list: Optional[ + Union[SamplingParams, Sequence[SamplingParams]] + ] = None, + ) -> List[OmniRequestOutput]: + logger.debug("[Orchestrator] generate() called") + if sampling_params_list is None: + raise ValueError("sampling_params_list is required for pipelined generation") + if len(sampling_params_list) != len(self.stage_list): + raise ValueError( + f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}" + ) + + # Normalize prompts to a list for per-request iteration + if not isinstance(prompts, (list, tuple)): + request_prompts: List[PromptType] = [prompts] + else: + request_prompts = list(prompts) + + final_outputs: List[OmniRequestOutput] = [] + + # Orchestrator keeps stage objects for input derivation + num_stages = len(self.stage_list) + + # Map from request_id to original prompt + request_id_to_prompt: Dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)} + + # Track per-request start time for end-to-end timing + _req_start_ts: Dict[int, float] = {} + _wall_start_ts: float = time.time() + _last_finish_ts: float = _wall_start_ts + + # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage) + final_stage_id_for_e2e = -1 + try: + for _sid, _st in enumerate(self.stage_list): + if getattr(_st, "final_output", False): + final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid) + if final_stage_id_for_e2e < 0: + final_stage_id_for_e2e = len(self.stage_list) - 1 + except Exception: + final_stage_id_for_e2e = len(self.stage_list) - 1 + # In-memory aggregators for this generate() call + stage_total_time_ms: List[float] = [0.0 for _ in range(num_stages)] + stage_total_tokens: List[int] = [0 for _ in range(num_stages)] + stage_req_counts: List[int] = [0 for _ in range(num_stages)] + transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {} + # Per-edge per-request sender timing to combine with receiver timing later + transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {} + e2e_total_ms: float = 0.0 + e2e_total_tokens: int = 0 + e2e_count: int = 0 + e2e_done: set[int] = set() + # Per-request overall aggregation + per_request: Dict[int, Dict[str, Any]] = {} + sum_per_request_transfer_ms: float = 0.0 + + # Seed stage-0 queue with all requests + logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)) + + for req_id, prompt in request_id_to_prompt.items(): + sp0: SamplingParams = sampling_params_list[0] # type: ignore[index] + task = { + "request_id": req_id, + "engine_inputs": prompt, + "sampling_params": sp0, + } + self.stage_list[0].submit(task) + _req_start_ts[req_id] = time.time() + logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id) + + # For each stage, forward results to next stage; collect finals at the end + # We pipeline by continually polling output queues in stage order + remaining_by_stage: List[int] = [len(request_prompts)] + [0] * (num_stages - 1) + completed_requests = 0 + total_requests = len(request_prompts) + + logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages) + while completed_requests < total_requests: + made_progress = False + for stage_id, stage in enumerate(self.stage_list): + result = stage.try_collect() + if result is None: + continue + + made_progress = True + req_id = result.get("request_id") + if "error" in result: + logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"]) + continue + + if result.get("type") == "stage_ready": + #Only happens when stage is initialized slower than expected, so we wait for a short time and try again + time.sleep(0.05) + continue + + engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm") + # Aggregate per-request metrics from stage worker if present + try: + _m = result.get("metrics") + if _m is not None: + stage_req_counts[stage_id] += 1 + stage_total_time_ms[stage_id] += float(_m.get("stage_gen_time_ms", 0.0)) + stage_total_tokens[stage_id] += int(_m.get("num_tokens_out", 0)) + # record per-request stage metrics + try: + rid_int = int(req_id) + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr_stages = pr["stages"] # type: ignore[index] + pr_stages[stage_id] = { + "stage_gen_time_ms": float(_m.get("stage_gen_time_ms", 0.0)), + "num_tokens_out": int(_m.get("num_tokens_out", 0)), + } + except Exception: + pass + # Also aggregate receiver-side transfer decode time for (stage_id-1 -> stage_id) + try: + if stage_id > 0: + key = (stage_id - 1, stage_id) + agg = transfer_agg.get(key) + if agg is None: + agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, + "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, + "sum_total_ms": 0.0, "total_count": 0.0} + transfer_agg[key] = agg + rx_b = float(_m.get("rx_transfer_bytes", 0.0)) + rx_ms = float(_m.get("rx_decode_time_ms", 0.0)) + in_flight_ms = float(_m.get("rx_in_flight_time_ms", 0.0)) + agg["sum_rx_bytes"] += rx_b + agg["sum_rx_ms"] += rx_ms + agg["rx_count"] += 1.0 + # If we have sender-side timing stored, compute combined (encode+enqueue + decode) + try: + edge_req = (stage_id - 1, stage_id, int(req_id)) + s = transfer_edge_req.get(edge_req) + if s is not None: + total_ms = float(s.get("tx_ms", 0.0)) + in_flight_ms + rx_ms + agg["sum_total_ms"] += total_ms + agg["total_count"] += 1.0 + # accumulate per-request transfer totals + try: + rid_int = int(req_id) + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index] + pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_b) # type: ignore[index] + except Exception: + pass + if getattr(self, "_stats_file", None): + try: + size_b = float(s.get("size_bytes", rx_b)) + _append_jsonl(self._stats_file, { # type: ignore[arg-type] + "type": "transfer_total_stats", + "from_stage": stage_id - 1, + "to_stage": stage_id, + "request_id": req_id, + "size_bytes": int(size_b), + "tx_time_ms": float(s.get("tx_ms", 0.0)), + "in_flight_time_ms": in_flight_ms, + "rx_decode_time_ms": rx_ms, + "total_time_ms": total_ms, + "total_time_per_kb_ms": total_ms / max(size_b / 1024.0, 1e-6) if size_b > 0 else 0.0, + }) + except Exception: + pass + except Exception: + pass + # Emit per-request RX stats to JSONL + if getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, { # type: ignore[arg-type] + "type": "transfer_rx_stats", + "from_stage": stage_id - 1, + "to_stage": stage_id, + "request_id": req_id, + "rx_bytes": int(rx_b), + "rx_decode_time_ms": rx_ms, + "in_flight_time_ms": in_flight_ms, + "rx_time_per_kb_ms": rx_ms / max(rx_b / 1024.0, 1e-6) if rx_b > 0 else 0.0, + }) + except Exception: + pass + except Exception: + pass + except Exception: + pass + logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id) + stage.set_engine_outputs(engine_outputs) + + if getattr(stage, "final_output", False): + final_outputs.append( + OmniRequestOutput( + stage_id=stage_id, + final_output_type=stage.final_output_type, # type: ignore[attr-defined] + request_output=engine_outputs, + ) + ) + logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id) + + # End-to-end timing and time-per-token for final output (only once per request at the designated final stage) + try: + if stage_id == final_stage_id_for_e2e and req_id not in e2e_done: + _t0 = _req_start_ts.get(req_id) + if _t0 is not None: + _t1 = time.time() + _last_finish_ts = max(_last_finish_ts, _t1) + _e2e_ms = (_t1 - _t0) * 1000.0 + # Count tokens from final stage outputs + def _count_tokens(_ros: List[Any]) -> int: # type: ignore[name-defined] + total = 0 + for _ro in _ros: + try: + outs = getattr(_ro, "outputs", None) + if outs and len(outs) > 0: + tokens = getattr(outs[0], "token_ids", None) + if tokens is not None: + total += len(tokens) + except Exception: + pass + return total + _num_tokens = _count_tokens(engine_outputs) + _time_per_token_ms = (_e2e_ms / _num_tokens) if _num_tokens > 0 else 0.0 + # Update E2E aggregators + e2e_total_ms += _e2e_ms + e2e_total_tokens += int(_num_tokens) + e2e_count += 1 + e2e_done.add(req_id) + # Write per-request overall record + try: + rid_int = int(req_id) + pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) + per_req_record = { + "type": "overall_request", + "request_id": rid_int, + "e2e_time_ms": _e2e_ms, + "num_tokens_out": int(_num_tokens), + "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)), + "transfers_total_bytes": int(pr.get("transfers_bytes", 0)), + "stages": pr.get("stages", {}), + } + sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0)) + if getattr(self, "_overall_stats_file", None): + try: + _append_jsonl(self._overall_stats_file, per_req_record) # type: ignore[arg-type] + except Exception: + pass + except Exception: + pass + if getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, { # type: ignore[arg-type] + "type": "orchestrator_request_e2e", + "request_id": req_id, + "final_stage_id": stage_id, + "e2e_time_ms": _e2e_ms, + "num_tokens_out": int(_num_tokens), + "e2e_time_per_token_ms": _time_per_token_ms, + }) + except Exception: + pass + except Exception: + pass + + next_stage_id = stage_id + 1 + if next_stage_id < num_stages: + next_stage: OmniStage = self.stage_list[next_stage_id] + next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) + sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index] + try: + # Measure transfer size and time (encode + enqueue) + size_bytes = 0 + try: + size_bytes = len(_ser(next_inputs)) + except Exception: + size_bytes = 0 + t0 = time.time() + ipc_payload = _encode( + next_inputs, + getattr(self, "_shm_threshold_bytes", 65536), + obj_key="engine_inputs", + shm_key="engine_inputs_shm", + ) + ipc_payload.update({ + "request_id": req_id, + "sampling_params": sp_next, + "sent_ts": time.time(), + }) + self.stage_list[next_stage_id].submit(ipc_payload) + t1 = time.time() + tx_ms = (t1 - t0) * 1000.0 + if self._enable_stats and getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, { + "type": "transfer_stats", + "from_stage": stage_id, + "to_stage": next_stage_id, + "request_id": req_id, + "size_bytes": int(size_bytes), + "tx_time_ms": tx_ms, + "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_ms, 1e-6) * 1000.0), + "used_shm": bool("engine_inputs_shm" in ipc_payload), + }) + except Exception: + pass + # Update in-memory transfer aggregator + try: + key = (stage_id, next_stage_id) + agg = transfer_agg.get(key) + if agg is None: + agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0, + "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0, + "sum_total_ms": 0.0, "total_count": 0.0} + transfer_agg[key] = agg + agg["sum_bytes"] += float(size_bytes) + agg["sum_ms"] += float(tx_ms) + agg["count"] += 1.0 + # Store sender-side timing for per-request combination + transfer_edge_req[(stage_id, next_stage_id, int(req_id))] = { + "tx_ms": float(tx_ms), + "size_bytes": float(size_bytes), + } + except Exception: + pass + except Exception: + self.stage_list[next_stage_id].submit({ + "request_id": req_id, + "engine_inputs": next_inputs, + "sampling_params": sp_next, + }) + logger.debug("[Orchestrator] Forwarded request %s to stage-%s", req_id, next_stage_id) + remaining_by_stage[next_stage_id] += 1 + else: + completed_requests += 1 + logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests) + + if not made_progress: + time.sleep(0.005) + logger.debug("[Orchestrator] All requests completed") + + # Summarize and print stats + try: + stage_summary: List[Dict[str, Any]] = [] + for sid in range(num_stages): + reqs = stage_req_counts[sid] + tokens = stage_total_tokens[sid] + total_ms = float(stage_total_time_ms[sid]) + avg_req = (total_ms / reqs) if reqs > 0 else 0.0 + avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 + stage_summary.append({ + "stage_id": sid, + "requests": int(reqs), + "tokens": int(tokens), + "total_time_ms": total_ms, + "avg_time_per_request_ms": avg_req, + "avg_tokens_per_s": avg_tok, + }) + + transfer_summary: List[Dict[str, Any]] = [] + for (src, dst), agg in transfer_agg.items(): + sum_bytes = float(agg.get("sum_bytes", 0.0)) + sum_ms = float(agg.get("sum_ms", 0.0)) + samples = int(agg.get("count", 0.0)) + tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0)) + sum_rx_ms = float(agg.get("sum_rx_ms", 0.0)) + samples_rx = int(agg.get("rx_count", 0.0)) + rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0 + sum_total_ms = float(agg.get("sum_total_ms", 0.0)) + samples_total = int(agg.get("total_count", 0.0)) + total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 + transfer_summary.append({ + "from_stage": src, + "to_stage": dst, + "samples": samples, + "total_bytes": int(sum_bytes), + "total_time_ms": sum_ms, + "tx_mbps": tx_mbps, + "rx_samples": samples_rx, + "rx_total_bytes": int(sum_rx_bytes), + "rx_total_time_ms": sum_rx_ms, + "rx_mbps": rx_mbps, + "total_samples": samples_total, + "total_transfer_time_ms": sum_total_ms, + "total_mbps": total_mbps, + }) + + e2e_avg_req = (e2e_total_ms / e2e_count) if e2e_count > 0 else 0.0 + e2e_avg_tok = (e2e_total_tokens * 1000.0 / e2e_total_ms) if e2e_total_ms > 0 else 0.0 + wall_time_ms = max(0.0, (_last_finish_ts - _wall_start_ts) * 1000.0) + summary: Dict[str, Any] = { + "e2e_requests": int(e2e_count), + # 按你的需求：e2e_total_time_ms 直接使用墙钟时间（非各请求之和） + "e2e_total_time_ms": float(wall_time_ms), + # 额外保留各请求 E2E 之和，供需要时参考 + "e2e_sum_time_ms": float(e2e_total_ms), + "e2e_total_tokens": int(e2e_total_tokens), + "e2e_avg_time_per_request_ms": e2e_avg_req, + "e2e_avg_tokens_per_s": e2e_avg_tok, + "wall_time_ms": wall_time_ms, + "final_stage_id": final_stage_id_for_e2e, + "stages": stage_summary, + "transfers": transfer_summary, + } + logger.info("[Summary] %s", summary) + if self._enable_stats and getattr(self, "_stats_file", None): + try: + _append_jsonl(self._stats_file, {"type": "orchestrator_summary", **summary}) # type: ignore[arg-type] + except Exception: + pass + if self._enable_stats and getattr(self, "_overall_stats_file", None): + try: + _append_jsonl(self._overall_stats_file, {"type": "overall_summary", **summary}) # type: ignore[arg-type] + except Exception: + pass + except Exception: + pass + + return final_outputs + + def _wait_for_stages_ready(self, timeout: int = 120) -> None: + deadline = time.time() + max(0, int(timeout)) + num_stages = len(self.stage_list) + while len(self._stages_ready) < num_stages and time.time() < deadline: + progressed = False + for stage_id, stage in enumerate(self.stage_list): + if stage_id in self._stages_ready: + continue + result = stage.try_collect() + if result is None: + continue + progressed = True + if result.get("type") == "stage_ready": + self._stages_ready.add(stage_id) + logger.debug("[Orchestrator] Stage-%s reported ready", stage_id) + else: + # No user data should arrive before seeding; ignore other messages + pass + if not progressed: + time.sleep(0.01) + if len(self._stages_ready) < num_stages: + not_ready = sorted(set(range(num_stages)) - set(self._stages_ready)) + logger.warning( + "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s", + len(self._stages_ready), num_stages, not_ready, + ) + # Provide actionable suggestions before shutdown + try: + suggestions = [ + "Verify GPU/device assignment in config (mrs.devices) is correct.", + "Check GPU/host memory availability; reduce model or batch size if needed.", + "Check model weights path and network reachability (if loading remotely).", + "Increase initialization wait time (init_sleep_seconds or call-site timeout).", + ] + if getattr(self, "_log_file", None): + suggestions.append( + f"Inspect per-stage log files for details: {self._log_file}.stage.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) + + diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index f820ef3d7..d7415b805 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -269,12 +269,17 @@ class Qwen2_5OmniForConditionalGeneration( # 2) Talker (if codec not provided) if self.model_stage == "talker": + # Mixed-mode support: in一个step内同时支持 Prefill*n 与 Decode*n。 + # 规则： + # - 使用特殊token将 Prefill 片段包裹为 [BOS][PAD...][EOS] + # - Decode 片段则为“单一非特殊token” + # - 若提供了 additional_information（可为按请求分割的 list 或拼接后的 tensor + 形状列表）， + # 则对 Prefill 片段按请求重建 thinker→talker 的输入嵌入； + # - 对 Decode 片段，若提供了 per-request 的 decode 辅助嵌入（可选），则进行叠加；否则保留原嵌入。 if input_ids is None and additional_information is None: input_ids = torch.zeros( - inputs_embeds.shape[0], - dtype=torch.long, - device=inputs_embeds.device, + inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device ) additional_information = {} self.thinker_reply_part = torch.zeros_like(inputs_embeds) @@ -282,63 +287,109 @@ class Qwen2_5OmniForConditionalGeneration( else: is_profile = False + # Ensure we have base embeddings when only ids are provided + if inputs_embeds is None and input_ids is not None: + inputs_embeds = self.talker.get_input_embeddings(input_ids) + + # ------- Request-scoped additional information (no cross-request concat) ------- + request_ids: Optional[list[str]] = kwargs.get("request_ids") # ordered + request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans") + addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id") + runtime_addi = kwargs.get("runtime_additional_information") + + # Normalize runtime_addi into a mapping by request_id for convenience + runtime_addi_by_req: dict[str, dict] = {} + if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids): + for i, rid in enumerate(request_ids): + if isinstance(rid, str) and isinstance(runtime_addi[i], dict): + runtime_addi_by_req[rid] = runtime_addi[i] + elif isinstance(request_ids, list) and isinstance(runtime_addi, dict): + for rid in request_ids: + if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict): + runtime_addi_by_req[rid] = runtime_addi[rid] + + # Containers to return per-request updates (e.g., thinker_reply_part_per_request) + update_by_req_id: dict[str, dict] = {} + + # ------- Prefill: span_len > 1 ------- if ( - input_ids is not None - and additional_information is not None - and not is_profile + not is_profile + and isinstance(request_ids, list) + and isinstance(request_token_spans, list) + and isinstance(addi_by_req, dict) ): - # read from additional_information dict - thinker_result = None - if additional_information is not None and isinstance( - additional_information, dict - ): - thinker_result = additional_information.get("thinker_result") - prompt_embeds = additional_information.get("prompt_embeds") - prompt_token_ids = additional_information.get("prompt_token_ids") - thinker_output_token_ids = additional_information.get( - "thinker_output_token_ids" - ) - else: - thinker_result = torch.zeros_like(inputs_embeds) - prompt_embeds = torch.zeros_like(inputs_embeds) - prompt_token_ids = torch.zeros( - inputs_embeds.shape[0], - dtype=torch.int64, - device=inputs_embeds.device, - ) - thinker_output_token_ids = torch.zeros( - inputs_embeds.shape[0], - dtype=torch.int64, - device=inputs_embeds.device, - ) + for idx_req, rid in enumerate(request_ids): + s, e = request_token_spans[idx_req] + span_len = int(e) - int(s) + if span_len <= 1: + continue + info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {} + if not isinstance(info, dict): + info = {} + pe = info.get("prompt_embeds") # Tensor [P,H] + tr = info.get("thinker_result") # Tensor [K,H] + ptoks = info.get("prompt_token_ids") # list[int] + otoks = info.get("thinker_output_token_ids") # list[int] - if thinker_result is None: - thinker_result = torch.zeros_like(inputs_embeds) - self.thinker_reply_part = thinker_result.squeeze(0) - if self.thinker_reply_part.shape[1] > 1: - self.thinker_reply_part = self.thinker_reply_part[1:, :] - input_ids, inputs_embeds = self._thinker_to_talker_prefill( - voice_type=voice_type, - output_prompt_embeds=thinker_result, - output_token_ids=thinker_output_token_ids, - thinker_prompt_embeds=prompt_embeds, - prompt_token_ids=prompt_token_ids, - ) - elif not is_profile: - input_ids, inputs_embeds = self._thinker_to_talker_decode_one_step( - output_prompt_embeds=( - self.thinker_reply_part[:1] - if self.thinker_reply_part.shape[0] >= 1 - else torch.zeros(1, self.thinker_reply_part.shape[1]) - .to(self._module_device(self.model)) - .to(torch.bfloat16) - + (-1.25 * 2 ** (-123)) - ), - output_token_ids=input_ids, - ) + if not isinstance(pe, torch.Tensor): + pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + if not isinstance(tr, torch.Tensor): + tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + if not isinstance(ptoks, (list, torch.Tensor)): + ptoks = [] + if not isinstance(otoks, (list, torch.Tensor)): + otoks = [] - if self.thinker_reply_part.shape[0] >= 1: - self.thinker_reply_part = self.thinker_reply_part[1:, :] + req_input_ids, req_embeds = self._thinker_to_talker_prefill( + voice_type=voice_type, + output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_token_ids=otoks, + thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)), + prompt_token_ids=ptoks, + ) + seg_len = min(span_len, req_embeds.shape[0]) + inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] + if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: + input_ids[s : s + seg_len] = req_input_ids + + # Prepare per-request reply queue for subsequent decode: drop first row + if tr.ndim == 2 and tr.shape[0] > 0: + update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous() + + # ------- Decode: span_len == 1 ------- + if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list): + for idx_req, rid in enumerate(request_ids): + s, e = request_token_spans[idx_req] + if (int(e) - int(s)) != 1: + continue + # choose step vector in priority order + step_vec = None + # A) runtime queue + q = None + if isinstance(rid, str): + q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request") + if isinstance(q, torch.Tensor) and q.numel() > 0: + step_vec = q[0:1] + new_q = q[1:].detach().to("cpu").contiguous() + update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q + else: + # B) per-request provided decode vector (optional) + info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {} + dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None + if isinstance(dv, torch.Tensor) and dv.numel() > 0: + step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1) + elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0: + # C) fallback shared pool + step_vec = self.thinker_reply_part[0:1] + self.thinker_reply_part = self.thinker_reply_part[1:] + + if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0: + one_id = input_ids[s : s + 1] + _, one_embed = self._thinker_to_talker_decode_one_step( + output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_token_ids=one_id, + ) + inputs_embeds[s] = one_embed[0] with torch.inference_mode(): talker_hidden = self.talker( @@ -346,8 +397,19 @@ class Qwen2_5OmniForConditionalGeneration( positions=positions[0], inputs_embeds=inputs_embeds, ) + multimodal_outputs: dict = None + # Return updates if any + if update_by_req_id: + multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id} - return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=None) + if sampling_metadata is not None: + # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size + sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448 + + return OmniOutput( + text_hidden_states=talker_hidden, + multimodal_outputs=multimodal_outputs, + ) if self.model_stage == "code2wav": code = ( @@ -359,8 +421,12 @@ class Qwen2_5OmniForConditionalGeneration( device=inputs_embeds.device, ) ) + + code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code + code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code + audio_tensor = self.generate_audio( - code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code, voice_type + code, voice_type ) return OmniOutput( text_hidden_states=None, multimodal_outputs={"audio": audio_tensor} diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index 337194f9d..2f69b3080 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -1,36 +1,50 @@ # stage config for running qwen2.5-omni with architecture of OmniLLM. stage_args: - stage_id: 0 + mrs: + process: true # Run this stage in a separate process + devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) + max_batch_size: 1 engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler - gpu_memory_utilization: 0.32 + gpu_memory_utilization: 0.8 enforce_eager: true # need to discuss trust_remote_code: true engine_output_type: latent # change the param name,such as pooling_output + enable_prefix_caching: false final_output: true final_output_type: text - stage_id: 1 + mrs: + process: true + devices: "1" + max_batch_size: 3 engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler - gpu_memory_utilization: 0.32 + gpu_memory_utilization: 0.8 enforce_eager: true trust_remote_code: true + enable_prefix_caching: false engine_output_type: latent engine_input_source: [0] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker - stage_id: 2 + mrs: + process: true + devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU + max_batch_size: 1 engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler - gpu_memory_utilization: 0.3 + gpu_memory_utilization: 0.15 enforce_eager: true trust_remote_code: true enable_prefix_caching: false @@ -38,3 +52,17 @@ stage_args: engine_input_source: [1] final_output: true final_output_type: audio + +# Top-level MRS config (concise): default windows and stage edges +mrs: + enabled: true + defaults: + window_size: -1 # Simplified: trigger downstream only after full upstream completion + max_inflight: 1 # Simplified: process serially within each stage + edges: + - from: 0 # thinker → talker: trigger only after receiving full input (-1) + to: 1 + window_size: -1 + - from: 1 # talker → code2wav: trigger only after receiving full input (-1) + to: 2 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 4dc10d4fd..ca5e02ba6 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/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] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 20a98e55a..2b1ce3509 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -99,9 +99,11 @@ class GPUARModelRunner(OmniGPUModelRunner): # TODO(woosuk): Avoid the copy. Optimize. self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) - # Reset per-step additional information collector + # Reset per-step additional information collector (deprecated concat path) if hasattr(self, "_forward_additional_information"): self._forward_additional_information = None + # New: per-request additional information for this step + per_req_additional_information: dict[str, dict] = {} # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only @@ -126,50 +128,27 @@ class GPUARModelRunner(OmniGPUModelRunner): self.inputs_embeds[start_offset : start_offset + overlay_len].copy_( src ) - # For additional_information: handle arbitrary keys + # Build per-request additional information (no cross-request concat) if addi_cpu is not None and isinstance(addi_cpu, dict): - # Lazy init collector dict - if ( - not hasattr(self, "_forward_additional_information") - or self._forward_additional_information is None - ): - self._forward_additional_information = {} - # Process tensors (slice by scheduled prompt range) and - # lists (append per-request) + req_info: dict[str, object] = {} for k, v in addi_cpu.items(): if isinstance(v, torch.Tensor): - # Slice along token dimension for prefill part - try: - seg = v[ - num_computed_tokens : num_computed_tokens - + overlay_len - ].to( - dtype=self.dtype, - device=self.device, - non_blocking=True, - ) - except Exception: - # Fallback: move whole tensor if slicing fails - seg = v.to( - dtype=self.dtype, - device=self.device, - non_blocking=True, - ) - prev_val = self._forward_additional_information.get(k) - self._forward_additional_information[k] = ( - torch.cat([prev_val, seg], dim=0) - if isinstance(prev_val, torch.Tensor) - else seg.clone() - ) - elif isinstance(v, list): - prev_val = self._forward_additional_information.get(k) - if prev_val is None: - self._forward_additional_information[k] = [v] - elif isinstance(prev_val, list): - self._forward_additional_information[k].append(v) + # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor + if overlay_len > 0: + try: + seg = v[ + num_computed_tokens : num_computed_tokens + overlay_len + ].detach().to("cpu").contiguous() + except Exception: + seg = v.detach().to("cpu").contiguous() + req_info[k] = seg else: - # Mixed types: wrap existing into list - self._forward_additional_information[k] = [prev_val, v] + req_info[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, list): + req_info[k] = v + else: + req_info[k] = v + per_req_additional_information[req_id] = req_info input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] @@ -395,15 +374,33 @@ class GPUARModelRunner(OmniGPUModelRunner): ): model_kwargs_extra = {} - # Only pass additional_information for the prefill part - if ( - hasattr(self, "_forward_additional_information") - and self._forward_additional_information is not None - and isinstance(self._forward_additional_information, dict) - ): - model_kwargs_extra["additional_information"] = ( - self._forward_additional_information + # Pass per-request additional information map for this step (no concat) + if per_req_additional_information: + model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information + # 始终传递每请求的运行期 additional_information（持久化在 request state） + try: + per_req_runtime_info = [] + for req_id in self.input_batch.req_ids: + req_state = self.requests.get(req_id) + info = ( + getattr(req_state, "additional_information_cpu", None) + if req_state is not None + else None + ) + per_req_runtime_info.append(info if isinstance(info, dict) else {}) + model_kwargs_extra["runtime_additional_information"] = ( + per_req_runtime_info ) + model_kwargs_extra["request_ids"] = self.input_batch.req_ids + # 传递每个请求在本步扁平序列中的 token 区间，便于模型按请求映射 decode/prefill + req_token_spans = [] + for req_index in range(len(self.input_batch.req_ids)): + start_offset = int(self.query_start_loc_cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + req_token_spans.append((start_offset, start_offset + sched_tokens)) + model_kwargs_extra["request_token_spans"] = req_token_spans + except Exception: + pass model_output = self.model( input_ids=input_ids, positions=positions, @@ -428,8 +425,42 @@ class GPUARModelRunner(OmniGPUModelRunner): hidden_states, multimodal_outputs = self.extract_multimodal_outputs( hidden_states ) - + # 模型侧可能返回 per-request 的 additional_information 更新（与模型无关的通用通道）。 + # 约定：multimodal_outputs["additional_information_update"] 为按 batch 顺序的 list[dict]， + # runner 将其合并到对应 request 的 additional_information_cpu，供后续 decode 使用。 + try: + if ( + isinstance(multimodal_outputs, dict) + and ( + "additional_information_update" in multimodal_outputs + or "additional_information_update_by_req_id" in multimodal_outputs + ) + ): + # 方案 A：按 batch 顺序的 list[dict] + updates_list = multimodal_outputs.get( + "additional_information_update" + ) + if isinstance(updates_list, list): + for idx, upd in enumerate(updates_list): + if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): + continue + req_id = self.input_batch.req_ids[idx] + self._merge_additional_information_update(req_id, upd) + # 方案 B：按 req_id 的 dict[str, dict] + updates_map = multimodal_outputs.get( + "additional_information_update_by_req_id" + ) + if isinstance(updates_map, dict): + for req_id, upd in updates_map.items(): + if not isinstance(upd, dict): + continue + if req_id not in self.requests: + continue + self._merge_additional_information_update(req_id, upd) + except Exception as e: + logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}") if not self.broadcast_pp_output: + # Common case. if not get_pp_group().is_last_rank: # Return the intermediate tensors. @@ -589,3 +620,26 @@ class GPUARModelRunner(OmniGPUModelRunner): invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, ) + + def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: + req_state = self.requests.get(req_id) + if req_state is None: + return + existing = getattr(req_state, "additional_information_cpu", {}) + if not isinstance(existing, dict): + existing = {} + merged = dict(existing) + for k, v in upd.items(): + if isinstance(v, torch.Tensor): + merged[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, list): + new_list = [] + for item in v: + if isinstance(item, torch.Tensor): + new_list.append(item.detach().to("cpu").contiguous()) + else: + new_list.append(item) + merged[k] = new_list + else: + merged[k] = v + setattr(req_state, "additional_information_cpu", merged) \ No newline at end of file