[Refactor] Stage Proc/Client for LLM and DiT

[chickeyton]

PLEASE FILL IN THE PR DESCRIPTION HERE.

Purpose

Refactor the following files

vllm_omni/diffusion/
inline_stage_diffusion_client.py
stage_diffusion_client.py
stage_diffusion_proc.py

vllm_omni/engine/
stage_client.py
stage_engine_core_client.py
stage_engine_core_proc_manager.py
stage_engine_core_proc.py
stage_pool.py

There will be subdirectories to contains the refactored code:

vllm_omni/diffusion/stage/
inline_stage_diffusion_client.py
stage_diffusion_client.py -> stage_diffusion_core_client.py
stage_diffusion_proc.py -> stage_diffusion_core_proc.py
stage_diffusion_core_proc_manager.py

vllm_omni/engine/stage/
stage_core_types.py
stage_client.py -> stage_core_client.py
stage_engine_core_client.py -> stage_llm_core_client.py
stage_engine_core_proc_manager.py -> stage_llm_core_proc_manager.py
stage_engine_core_proc.py -> stage_llm_core_proc.py
stage_pool -> stage_replica_pool.py

• Better naming, accurate and informative
• Remove unused, duplicated or redundant code
• More elegant OOP structures

Inheritance Diagrams

flowchart BT
    ECP["EngineCoreProc"]
    SLCP["StageLLMCoreProc"]
    SDCP["StageDiffusionCoreProc"]

    CEPM["CoreEngineProcManager"]
    SLCPM["StageLLMCoreProcManager"]
    SDCPM["StageDiffusionCoreProcManager"]

    SLCP --> ECP
    SLCPM --> CEPM
    SLCPM -. "contains" .-o SLCP
    SDCPM -. "contains" .-o SDCP

Loading

flowchart BT
    AMPC["AsyncMPClient"]
    DAMPC["DPLBAsyncMPClient"]

    SCCB["StageCoreClientBase"]
    SLCCB["StageLLMCoreClientBase"]
    SLCC["StageLLMCoreClient"]
    DSLCC["DPLBStageLLMCoreClient"]
    SDCC["StageDiffusionCoreClient"]
    ISDC["InlineStageDiffusionClient"]

    SLCC  --> SCCB
    SLCCB --> SCCB
    DSLCC --> SLCCB 
    DSLCC --> DAMPC
    SLCC --> AMPC 
    SLCC --> SLCCB
    SDCC --> SCCB
    ISDC --> SCCB

Loading

flowchart BT
    ECR["EngineCoreRequest"]
    ECO["EngineCoreOutput"]
    ECOS["EngineCoreOutputs"]

    SCR["StageCoreRequest"]
    SCO["StageCoreOutput"]
    SCOS["StageCoreOutputs"]

    SLCR["StageLLMCoreRequest"]
    SLCO["StageLLMCoreOutput"]
    SLCOS["StageLLMCoreOutputs"]

    SDCR["StageDiffusionCoreRequest"]
    SDCO["StageDiffusionCoreOutput"]
    SDCOS["StageDiffusionCoreOutputs"]

    SLCR --> ECR
    SLCR --> SCR

    SLCO --> ECO
    SLCO --> SCO

    SLCOS --> ECOS
    SLCOS --> SCOS

    SDCR --> SCR
    SDCO --> SCO
    SDCOS --> SCOS

Loading

Changes

stage_core_types.py

• Add StageCoreRequest, StageCoreOutput, StageCoreOutputs
• Add StageLLMCoreRequest, StageLLMCoreOutput, StageLLMCoreOutputs
• Add StageDiffusionCoreRequest, StageDiffusionCoreOutput, StageDiffusionCoreOutputs

stage_client.py -> stage_core_client.py

• StagePoolClient, StagePoolLLMClient, StagePoolDiffusionClient not used, remove them
• rename StageClientBase to StageCoreClientBase
• update interface of StageCoreClientBase

stage_engine_core_client.py -> stage_llm_core_client.py

