[Core]Add Diffusion executor by natureofnature · Pull Request #865 · vllm-project/vllm-omni

natureofnature · 2026-01-20T06:36:48Z

Purpose

This PR aligns the diffusion execution stack with vLLM’s model executor structure, extracting executor responsibilities from the engine and making the executor pluggable. The refactor sets the foundation for worker-actor based execution and enables swapping in RL-specific executors. In the future, ray based worker actor could then be easily added accordingly to support distributed execution.
Key goals:

Align with vLLM model executor design
Enable future worker-actor implementations
Support RL workflows via executor replacement ([Feature]: Support ray backend support for Omni Diffusion Worker #796, [RFC]: Reinforcement learning support on vllm-omni #778)

Test Plan

Test Result

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

natureofnature · 2026-01-20T07:27:23Z

@codex review

chatgpt-codex-connector · 2026-01-20T07:32:30Z

Codex Review: Didn't find any major issues. Another round soon, please!

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

natureofnature · 2026-01-20T07:48:49Z

@codex review

chatgpt-codex-connector · 2026-01-20T07:53:05Z

Codex Review: Didn't find any major issues. 🚀

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

hsliuustc0106 · 2026-01-20T08:11:40Z

@SamitHuang @ZJY0516 @wtomin PTAL

ZJY0516 · 2026-01-20T08:47:31Z

cc @knlnguyen1802

ZJY0516

overll LGTM. Let's wait for #847

ZJY0516 · 2026-01-20T09:12:41Z

vllm_omni/diffusion/executor/abstract.py

+                return MultiprocDiffusionExecutor
+
+            try:
+                executor_class = resolve_obj_by_qualname(backend)


which scenario we need a class name instead of "mp" or "ray"?

For example, as for RL case, @knlnguyen1802 will define his own executor and worker to enable user defined worker functions. (#686) @ZJY0516

I think for clear you can explicit it as an "external_backend". See example from vllm https://github.com/vllm-project/vllm/blob/8be263c3fb1f98d85bd6a06d52e6036057f8814e/vllm/v1/executor/abstract.py#L73

I think for clear you can explicit it as an "external_backend". See example from vllm https://github.com/vllm-project/vllm/blob/8be263c3fb1f98d85bd6a06d52e6036057f8814e/vllm/v1/executor/abstract.py#L73

I updated the code to align with vllm for this function, and I suppose you should use the branch isinstance(distributed_executor_backend, str)

NickLucche · 2026-01-20T10:19:54Z

vllm_omni/diffusion/diffusion_engine.py


    def close(self) -> None:
-        self._finalizer()
+        if hasattr(self, "executor"):


why wouldn't the engine have an executor attribute here?

This is defensive programming to handle cases where DiffusionEngine initialization fails before the executor is fully assigned.

NickLucche · 2026-01-20T10:23:42Z

vllm_omni/diffusion/scheduler.py

-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            cls._instance = super().__new__(cls)
-        return cls._instance


why is the singleton pattern being dropped?

I think the scheduler manages resources that should be specific to a single DiffusionEngine instance, not global to the entire process.

NickLucche · 2026-01-20T10:27:26Z

vllm_omni/diffusion/executor/multiproc_executor.py

+        self.scheduler = Scheduler()
+        self.scheduler.initialize(self.od_config)


doesn't the scheduler belong with the engine, if aligning with vllm design?

Currently, the Scheduler effectively acts as a communicator between the engine and worker via its internal MessageQueue, and the plugged executor needs it to do the communication as well. That's why it's kept into the executor. Since the executor is owned by the Engine, the scheduler remains an instance-level resource, strictly aligning with vLLM's design.
In the future, I suggest we move the communication channel from the scheduler to the executor to better separate their concerns. However, this commit focuses on architecture updates , and I prefer not to modify the internal workflow too drastically at this stage.

keep it in a future pr

natureofnature · 2026-01-21T03:14:33Z

@codex review

chatgpt-codex-connector · 2026-01-21T03:18:38Z

Codex Review: Didn't find any major issues. Already looking forward to the next diff.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Signed-off-by: wzliu <wzliu@connect.hku.hk>

knlnguyen1802 · 2026-01-21T09:31:23Z

vllm_omni/diffusion/executor/multiproc_executor.py

+        self.scheduler = Scheduler()
+        self.scheduler.initialize(self.od_config)


This 2 line can be merge into 1. The initialize can be move inside the constructor

knlnguyen1802

Overall, design and logic LGTM, thanks for the work.

hsliuustc0106

lgtm, please update the desgin doc for diffusion module @SamitHuang

natureofnature requested a review from hsliuustc0106 as a code owner January 20, 2026 06:36

natureofnature force-pushed the diffusion_executor branch from 8b9dc2a to 68b12b3 Compare January 20, 2026 07:18

natureofnature force-pushed the diffusion_executor branch from 68b12b3 to 58fbd80 Compare January 20, 2026 07:46

ZJY0516 reviewed Jan 20, 2026

View reviewed changes

NickLucche reviewed Jan 20, 2026

View reviewed changes

natureofnature requested review from NickLucche and knlnguyen1802 January 21, 2026 03:47

ZJY0516 added the ready label to trigger buildkite CI label Jan 21, 2026

ZJY0516 approved these changes Jan 21, 2026

View reviewed changes

natureofnature added 3 commits January 21, 2026 16:50

split diffusion engine to engine + executor for vllm compatible

012ff22

Signed-off-by: wzliu <wzliu@connect.hku.hk>

abstract class aligned with vllm

3a7c049

Signed-off-by: wzliu <wzliu@connect.hku.hk>

add __init__ to avoid doc error

f8782d1

Signed-off-by: wzliu <wzliu@connect.hku.hk>

natureofnature force-pushed the diffusion_executor branch from f26c4e3 to f8782d1 Compare January 21, 2026 08:51

knlnguyen1802 reviewed Jan 21, 2026

View reviewed changes

knlnguyen1802 approved these changes Jan 21, 2026

View reviewed changes

natureofnature requested a review from ZJY0516 January 21, 2026 10:43

hsliuustc0106 approved these changes Jan 21, 2026

View reviewed changes

hsliuustc0106 merged commit c9d7cd1 into vllm-project:main Jan 21, 2026
7 checks passed

natureofnature mentioned this pull request Jan 23, 2026

[Feature]: Multi-Node Serving: Support for multi-node setups using the Mooncake Connector. JiusiServe/vllm-omni#31

Open

1 task

		self.scheduler = Scheduler()
		self.scheduler.initialize(self.od_config)

Comments

Conversation

natureofnature commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

natureofnature commented Jan 20, 2026

Uh oh!

chatgpt-codex-connector bot commented Jan 20, 2026

Uh oh!

natureofnature commented Jan 20, 2026

Uh oh!

chatgpt-codex-connector bot commented Jan 20, 2026

Uh oh!

hsliuustc0106 commented Jan 20, 2026

Uh oh!

ZJY0516 commented Jan 20, 2026

Uh oh!

ZJY0516 left a comment

Choose a reason for hiding this comment

Uh oh!

ZJY0516 Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

natureofnature Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

natureofnature Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

natureofnature Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

natureofnature commented Jan 21, 2026

Uh oh!

chatgpt-codex-connector bot commented Jan 21, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

knlnguyen1802 left a comment

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

natureofnature commented Jan 20, 2026 •

edited

Loading

ZJY0516 Jan 20, 2026 •

edited

Loading

natureofnature Jan 21, 2026 •

edited

Loading

natureofnature Jan 21, 2026 •

edited

Loading

natureofnature Jan 21, 2026 •

edited

Loading