anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

adjust the doc

Browse Source 
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>
This commit is contained in:
tzhouam authored and 齐保元 committed 2025-11-14 18:44:29 +08:00
1 parent 30dd7aad34
commit a4e8327797
1 file changed
+70 -119
Download Patch File Download Diff File Expand all files Collapse all files
docs/mrs_design.md
+70 -119
View File
@@ -1,9 +1,5 @@
## Multi-Request Streaming (MRS) on a Single Machine
- Author: Team
- Date: 2025-10-27
- Version: v0.3 (simplified: copy-based IPC / serial per stage / full-trigger)
### 1. Background & Scope
- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved.
- 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.
@@ -12,77 +8,30 @@
### 2. Key Constraints
- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`).
- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version.
- Full-trigger across stages: unify `window_size = -1` for all edges — downstream starts only after upstream completes for that request.
- Serial within a stage: at most one in-flight request per stage (max_inflight=1). No batching/concurrency inside a stage in this version.
- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0).
### 3. Architecture Overview
- Processes & IPC queues
- Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue.
- Inter-stage connection via IPC: control plane (Unix Domain Socket / Pipe) + data plane (shared memory ring buffer in future; copy-based in this version).
- MPMC queue semantics; control plane passes light-weight pointers/metadata; data plane carries large payloads.
- 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 & 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
runtime:
enabled: true
defaults:
window_size: -1 # Full-trigger across stages
max_inflight: 1 # Serial per stage
```
- 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: blocking reads/writes based on queue capacity, blocking upstream on high watermark.
- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability.
### 8. Scheduling & Cancellation (simplified)
- Serial FIFO: each stage `max_inflight=1`; no cross-request micro-batching or concurrency.
- Cancel/timeout: ongoing work should exit early; remaining queued requests stay pending.
- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency.
- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue.
#### Short sequence example (req0/req1, stage A→B)
1) t0: stage A handles req0
@@ -95,79 +44,77 @@ runtime:
- 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.
- `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`
- 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 & Window State (pseudo-code)
### 10. Worker Template (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 []
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)
```
### 11. Configuration
- `runtime.enabled`, `runtime.ar.default_window_size`, `runtime.ar.min_flush_interval_ms`, `runtime.ar.max_segment_tokens`, `runtime.ar.scheduler_policy`.
- Devices/processes:
- Per-stage: `runtime` block on each stage with fields like `process`, `devices`, `max_batch_size`.
- Top-level: `runtime.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.
- 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.
— 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)
@@ -223,12 +170,16 @@ Config (preserve & extend)
- 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
- Windows: keep `window_size=-1` semantics (downstream triggers only after upstream finishes the request)
Migration (execution order)
1) Enhanced `OmniStage` with `init_stage_worker/stop_stage_worker/submit/try_collect/_stage_worker`.