• rename StageEngineCoreClientBase to StageLLMCoreClientBase (subclass from StageCoreClientBase)
• rename StageEngineCoreClient to StageLLMCoreClient
• rename DPLBStageEngineCoreClient to DPLBStageLLMCoreClient
• make StageLLMCoreClientBase subclass from StageCoreClientBase
• rename _default_process_engine_inputs to _default_process_core_inputs as a static function of StageLLMCoreClientBase
• remove the monkey patch applied to vllm as which exists in patch.py already:
• replace the references of EngineCoreRequest and EngineCoreOutputs with StageLLMCoreRequest and StageLLMCoreOutputs
• rename _kv_sender_host to _core_host, _resolve_contact_host to _resolve_core_host
• use concrete typing for arguments of engine_manager (renamed to proc_manager) and coodinator in make_async_mp_client:

proc_manager: StageLLMCoreProcManager | None
coordinator: DPCoordinator | None

engine/init.py

• remove OmniEngineCoreOutput and OmniEngineCoreOutputs

patch.py

• replace the references of OmniEngineCoreOutput, OmniEngineCoreOutputs and OmniEngineCoreRequest by StageLLMCoreOutput,
  StageLLMCoreOutputs and StageLLMCoreRequest

stage_engine_core_proc_manager.py -> stage_llm_core_proc_manager.py

• rename StageEngineCoreProcManager to StageLLMCoreProcManager
• in __init__, argument start_index (the configurated dp rank) and local_start_index are always set to 0 by callers, to be removed
• rename local_engine_count to local_proc_count
• remove local_dp_rank as not needed by StageLLMCoreProc

stage_engine_core_proc.py -> stage_llm_core_proc.py

• rename StageEngineCoreProc to StageLLMCoreProc
• Argument local_dp_rank in run_stage_core is not used, to be removed
• remove the monkey patch for vllm, which is already exists in patch.py

stage_pool.py -> stage_replica_pool.py

• rename StagePool to StageReplicaPool

inline_stage_diffusion_client.py

• make InlineStageDiffusionClient subclass from StageClientBase

stage_diffusion_client.py -> stage_diffusion_core_client.py

• rename StageDiffusionClient to StageDiffusionCoreClient and subclass from StageClientBase

stage_diffusion_proc.py -> stage_diffusion_core_proc.py

• rename StageDiffusionProc to StageDiffusionCoreProc
• move StageDiffusionProcManager to stage_diffusion_core_proc_manager.py

stage_diffusion_core_proc_manager.py

• rename StageDiffusionProcManager to StageDiffusionCoreProcManager

Test Plan

vLLM Version: v0.25.0

Environment

Item	Value
Host	Test server 1
GPU	1× NVIDIA L20X (both stages colocated on one GPU)
vLLM	0.25.0

Model

• Model: ByteDance-Seed/BAGEL-7B-MoT
• Topology: two-stage (default)
  • Stage 0 — Thinker (AR / LLM): runs on the vLLM AR engine via the relocated StageLLMCoreProc.
  • Stage 1 — DiT (Diffusion): runs out-of-process via the relocated StageDiffusionCoreProc / StageDiffusionCoreClient / StageDiffusionCoreProcManager; the KV/latent handoff from Stage 0 crosses the ZMQ boundary through SharedMemoryConnector.
• Task: TI2I / image-edit — an input image plus a text instruction produce an edited image.

Replica & parallel settings

Setting	Stage 0 (Thinker)	Stage 1 (DiT)
Replicas	1	1
devices	"0"	"0" (colocated with Stage 0)
Tensor parallel	1	1
Sequence / CFG / USP parallel	—	none (DiffusionParallelConfig() defaults)
Data parallel	1 (DP > 1 unsupported for the omni stage LLM engine)	—
max_num_seqs	3	1
max_num_batched_tokens	32768	32768
gpu_memory_utilization	0.45	(default)
enforce_eager	—	true
enable_prefix_caching	false	false

No tensor / sequence / data parallelism was used — this is the single-GPU,
single-replica-per-stage default. The diffusion side uses DiffusionParallelConfig()
with its default (TP=1, no SP/CFG/USP parallel).

YAML config (vllm_omni/deploy/bagel.yaml)

async_chunk: false

