vllm-omni/docs/analysis/edge_branch_plan.md

Plan: an edge branch for vLLM-Omni

Target: a development branch in ../vllm-omni (main @ 7be014bc) that makes vLLM-Omni itself run, and perform acceptably, on edge devices. This is the "S1/S2/S4 inside vLLM-Omni" part of vllm_omni_edge_feasibility.md; the C++ orchestrator work on llama.cpp/MNN/ExecuTorch (S3) stays outside this branch and is covered by implementation_roadmap.md. Every touch point below was located in the checkout; nothing has been changed yet.

0. Ground rules

• Branch: edge/omni-edge off main@7be014bc, repo-local identity tzhouam <tzhouam@connect.ust.hk>, credential helper via gh (per /data/zhoutaichang/CLAUDE.md). Work packages land as separate sub-branches (edge/wp1-deploy-profiles, ...) so each can become an upstream PR on its own.
• Environment: a fresh uv venv for this checkout (uv pip install -e . per the repo's AGENTS.md/docs), because the machine's Python imports /home/zhoutaichang/feature/vllm-omni-main (be335a86f), not this tree. This is a dependency installation and needs your go-ahead.
• Upstream policy: vLLM-Omni's .claude/skills (precheck-pr, review-pr, vllm-omni-test) define PR shape and test levels (core_model L1/L2, advanced_model L3, full_model L4; cpu platform marker exists in pyproject.toml). Each WP below names the tests it adds at those levels.
• GPU runs go through gpu run --timeout --note; model downloads and installs are permission-gated.

1. Work packages

graph LR
  WP0[WP0 branch + env + baseline harness] --> WP1[WP1 edge deploy profiles]
  WP0 --> WP2[WP2 cold-start reduction]
  WP0 --> WP3[WP3 CPU platform]
  WP1 --> WP4[WP4 Jetson / aarch64 CUDA]
  WP2 --> WP4
  WP1 --> WP5[WP5 codec-token streaming mode]
  WP3 --> WP6[WP6 dependency slimming]
  WP1 --> WP7[WP7 device-direct connector]
  WP5 --> WP8[WP8 benchmarks + CI lanes]
  WP3 --> WP8
  WP4 --> WP8

WP0 — Branch, environment, baseline harness (2–3 days)

• Create the branch and the uv environment; confirm python -c "import vllm_omni" resolves to this checkout.
• Promote the session driver experiments/vllm_omni_tts/tts_stream_bench.py into benchmarks/tts/stream_latency_bench.py (reuses the example's _build_inputs; records TTFA, per-chunk arrivals, uninterrupted-playback latency, RTF streamed+tail, init time, RSS, nvidia-smi peak). JSON schema = the one in experiments/vllm_omni_tts/bench_results.json.
• Re-run the Qwen3-TTS 1.7B baseline on this checkout (not the installed one) and add the 0.6B variant (driver flag to override the hard-coded model ids in _build_inputs).
• Acceptance: baseline JSON for 0.6B and 1.7B on L20X committed under benchmarks/tts/results/; numbers within noise of experiments.md §2 for 1.7B.

WP1 — Edge deploy profiles (3–5 days)

Touch points: vllm_omni/deploy/ (90 YAMLs; platform overlays under platforms: exist in 23 of them), vllm_omni/config/stage_config.py (enforce_eager L408, compilation_config L418, gpu_memory_utilization L402, devices L386, async_chunk L548, inline_diffusion L273), vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py (parse_chunk_ramp L42, ramp_chunk_size L79).

• Add vllm_omni/deploy/edge/qwen3_tts_edge.yaml (and one per TTS family that passes Phase 0 parity): both stages on one device, gpu_memory_utilization sized from measured need (not 0.3), codec_chunk_ramp: [4, 4, 8, 16, 25] on by default (removes the ≈80 ms first-gap measured in experiments.md), decode_batch_max_size: 1, max_num_seqs small, enforce_eager documented as a switch per stage, sampling defaults unchanged.
• Add a profile: edge selector so --deploy-config is not needed (a DeployConfig lookup by (model, profile) in config/config_factory.py).
• Tests: L1 tests/config/test_edge_profiles.py (YAML resolves, schedule parses, ramp ends at codec_chunk_frames); L3 e2e TTFA/playback-gap regression on CUDA using WP0's harness.
• Acceptance: uninterrupted-playback latency ≤ TTFA + 60 ms on L20X with the ramp (versus ≈190 ms today); no RTF regression > 10 %.

WP2 — Cold-start reduction (1–2 weeks)

Touch points: vllm_omni/engine/stage_runtime.py (parallel replica init exists near L460; sequential today for co-located stages), vllm_omni/engine/async_omni_engine.py (_wait_for_orchestrator_init, 300 s default that failed run 1), vllm_omni/worker/gpu_ar_model_runner.py (talker-MTP graph capture, L546 in the installed log), model_executor/models/qwen3_tts/{segmented_graph_wrapper.py,cuda_graph_decoder_wrapper.py}, common/qwen3_code_predictor.py (warm-up buckets + torch.compile).

• Instrument init with a timeline (process spawn, weight load, profile run, KV alloc, graph capture, compile, predictor warm-up) emitted as JSON; make it a benchmarks/tts/init_profile.py.
• Measure the four configurations: default; enforce_eager both stages; compile cache warm; parallel stage init. Keep whichever combination meets the gate and make it the edge profile default.
• Persist and reuse compile/graph artifacts (VLLM_CACHE_ROOT already caches AOT compile; add the code-predictor warm-up buckets and code2wav graph shapes to a versioned cache keyed by checkpoint hash + deploy profile).
• Raise the default orchestrator timeout when the edge profile is active, or scale it from the measured timeline.
• Acceptance: ready-to-serve ≤ 30 s on L20X for Qwen3-TTS 0.6B with warm caches, ≤ 90 s cold; init timeline test at L3.

WP3 — CPU platform (3–5 weeks)

Touch points: vllm_omni/platforms/__init__.py (builtin_omni_platform_plugins L123: cuda/rocm/npu/xpu/musa), vllm_omni/platforms/interface.py (OmniPlatform surface: get_omni_ar_worker_cls L61, get_omni_generation_worker_cls L65, get_diffusion_worker_cls L175, get_device_count L266, memory/synchronize/autocast/set_device_control_env_var L329; UnspecifiedOmniPlatform L385), vllm_omni/platforms/cuda/platform.py (432 lines, the smallest template), vllm_omni/worker/{base.py,gpu_ar_worker.py,gpu_ar_model_runner.py,gpu_generation_worker.py,gpu_generation_model_runner.py,mixins.py}, requirements/cpu.txt (today: common.txt + onnxruntime only), setup.py (already accepts device cpu).

Design: vLLM's CPUWorker(Worker) and CPUModelRunner(GPUModelRunner) (vllm/v1/worker/cpu_worker.py, cpu_model_runner.py) subclass the same GPU classes that vLLM-Omni's OmniGPUWorkerBase/GPUARModelRunner extend, so the CPU platform is built by composition rather than a rewrite:

• vllm_omni/platforms/cpu/platform.py: CPUOmniPlatform(OmniPlatform) returning vllm_omni.worker.cpu_ar_worker.CPUARWorker, cpu_generation_worker.CPUGenerationWorker, a CPU diffusion worker (initially "unsupported" with a clear error), get_device_count() = 1, device_control_env_var = none, memory functions from psutil, supports_talker_mtp_graph_capture() = False, supports_torch_inductor() = vLLM CPU default, autocast = bf16/fp16 per vLLM CPU dtype rules.
• vllm_omni/worker/cpu_ar_worker.py: CPUARWorker(CPUWorker, OmniWorkerMixin) with the omni hooks from gpu_ar_worker.py factored into mixins.py where they are not already; cpu_ar_model_runner.py: CPUARModelRunner(CPUModelRunner) + the omni runner mixin (omni_connector_model_runner_mixin.py); same pair for the generation stage (code2wav on CPU: no CUDA graphs, eager decoder, fp32).
• Register "cpu" in builtin_omni_platform_plugins behind VLLM_TARGET_DEVICE=cpu or vLLM's CPU platform detection; keep UnspecifiedOmniPlatform for the no-device case.
• requirements/cpu.txt: real runtime deps (soundfile, torchaudio CPU, onnxruntime) and the vLLM CPU wheel line; docs page docs/getting_started/installation/cpu.md.
• Model-level: the Qwen3-TTS decoder's Triton Snake/GroupNorm kernels have torch fallbacks (verified in the model report); add a CPU branch in CustomOp.dispatch_forward users where forward_native is missing; CosyVoice3's TRT default (COSYVOICE3_TRT) must be off on CPU.
• Tests: L1 tests/platforms/test_cpu_platform.py (registration, worker classes, device count) and per-family CPU unit tests already exist under tests/model_executor/models/<family>/; L3 e2e Qwen3-TTS 0.6B on CPU (pytest -m "core_model and cpu" lane).
• Validation order: x86 on this host first (the CPU-only path needs no GPU reservation), then aarch64 Linux (vLLM ships aarch64 CPU wheels; KleidiAI W4A8 for the talker).
• Acceptance: Qwen3-TTS 0.6B end-to-end on 8 x86 threads with RTF ≤ 1.0 and TTFA ≤ 500 ms; same test green on one aarch64 box; RSS ≤ 4 GB.

WP4 — Jetson / aarch64 CUDA (2–3 weeks, hardware-gated)

Touch points: none in vLLM-Omni beyond the edge profile and docs; the work is environment: vLLM aarch64+CUDA build or wheel (vLLM detects Tegra via /etc/nv_tegra_release; CUDA arch 8.7/11.0 are in its CMakeLists.txt), torch for JetPack, recipes/edge/jetson.md.

• Run WP0's harness and WP2's init profile on Orin; capture memory with unified-memory caveats (gpu_memory_utilization semantics differ when the GPU shares RAM with the CPU: add a max_memory_gb absolute option to the edge profile in stage_config.py).
• Acceptance: Qwen3-TTS 0.6B RTF ≤ 0.5 and ready-to-serve ≤ 90 s on Orin; documented power mode; blocker if no board is available.

WP5 — Codec-token streaming mode (S4 wire contract) (2–3 weeks)

Touch points: vllm_omni/entrypoints/openai/serving_speech_stream.py (OmniStreamingSpeechHandler, _generate_and_send L260 consumes PCM chunks), vllm_omni/entrypoints/openai/tts_adapters/{base.py,capabilities.py} (register_tts_adapter), vllm_omni/model_executor/models/output_templates.py, vllm_omni/data_entry_keys.py (codes key), vllm_omni/deploy/edge/qwen3_tts_talker_only.yaml (stage 0 only, final_output_type codes).

• Server: a mode: "codec_tokens" session option on /v1/audio/speech/stream that runs the talker-only profile and emits frames {seq, n_frames, codes: int16[n_frames×16], eos} with explicit ordering, EOS, and cancel semantics; abort path unchanged (AsyncOmni.abort).
• Adapter capability flag supports_codec_stream per family (Qwen3-TTS first; MOSS next since its codec is transformer-only).
• Client reference: examples/online_serving/text_to_speech/edge_decoder_client.py decoding on device with the exported Qwen3TTSTokenizerV2Decoder (plain torch CPU first; ExecuTorch .pte second), carrying the 72-frame context and reset() on cancel.
• Tests: L1 protocol tests (ordering, EOS, cancel mid-chunk); L3 e2e PCM parity between server-decoded and client-decoded audio (SNR ≥ 40 dB).
• Acceptance: client-observed TTFA ≤ server TTFA + RTT + first-chunk decode; bandwidth ≈ 400 B/s payload + framing; interruption stops audio within one chunk.

WP6 — Dependency and import slimming (1–2 weeks)

Touch points: pyproject.toml optional-dependencies, requirements/common.txt (diffusers pin, x-transformers, openai-whisper, onnxruntime), vllm_omni/__init__.py and patch.py import graph, vllm_omni/diffusion/ imports pulled in by AR-only paths.

• Add an edge extra that excludes diffusion/video dependencies; make diffusion, forced-aligner, and VAD imports lazy on the TTS path; measure import vllm_omni time and RSS before/after.
• Acceptance: python -c "import vllm_omni" ≤ 3 s and no diffusers import for a TTS-only deploy; test at L1 with an import-guard test.

WP7 — Device-direct connector for co-located stages (2 weeks, optional)

Touch points: vllm_omni/distributed/omni_connectors/connectors/{base.py,shm_connector.py} (OmniConnectorBase.put/get), utils/serialization.py (OmniMsgpackEncoder, zero-copy TODO), deploy YAML connectors: section.

• CudaIpcConnector: put exports a CUDA IPC handle (torch.multiprocessing reductions) for same-GPU stages instead of D2H + msgpack; get maps it; metadata carries handle + shape + dtype; fallback to shared memory for non-tensor payloads. Mori's intra-node xgmi path shows the pattern for AMD; this covers single-GPU NVIDIA co-location (Jetson, single-card servers).
• Acceptance: per-chunk hop latency measured with the WP0 harness drops for the Talker → Code2Wav edge; no change in audio; L1 unit test with a fake two-process exchange.

WP8 — Benchmarks and CI lanes (ongoing)

• benchmarks/tts/stream_latency_bench.py and init_profile.py results committed per profile; a cpu Buildkite lane running pytest -m "core_model and cpu" for WP3; a nightly edge job for the L3 e2e tests on CUDA; recipes under recipes/edge/.

2. Order and effort

Order	WP	Effort	Depends on	Needs hardware / permission
1	WP0	2–3 d	—	uv env install (permission); GPU via scheduler
2	WP1	3–5 d	WP0	GPU
3	WP2	1–2 w	WP0	GPU
4	WP3	3–5 w	WP0	none for x86 CPU; aarch64 box later
5	WP5	2–3 w	WP1	GPU; ExecuTorch install for the .pte client (permission)
6	WP6	1–2 w	WP3	none
7	WP4	2–3 w	WP1, WP2	Jetson board (blocker today)
8	WP7	2 w	WP1	GPU
9	WP8	ongoing	all	CI access

Critical path for a demonstrable result: WP0 → WP1 → WP2 gives a measurably better edge-server tier on the hardware we have; WP3 is the first item that reaches devices without a discrete GPU and is the largest.

3. Risks

• vLLM drift: vLLM-Omni imports 284 vllm.* modules and monkeypatches request/output classes; the CPU worker composition in WP3 depends on CPUWorker/CPUModelRunner internals that changed between the installed 0.28.0 and this checkout's target. Pin the vLLM commit in the branch's requirements/ and re-validate at each vLLM minor.
• Upstream acceptance: each WP is written as an upstream-shaped PR (tests at L1/L3, precheck skill), but WP3 and WP7 are large; expect design review through the repo's imdesign/review-pr flow before implementation.
• Model coverage: the edge profiles are validated per family; families relying on TensorRT (CosyVoice3) or external packages (Fish fish-speech) stay CUDA-only until their CPU fallbacks are verified.
• Hardware: WP4 and the aarch64 half of WP3 cannot start without boards; WP3-x86 and WP5 can.

4. Definition of done for the branch

• Edge profiles, CPU platform, and codec-token streaming merged on edge/omni-edge, each with L1 tests green locally and L3 e2e green on the CUDA lane.
• Published benchmark JSON for: Qwen3-TTS 0.6B/1.7B on L20X (default vs edge profile), 0.6B on x86 CPU and on one aarch64 box, and (if a board arrives) Jetson.
• Docs: docs/getting_started/installation/cpu.md, recipes/edge/*.md, and an updated docs/design/feature/async_chunk.md section on ramp defaults.

5. Implementation status (branch edge/omni-edge, 2026-09-09)

The branch lives in /data/zhoutaichang/embedding_infer/vllm-omni (from main@7be014bc, one commit per work package, not pushed). Everything below was measured on this host (one NVIDIA L20X at a time through gpu run --gpus 1, or CPU-only on NUMA node 1); the JSON is committed under benchmarks/tts/edge_results/ and rendered in experiments/edge_branch/results_summary.md. Section numbers refer to experiments.md §8.

WP	Status	Where	Result / decision
WP0 baseline harness	done	benchmarks/tts/stream_latency_{bench,metrics}.py, tests/benchmarks/	1.7B default: TTFA p50 32 ms, playback-start 116 ms, stall 85 ms, RTF 0.10; 0.6B: 28 / 95 / 67 ms. Reproduces the study's ≈80 ms first-chunk gap.
WP1 edge profile + deploy_profile	done	vllm_omni/deploy/edge/qwen3_tts.yaml, entrypoints/utils.py, --deploy-profile	Ramp [2,4,8,16,25] + small batches: stall 0 ms, playback-start = TTFA (44 ms for 1.7B, 41 ms for 0.6B), RTF unchanged. TTFA +12 ms because the first chunk is 2 frames.
WP2 init timeline + cold start	done	engine/init_timeline.py, benchmarks/tts/init_profile.py	Warm 1.7B: default 152 s, eager stage 1 150 s, eager both 111 s (first TTFA 15 s, rejected), parallel_stage_init 94.6 s (adopted via the profile's orchestrator: block), cold caches 355 s. Dominant phases: engine-core init 51–75 s per stage (model load 9–13 s, talker-MTP graph capture 27 s, predictor compile 12 s, code2wav graphs 4–6 s); the ≤30 s warm target is not met on this stack, the remainder is vLLM's own per-stage init.
WP3 CPU platform	done (no vLLM change needed)	vllm_omni/platforms/cpu/, platforms: cpu: overlay, docs/getting_started/installation/cpu.md	0.6B end to end on the +cpu wheel. Thread sweep: 8 threads TTFA 426 ms / RTF 2.8; 16: 292 ms / 1.8; 16 + inductor: 234 ms / 1.4; 32: 252 ms / 1.7. Attribution: the talker step is 136–250 ms per frame (170–310 % of the 80 ms budget); scheduler + hop overhead < 1 %. Not real-time on this CPU class.
WP4 hardware adaptation	done (emulated)	vllm_omni/edge/{hardware_probe,adapt,calibrate,auto_profile,decoder_export}.py, deploy/edge/hardware/*.yaml, benchmarks/tts/hw_emulation.py	deploy_profile="auto" classifies and sizes every emulated class; Jetson-8 GB emulation (eager, absolute KV budget) TTFA 181 ms / RTF 0.57, Jetson-32 GB 68 ms / 0.09, 8 GB discrete 47 ms / 0.09; CPU classes as in WP3. Decoder export: windowed torch.export verified (exact vs eager, 460 MB fp32 per chunk size); ExecuTorch lowering written but not executed (install permission-gated).
WP5 codec-token streaming	done (parity partly)	deploy/qwen3_tts_talker_only.yaml, entrypoints/openai/codec_stream.py, serving_speech{,_stream}.py, client + parity scripts, docs/serving/codec_token_streaming.md	Streams per step after three runtime fixes found by measurement (FINAL_ONLY coercion, list payloads, delta rows): client TTFA 152–156 ms, stall 0, ramp 2/4/8/16/25 frames, client RTF 0.5 with CPU decode. Client decoder == exact decoder (72–76 dB); end-to-end SNR vs the two-stage server is not meaningful because the two deployments sample different tokens for the same seed (§8.6).
WP6 import slimming	done	lazy diffusion imports, requirements/{core,media,*-edge}.txt, VLLM_OMNI_EDGE_BUILD	AR/TTS path no longer imports diffusers/x_transformers/whisper (guard test); the remaining ≈20 s import overhead is transformers itself.
WP7 CUDA-IPC connector	done, not adopted	connectors/cuda_ipc_connector.py, deploy/qwen3_tts_cuda_ipc.yaml	Talker→code2wav put→get hop: 2.4 ms with shared memory vs 2.9 ms (p95 10.7 ms) with CUDA-IPC (§8.7); below the ≥1 ms gain gate, so the edge profile keeps SharedMemoryConnector.
WP8 edge scheduling	attribution done; gated items not triggered	vllm_omni/edge/step_stats.py, vllm_omni/edge/scheduling.py, --step-stats	Per-step counters show scheduler + chunk hop + dispatch < 2 % of the frame budget on GPU and CPU; the model step is the cost. Event-driven hand-off, inline AR stage and deadline co-scheduling stay unimplemented by the plan's 5 % gate; the no-preemption admission rule is implemented (exact KV bytes/token) and tested.

Decisions recorded.

• Init knobs: keep CUDA graphs on both stages; enable parallel_stage_init in the edge profile (run 20260909-112200). Eager stage 1 buys nothing; eager on both cuts 41 s of init but costs 15 s of first-request latency.
• Scheduling: no scheduler replacement. The dead connector_get_sleep_s knob (no consumer in the tree) and the 1 ms receive-loop back-off were confirmed; polling costs ≤1 ms per chunk, so an event-driven hand-off would recover ≤1.3 % of the frame budget and was not built. The admission rule (vllm_omni/edge/scheduling.py) is used by the hardware adaptation to size max_num_seqs so vLLM never preempts by recompute.
• Native rewrite: not warranted by measurement. On the L20X the whole non-model overhead per frame is a few milliseconds; on CPU the talker forward alone is 2–3× the frame budget. The escape hatch remains the C++ S3 orchestrator on llama.cpp/MNN/ExecuTorch for devices where vLLM cannot run at all (phones), which the codec-token stream and the decoder export now serve.
• CUDA-IPC: measured no gain on the talker → code2wav edge (the chunk tensor is a small host tensor; the IPC export/ack adds 0.5 ms and a worse tail), so the connector only pays off when a hop carries large device tensors (Qwen3-Omni hidden states); it stays available behind qwen3_tts_cuda_ipc.yaml.

Not done / deviations. ExecuTorch was not installed (permission gate; the export path is unit-tested with a fake decoder and the .pt2 export with the real one). The torch-profiler attribution on CPU was replaced by the step counters (same question, lower overhead). No real edge hardware was used; Jetson/ARM numbers are emulations of the derivation logic, not of the silicon. The customvoice regression e2e passed; the codec-stream e2e passed after the client fix, and three further fixes were needed before frames actually streamed per step (§8.6) — the e2e test checks frames and decoding, not cadence, which is a gap to close.