From f22944eea41d79f367a9785f4a3a174cf3428c57 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Fri, 12 Dec 2025 01:52:53 -0500 Subject: [PATCH] [Doc] fix docs Feature Design and Module Design (#283) Signed-off-by: hsliu Signed-off-by: Hongsheng Liu Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- docs/.nav.yml | 10 +- docs/api/README.md | 26 +-- docs/contributing/profiling.md | 3 + ...r_design.md => disaggregated_inference.md} | 2 +- .../design/feature/multi_request_streaming.md | 47 +++++ .../feature/multi_request_streaming_design.md | 188 ------------------ docs/design/feature/ray_based_execution.md | 2 +- docs/design/index.md | 10 +- .../{ar_module_design.md => ar_module.md} | 0 .../{dit_module_design.md => dit_module.md} | 0 ..._module_design.md => entrypoint_module.md} | 0 vllm_omni/core/dit_cache_manager.py | 13 -- 12 files changed, 71 insertions(+), 230 deletions(-) create mode 100644 docs/contributing/profiling.md rename docs/design/feature/{omni_connector_design.md => disaggregated_inference.md} (98%) create mode 100644 docs/design/feature/multi_request_streaming.md delete mode 100644 docs/design/feature/multi_request_streaming_design.md rename docs/design/module/{ar_module_design.md => ar_module.md} (100%) rename docs/design/module/{dit_module_design.md => dit_module.md} (100%) rename docs/design/module/{entrypoint_module_design.md => entrypoint_module.md} (100%) delete mode 100644 vllm_omni/core/dit_cache_manager.py diff --git a/docs/.nav.yml b/docs/.nav.yml index daf4fa736..699aefc1e 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -36,13 +36,13 @@ nav: - design/index.md - design/architecture_overview.md - Feature Design: - - design/feature/multi_request_streaming_design.md - - design/feature/omni_connector_design.md + - design/feature/disaggregated_inference.md + - design/feature/multi_request_streaming.md - design/feature/ray_based_execution.md - Module Design: - - design/module/ar_module_design.md - - design/module/dit_module_design.md - - design/module/entrypoint_module_design.md + - design/module/ar_module.md + - design/module/dit_module.md + - design/module/entrypoint_module.md - Docs Guide: contributing/DOCS_GUIDE.md - API Reference: - api/README.md diff --git a/docs/api/README.md b/docs/api/README.md index 59f7060e1..197fd1568 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -1,8 +1,15 @@ # Summary -## Entry Points -Main entry points for vLLM-Omni inference and serving. +## Configuration + +Configuration classes. + +- [vllm_omni.config.model.OmniModelConfig][] + +## EntryPoints + +Main entrypoints for vLLM-Omni inference and serving. - [vllm_omni.entrypoints.async_omni.AsyncOmni][] - [vllm_omni.entrypoints.async_omni.AsyncOmniStageLLM][] @@ -48,7 +55,6 @@ Engine classes for offline and online inference. Core scheduling and caching components. -- [vllm_omni.core.dit_cache_manager.DiTCacheManager][] - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][] - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][] - [vllm_omni.core.sched.output.OmniNewRequestData][] @@ -91,28 +97,14 @@ Model execution components. - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchEmbed][] - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchMerger][] -## Configuration - -Configuration classes. - -- [vllm_omni.config.model.OmniModelConfig][] -- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][] - ## Workers Worker classes and model runners for distributed inference. - [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][] - [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][] -- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][] -- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][] - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][] - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][] - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][] - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][] - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][] -- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][] -- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][] -- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][] -- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][] -- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][] diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md new file mode 100644 index 000000000..549a13495 --- /dev/null +++ b/docs/contributing/profiling.md @@ -0,0 +1,3 @@ +# Profiling vLLM-Omni (update soon) + +Profiling is only intended for vLLM-Omni developers and maintainers to understand the proportion of time spent in different parts of the codebase. **vLLM-Omni end-users should never turn on profiling** as it will significantly slow down the inference. diff --git a/docs/design/feature/omni_connector_design.md b/docs/design/feature/disaggregated_inference.md similarity index 98% rename from docs/design/feature/omni_connector_design.md rename to docs/design/feature/disaggregated_inference.md index 442ae6279..2dae03ff8 100644 --- a/docs/design/feature/omni_connector_design.md +++ b/docs/design/feature/disaggregated_inference.md @@ -1,4 +1,4 @@ -# OmniConnector: a unified connector interface of disaggregated inference for omni-modality models +# Disaggregated Inference for Omni-Modality Models This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines. diff --git a/docs/design/feature/multi_request_streaming.md b/docs/design/feature/multi_request_streaming.md new file mode 100644 index 000000000..e3aeb5b12 --- /dev/null +++ b/docs/design/feature/multi_request_streaming.md @@ -0,0 +1,47 @@ +## Multi-Request Streaming (MRS) on a Single Machine + +### 1. Background & Scope +- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved. +- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling. +- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead. + +### 2. Key Constraints +- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`). +- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version. +- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0). + +### 3. Architecture Overview +- Processes & IPC queues + - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue. + - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects. + - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side. +- Device visibility + - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets. + - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit. +- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work. +- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately. +- Scheduling + - A downstream stage triggers only after the upstream completes the request. + - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided. + +### 4. IPC Implementation (simplified: copy-based) +- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data). +- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version. +- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability. + +### 5. Scheduling & Cancellation (simplified) +- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency. +- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue. + +#### Short sequence example (req0/req1, stage A→B) +1) t0: stage A handles req0 +2) t1: req0 completes on A → enters B; A immediately starts req1 +3) t2: B handles req0 while A handles req1 (parallel across stages) + +### 6. Integration Points (by file) +- `vllm_omni/entrypoints/omni.py` (Orchestrator) + - Class `Omni` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers. + - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow. + - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled. + - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM. + - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions. diff --git a/docs/design/feature/multi_request_streaming_design.md b/docs/design/feature/multi_request_streaming_design.md deleted file mode 100644 index 50e3a8c15..000000000 --- a/docs/design/feature/multi_request_streaming_design.md +++ /dev/null @@ -1,188 +0,0 @@ -## Multi-Request Streaming (MRS) on a Single Machine - -### 1. Background & Scope -- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved. -- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling. -- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead. - -### 2. Key Constraints -- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`). -- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version. -- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0). - -### 3. Architecture Overview -- Processes & IPC queues - - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue. - - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects. - - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side. -- Device visibility - - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets. - - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit. -- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work. -- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately. -- Scheduling - - A downstream stage triggers only after the upstream completes the request. - - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided. - -### 7. IPC Implementation (simplified: copy-based) -- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data). -- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version. -- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability. - -### 8. Scheduling & Cancellation (simplified) -- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency. -- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue. - -#### Short sequence example (req0/req1, stage A→B) -1) t0: stage A handles req0 -2) t1: req0 completes on A → enters B; A immediately starts req1 -3) t2: B handles req0 while A handles req1 (parallel across stages) - -### 9. Integration Points (by file) -- `vllm_omni/entrypoints/omni_llm.py` (Orchestrator) - - Class `OmniLLM` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers. - - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow. - - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled. - - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM. - - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions. - -### 9.1 Process Device Visibility -- Device binding: set `CUDA_VISIBLE_DEVICES` before process start, or call `torch.cuda.set_device` early in init; configs can be single/multi-GPU. -- Cross-device transfer (simplified): via CPU copies; zero-copy paths are out of scope for now. - -### 10. Worker Template (pseudo-code) -```python -def stage_worker(input_q, output_q, runtime, shm_threshold): - # Device binding - set_stage_gpu_devices(runtime.devices) - engine = OmniStageLLM(model, **engine_args) - batch_seq = 0 - while True: - first = input_q.get() - if first is None: - break - # Batch intake (up to runtime.max_batch_size) - batch = [first] - while len(batch) < int(runtime.max_batch_size or 1) and not input_q.empty(): - nxt = input_q.get() - if nxt is None: - input_q.put(None) - break - batch.append(nxt) - # Decode IPC payload and sampling params - request_ids = [] - engine_inputs = [] - rx_bytes = {} - rx_decode_ms = {} - for t in batch: - ein, rx = maybe_load_from_ipc_with_metrics(t, "engine_inputs", "engine_inputs_shm") - request_ids.append(t["request_id"]) - engine_inputs.extend(ein if isinstance(ein, list) else [ein]) - rx_bytes[t["request_id"]] = rx["rx_transfer_bytes"] - rx_decode_ms[t["request_id"]] = rx["rx_decode_time_ms"] - # Generate and dispatch (grouped by request_id) - batch_seq += 1 - outputs = list(engine.generate(engine_inputs, batch[0]["sampling_params"], use_tqdm=False)) - grouped = group_by_request_id(outputs, request_ids) - for rid in request_ids: - r_out = grouped.get(rid, []) - use_shm, payload = maybe_dump_to_shm(r_out, shm_threshold) - msg = { - "request_id": rid, - "metrics": { - "batch_id": batch_seq, - "rx_transfer_bytes": int(rx_bytes.get(rid, 0)), - "rx_decode_time_ms": float(rx_decode_ms.get(rid, 0.0)), - # omitted: generation latency and token stats - }, - } - if use_shm: - msg["engine_outputs_shm"] = payload - else: - msg["engine_outputs"] = payload - output_q.put(msg) -``` - -### 12. Observability -- Metrics (as implemented): - - per-request (emitted by stages): `num_tokens_out`, `stage_gen_time_ms`, `rx_transfer_bytes`, `rx_decode_time_ms`, `rx_in_flight_time_ms`, `batch_id` - - orchestrator aggregates: E2E latency, tokens/s (written only when stats are enabled) - - optional per-stage JSONL: `{log_file}.stage{stage_id}.stats.jsonl` -- Logs/tracing: - - optional per-stage log files: `{log_file}.stage{stage_id}.log` - - the orchestrator can log readiness, forward size/time, and summary information - -### 15. Risks & Mitigations -— End-to-end latency: full-trigger sacrifices interactivity; start simple/correct, add windowed streaming later. -— CPU copy overhead: copy-based IPC for maintainability; add SHM/zero-copy later as an optimization path. - -### 16. Refactor: Sink LLM init and process into Stage (diagram) - -Objective: without changing external call patterns, encapsulate per-stage LLM init, process creation, and worker logic inside `Stage`, so `PipelinedOmniLLM` focuses on orchestration (seeding, polling, forwarding, collecting) while keeping multi-process, shared memory, and device mapping capabilities. - -Key changes (vs. current): -- Before: Orchestrator created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there. -- After: enhanced `OmniStage` owns its subprocess (`init_stage_worker`), with `_stage_worker` as the entry; device setup, LLM init, batching, and SHM live in `OmniStage`. - -Classes & responsibilities -- OmniLLM (orchestrator) - - Build `OmniStage` list in parallel (preserve `process_engine_inputs` wiring) - - Connect adjacent `Stage` input/output queues; seed requests to stage 0 - - Poll stage outputs, decode results, call `process_engine_inputs`, then encode and forward to next stage - - Termination & cleanup: distribute shutdown signals; join/terminate subprocesses -- Stage (stage unit) - - Members: `stage_config` (with `engine_args` and `runtime`), `in_q/out_q`, subprocess handle, stats - - API: - - `init_stage_worker(...)`: spawn subprocess - - `stop_stage_worker()`: graceful exit - - `submit(ipc_payload)`: submit to input queue (may use shared memory) - - `try_collect() -> Optional[payload]`: non-blocking get from output queue - - `process_engine_inputs(stages, prompts)`: reuse/delegate existing logic - - Subprocess entry `worker_main()`: - - `set_stage_gpu_devices` (`CUDA_VISIBLE_DEVICES` & logical index mapping) - - build `OmniStageLLM(model, **engine_args)` - - while-loop: - - take first task + up to `runtime.max_batch_size-1` non-blocking; support `max_batch_size>=1` - - `maybe_load_from_ipc_with_metrics` per `engine_inputs` in batch - - `stage_engine.generate(batched_inputs, sampling_params)` (window `-1`) - - write per-request aggregated outputs to output queue after `maybe_dump_to_shm` - - log start/finish; report `error` per request on exceptions - -Data flow (sketch) -``` -OmniLLM - ├─ Stage[0].start_process(in0, out0) - ├─ Stage[1].start_process(in1, out1) - ├─ ... - └─ Orchestrator loop: - - stage0.submit(encode(engine_inputs)) - - for each stage: - res = stage.try_collect() - outputs = decode(res) - stage.set_engine_outputs(outputs) - if !final: next_inputs = next_stage.process_engine_inputs(stages, original_prompt) - next_stage.submit(encode(next_inputs)) -``` - -Config (preserve & extend) -- `runtime.devices`: - - Comma-separated visibility list (e.g., "2,5,7"): set visible set; default logical index 0 (first GPU) - - Number/string: treat as logical index (0/1-based per implementation) mapping to physical GPUs -- `runtime.max_batch_size`: max requests taken per intake (default 1) -- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; injected into stages) -- Stage link and inputs: - - `engine_input_source`: list of upstream stage_id(s) from which the downstream stage takes inputs (typically a single source) - - `custom_process_input_func`: function to construct downstream engine_inputs from upstream outputs -- Output flags: - - `final_output`, `final_output_type`: marks that a stage produces the final external output and its type - -Kept & strengthened -- Multi-process: each stage in its own process; parallel across stages; serial/batching inside a stage controlled by `max_batch_size` -- Shared memory: large objects via SHM with strict `memoryview -> close -> unlink` -- Device mapping: unified `set_stage_gpu_devices`, support comma lists and logical indices - -Migration (execution order) -1) Enhanced `OmniStage` with `init_stage_worker/stop_stage_worker/submit/try_collect/_stage_worker`. -2) `OmniLLM` holds `OmniStage` instances, replacing the older pipelined class; uses Stage APIs. -3) Reuse `entrypoints/stage_utils`: device/SHM/encode-decode utilities live here. -4) Regression: single/multi-request, various `max_batch_size`, SHM threshold, device mapping, multi-stage parallelism, error paths, and shutdown flows. diff --git a/docs/design/feature/ray_based_execution.md b/docs/design/feature/ray_based_execution.md index 38f0dc95c..fa793aca5 100644 --- a/docs/design/feature/ray_based_execution.md +++ b/docs/design/feature/ray_based_execution.md @@ -1,4 +1,4 @@ -# VLLM-Omni Distributed Utilities +# Distributed utils This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends. diff --git a/docs/design/index.md b/docs/design/index.md index 772e491a3..c5bf7af47 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -8,12 +8,12 @@ This section contains design documents and architecture specifications for vLLM- ## Feature Design Documents -- [Multi-Request Streaming](feature/multi_request_streaming_design.md) -- [Omni Connector Design](feature/omni_connector_design.md) +- [Disaggregated Inference](feature/disaggregated_inference.md) +- [Multi-Request Streaming](feature/multi_request_streaming.md) - [Ray-based Execution](feature/ray_based_execution.md) ## Module Design Documents -- [AR Module Design](module/ar_module_design.md) -- [DIT Module Design](module/dit_module_design.md) -- [Entrypoint Module Design](module/entrypoint_module_design.md) +- [AR Module](module/ar_module.md) +- [DIT Module](module/dit_module.md) +- [Entrypoint Module](module/entrypoint_module.md) diff --git a/docs/design/module/ar_module_design.md b/docs/design/module/ar_module.md similarity index 100% rename from docs/design/module/ar_module_design.md rename to docs/design/module/ar_module.md diff --git a/docs/design/module/dit_module_design.md b/docs/design/module/dit_module.md similarity index 100% rename from docs/design/module/dit_module_design.md rename to docs/design/module/dit_module.md diff --git a/docs/design/module/entrypoint_module_design.md b/docs/design/module/entrypoint_module.md similarity index 100% rename from docs/design/module/entrypoint_module_design.md rename to docs/design/module/entrypoint_module.md diff --git a/vllm_omni/core/dit_cache_manager.py b/vllm_omni/core/dit_cache_manager.py deleted file mode 100644 index e07c4a167..000000000 --- a/vllm_omni/core/dit_cache_manager.py +++ /dev/null @@ -1,13 +0,0 @@ -""" -DiT Cache Manager for vLLM-Omni. - -This module provides caching functionality for DiT (Diffusion Transformer) models -to optimize inference performance and memory usage. -""" - - -class DiTCacheManager: - """Manages DiT-specific caching for optimized inference.""" - - def __init__(self, config): - pass