stages:
  - stage_id: 0
    max_num_batched_tokens: 32768
    max_num_seqs: 3
    gpu_memory_utilization: 0.45
    trust_remote_code: true
    enable_prefix_caching: false
    devices: "0"
    default_sampling_params:
      temperature: 0.4
      top_p: 0.9
      top_k: 1
      max_tokens: 2048
      seed: 52
      detokenize: true
      repetition_penalty: 1.05

  - stage_id: 1
    max_num_batched_tokens: 32768
    max_num_seqs: 1
    enforce_eager: true
    trust_remote_code: true
    enable_prefix_caching: false
    devices: "0"
    input_connectors:
      from_stage_0: shared_memory_connector
    default_sampling_params:
      seed: 52

connectors:
  shared_memory_connector:
    name: SharedMemoryConnector

  rdma_connector:
    name: MooncakeTransferEngineConnector
    extra:
      host: "auto"
      zmq_port: 50051
      protocol: "rdma"
      device_name: ""
      memory_pool_size: 4294967296
      memory_pool_device: "cpu"

Launch / run command

# Server
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
    --deploy-config vllm_omni/deploy/bagel.yaml

# Request (img2img / TI2I)
IMAGE_BASE64=$(base64 -w 0 bagel_02_input.png)
curl http://localhost:8091/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{
      "role": "user",
      "content": [
        {"type": "text", "text": "<|im_start|>Turn the main subject into a polished gold sculpture.<|im_end|>"},
        {"type": "image_url", "image_url": {"url": "data:image/png;base64,'"${IMAGE_BASE64}"'"}}
      ]
    }],
    "modalities": ["image"],
    "height": 512, "width": 512,
    "num_inference_steps": 30, "seed": 42
  }'

Example (1 of 10 jobs)

Input text (edit instruction):

Subject re-rendered as a polished gold sculpture; city-street setting preserved.

Input image — bagel_02_input.png (a green fire hydrant with a yellow cap and white leaf motif on a city sidewalk):

Output image — bagel_02.png, 512×512:

Test Result

Loaded in 175.4s  class=BagelPipeline  stages=2
[JOB  1] OK    21.9s  size=(512, 512)  -> bagel_01.png   (first job incl. warmup)
[JOB  2..10] OK ~8.6-8.8s each         -> bagel_02..10.png
BAGEL TI2I DONE: 10/10 succeeded

• TI2I e2e: 10/10 images generated, all 512×512, subject preserved with the
  requested background/style edit applied.
• Confirms the relocated out-of-process diffusion stage
  (StageDiffusionCoreProc/Client/ProcManager) and Thinker stage
  (StageLLMCoreProc) work end-to-end through the orchestrator after the move.
• Also validated on this environment: 175/176 unit tests pass across the
  relocated stage modules, and single-stage Qwen-Image-Edit = 3/3.

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[chickeyton]

Pre-check report

• Mode: full
• Type: general (refactor) — relocates stage proc/client into vllm_omni/engine/stage/ + vllm_omni/diffusion/stage/; the diffusion/ touches are module relocations, not a new diffusion model.
• Base: origin/main @ 2bb61890
• HEAD: 43acd350 · Diff: 49 files · 19 commits (all DCO-signed)
• Re-run: 2026-07-27 (after fixing both code findings)

Dimension summary

Dimension	Result	Change since first run
PR title format	✗	unchanged (process, not code)
Code quality — kwargs	✓	—
Code quality — broad except	✓	— (relocated; metrics fix removed one)
Code quality — Any hints	✓	⚠ → ✓ (2 msgspec fields now documented-intentional)
Code quality — hot-path copy	✓	—
Code quality — event-loop blocking	✓	—
Metrics-helper divergence from main	✓	⚠ → ✓ (adopted centralized helpers)
Commit hygiene (squash)	⚠	unchanged (now 19 commits)
Dead code	✓	—
Test coverage	✓	— (176 unit still green after fix)
Rebase / CI gates	✓	19 ahead / 0 behind; full pre-commit clean
Scope coherence	✓	—

Verdict: 1 blocking (PR title) · 1 warning (squash). Both code issues from the first run are fixed and validated. The only remaining items are process (title + squash), not code.

---

Fixed since first run

✅ Metrics-helper divergence → resolved (43acd350)

