- 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.
- 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)
- 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
defstage_worker(input_q,output_q,ctx):
whilectx.alive:
item=input_q.get()# blocking / with timeout
ifitem.is_cancelled():
continue
withcuda_stream(ctx.stream):
out_segments=run_stage_compute(item)# produce segments
- 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.
— 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.
- 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`.