The relocated pool no longer carries the inline helpers that #5168 centralized.
stage_replica_pool.py now imports and uses count_audio_frames, count_image_pixels,
coerce_positive_int_scalar, iter_mm_outputs (from vllm_omni.metrics(.utils)) and
defs.resolve_audio_sample_rate[_or_none], and drops the duplicated inline
_count_audio_frames / _coerce_int_scalar / _count_image_value_pixels /
_iter_multimodal_outputs. _collect_audio_metrics / _infer_audio_sample_rate gained the
same use_default_sample_rate / use_default plumbing main adopted.

• Net −58 lines (50 insertions / 108 deletions).
• Verified: grep confirms 0 inline helper definitions remain; metrics-focused tests
  (-k "audio or sample_rate or image or metric") = 3 passed; full stage suite =
  176 passed, 0 failed on server 1.

✅ Loose Any msgspec fields → resolved (43acd350)

Investigation showed StageDiffusionCoreRequest.prompt and StageDiffusionCoreOutput.output
are msgspec wire fields — the client does encoder.encode(request) and the proc does
msgspec.convert(msg, StageDiffusionCoreRequest). They carry model-dependent / opaque payloads
(prompts with images, the full OmniRequestOutput) handled by the Omni msgpack enc/dec hooks.
Narrowing them would make msgspec.convert reject valid model-specific shapes on the wire,
so Any is the correct type — not laziness. Both fields now carry a comment documenting this,
turning an unexplained Any into a deliberate, reviewable one. (The other two Any in new
code — __init__(*args, **kwargs) vLLM-compat forwarding and collective_rpc_async -> Any —
were always justified.)

---

Remaining (process, not code)

✗ PR title format (blocking at PR-creation time)

No commit carries a conventional prefix. Open the PR with, e.g.:

[Core] Refactor stage proc/client into engine/stage & diffusion/stage subpackages

⚠ Commit hygiene — squash before merge

19 commits, several low-signal (enhancement, bugfix, fix arg name). Collapse into a small,
well-described set on squash-merge, preserving Signed-off-by.

---

Passing (✓) — evidence

• Code quality: kwargs = only vLLM-compat super().__init__(*args, **kwargs) + test asserts;
  broad-except = relocated backlog (the metrics fix removed one inline except Exception);
  Any = all justified/documented; zero new .clone()/deepcopy/time.sleep/blocking-HTTP/lock-across-await;
  zero new SimpleNamespace test fakes.
• Dead code: refactor deletes 7 legacy modules + the Protocol layer; ruff F401 clean; every
  new stage/ and diffusion/stage/ module is imported in production and exercised by tests.
• Test coverage: 17 test files; validated on server 1 (L20X) — 176 unit tests pass,
  plus BAGEL-7B-MoT two-stage TI2I e2e = 10/10 and single-stage Qwen-Image-Edit = 3/3.
• Rebase / CI gates: 19 ahead / 0 behind origin/main (2bb61890); full pre-commit suite
  passes (ruff, typos, whitespace, pytest-marks, pickle-imports); all 19 commits DCO-signed.
• Scope: the non-stage touches are all OmniEngineCore*→StageLLMCore* rename propagation.

---

Recommended actions before opening the PR

1. Set a [Core] PR title (blocking).
2. (Optional) Squash the 19 commits into a clean set, preserving sign-offs.
3. Write a PR body: motivation (dedupe stage runtime), what moved where, the
   OmniEngineCore*→StageLLMCore* rename, adoption of vllm_omni.metrics helpers, and the
   validation evidence (176 unit + BAGEL 10/10).

Report is advisory and local-only — nothing was posted to GitHub.

[hsliuustc0106]

will the change the usage of diffusion DP? for example, we currently work heavily with distributed layerwise offload, it rereuiqs work with DP replica

[chickeyton]

will the change the usage of diffusion DP? for example, we currently work heavily with distributed layerwise offload, it rereuiqs work with DP replica

There is no functional nor behavioural change, just code relocation and renaming

[hsliuustc0106]

please review the agent codes by yourself with comments first

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit fd880bfbd074 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[lishunyang12]

Use msgspec.convert(self._decoder.decode(frame), StageEncodeCoreOutputs); decode doesn't accept a type argument.

[lishunyang12]

Update test_stage_runtime_passes_log_stats_to_llm_replica_launch to use StageLLMCoreClientBase; the stale name raises AttributeError.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 7 days

@chickeyton this pull request has had no human commit, comment or review since 2026-09-16. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.