Skip to content

[core]refactor communication layer: PR1(Added Refactor Infra Only)#1555

Merged
hsliuustc0106 merged 7 commits into
vllm-project:mainfrom
natureofnature:refactor/communication_layer
Apr 13, 2026
Merged

[core]refactor communication layer: PR1(Added Refactor Infra Only)#1555
hsliuustc0106 merged 7 commits into
vllm-project:mainfrom
natureofnature:refactor/communication_layer

Conversation

@natureofnature
Copy link
Copy Markdown
Contributor

@natureofnature natureofnature commented Feb 27, 2026

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Refactor the vLLM-Omni communication layer by moving all data-plane connector.put()/connector.get() calls from Orchestrator, OmniStage, and Scheduler into a unified OmniConnectorModelRunnerMixin at the Model Runner level, while keeping scheduling coordination logic (e.g., WAITING_FOR_CHUNK state) in the Scheduler.

This is the first PR for communication refactoring, which does not change existing workflow. The core files include Omni connector mixin and scheduler coordinator.

Architecture

architecture

Status Summary

  • Main control-plane vs data-plane splitted.
  • OmniConnectorModelRunnerMixin as the main communication hub for the refactored path.
  • Qwen2.5, Qwen 3 omni (aysnc & non async), Bagel works
  • Batch and async chunk flows are mostly unified
  • chunk life cycle management in model runner to avoid chunk data transmission between model runner and scheduler
  • Legacy compatibility paths to be removed.
  • Fully absorb kv cache transfer into the mixin.
  • Heterogeneous rank support (kv cache transfer/full payload).
  • Heterogeneous rank support (async chunk).
  • Code Simplification and Reduction
  • Rebase with main

Test Plan

pytest -s -v tests/e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model

Test Result

Screenshot from 2026-03-29 10-13-23
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. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
  • The test results. Please paste the results comparison before and after, or the e2e results.
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
  • (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)

Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5cd54ca33f

ℹ️ 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".

Comment thread vllm_omni/entrypoints/omni.py Outdated
Comment thread vllm_omni/entrypoints/async_omni.py Outdated
Comment thread vllm_omni/worker/omni_connector_model_runner_mixin.py Outdated
Copy link
Copy Markdown
Collaborator

@lishunyang12 lishunyang12 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left a few comments. This is a big refactor -- worth getting the threading right before it goes in.

Comment thread vllm_omni/worker/omni_connector_model_runner_mixin.py
Comment thread vllm_omni/worker/omni_connector_model_runner_mixin.py
Comment thread vllm_omni/worker/omni_connector_model_runner_mixin.py
Comment thread tests/e2e/online_serving/test_refactored_communication.py Outdated
Comment thread tests/core/sched/test_chunk_scheduling_coordinator.py
if self.chunk_transfer_adapter:
self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
if self.chunk_coordinator:
oco = self._latest_omni_connector_output
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the variable names are too vague

# Core scheduling methods
# ------------------------------------------------------------------ #

def process_pending_chunks(
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to put the processing of Async Chunk and non-async chunk code blocks in the same function? Currently, some processing seems replicated in process_pending_chunks & process_pending_batch_inputs. Could we use a single queue, waiting_for_input_waiting, to handle both waiting_for_input and waiting_for_chunk_waiting simultaneously, and only hold one request status (e.g., WAITING_FOR_INPUT) for both asynchronous and non-asynchronous chunk modes?

Comment thread vllm_omni/core/sched/output.py Outdated
finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
# Requests that need to be registered for chunk recv by the Model Runner's
# background thread. Populated by ChunkSchedulingCoordinator.
pending_chunk_registrations: list = field(default_factory=list)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as above. can we keep only one for registration?

Comment thread vllm_omni/entrypoints/async_omni.py Outdated
seed_input["prompt_token_ids"] = [0] * next_prompt_len
seed_input["multi_modal_data"] = seed_input["mm_processor_kwargs"] = None
for ds_stage_id in range(1, len(self.stage_list)):
sp_ds = sampling_params_list[ds_stage_id]
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the variable name is too vague

@lishunyang12
Copy link
Copy Markdown
Collaborator

lishunyang12 commented Mar 4, 2026

Hey @natureofnature — I see two new commits since my review, but my 4 inline comments from 2/28 still seem open (busy-spin in recv_loop, shallow copy, dead _finished_save_reqs, duplicate MockQueue in e2e test). Could you take a look when you get a chance?

@natureofnature
Copy link
Copy Markdown
Contributor Author

natureofnature commented Mar 4, 2026

Hey @natureofnature — I see two new commits since my review, but my 4 inline comments from 2/28 still seem open (busy-spin in recv_loop, shallow copy, dead _finished_save_reqs, duplicate MockQueue in e2e test). Could you take a look when you get a chance?

I'm still working on it, and will fix them once the basic workflow functions well. @lishunyang12

@natureofnature
Copy link
Copy Markdown
Contributor Author

@codex review

@natureofnature natureofnature force-pushed the refactor/communication_layer branch from 04a89ed to 3d3e6e8 Compare March 9, 2026 02:20
@natureofnature
Copy link
Copy Markdown
Contributor Author

@tzhouam @divyanshsinghvi PTAL

@natureofnature natureofnature force-pushed the refactor/communication_layer branch 2 times, most recently from 391544a to bda4aac Compare March 10, 2026 14:23
@natureofnature natureofnature changed the title [WIP][core]refactor commnication layer [WIP][core]refactor communication layer Mar 11, 2026
@Gaohan123 Gaohan123 added this to the v0.18.0 milestone Mar 11, 2026
@natureofnature natureofnature force-pushed the refactor/communication_layer branch 4 times, most recently from 4b8fd82 to e9601a4 Compare March 18, 2026 06:43
Copy link
Copy Markdown
Collaborator

@hsliuustc0106 hsliuustc0106 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate Status: ❌ BLOCKED

Check Status
DCO
Docs
Mergeable ❌ CONFLICTING

Prior Feedback Unaddressed (2/28)

  1. Dead state: _finished_save_reqs initialized but never read
  2. Silent connector failure: Returns None instead of failing fast
  3. Busy-spin: time.sleep(0.001) placement still suboptimal
  4. Duplicate MockQueue: Same class in two test files

Test Evidence Missing

Claims "Qwen2.5, Qwen 3 omni, Bagel works" but provides no test commands or output. Please add reproduction steps.


Minor: MRO Inconsistency

# gpu_ar_model_runner.py:154
class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):  # ✓

# gpu_generation_model_runner.py:45
class GPUGenerationModelRunner(OmniConnectorModelRunnerMixin, OmniGPUModelRunner):  # ✗

Please resolve conflicts and address the 4 prior items.

@natureofnature natureofnature force-pushed the refactor/communication_layer branch 4 times, most recently from 2b2b419 to 776d5ad Compare March 19, 2026 02:59
@natureofnature natureofnature marked this pull request as draft March 19, 2026 06:57
@natureofnature natureofnature force-pushed the refactor/communication_layer branch from 776d5ad to 5b24a97 Compare March 20, 2026 08:11
Copy link
Copy Markdown
Collaborator

@hsliuustc0106 hsliuustc0106 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this PR to too large, can we test this in one model intensively before moving to all models?

Comment thread vllm_omni/entrypoints/talker_prompt_utils.py Outdated
@natureofnature natureofnature force-pushed the refactor/communication_layer branch 3 times, most recently from e6d6219 to eb7de9b Compare March 30, 2026 02:26
@natureofnature natureofnature changed the title [WIP][core]refactor communication layer [core]refactor communication layer: PR1(Added Refactor Infra Only) Mar 30, 2026
@natureofnature
Copy link
Copy Markdown
Contributor Author

this PR to too large, can we test this in one model intensively before moving to all models?
I split the pr to two parts, this is the first part, which does not change existing workflow. @hsliuustc0106

Comment thread vllm_omni/worker/payload_span.py
Comment thread vllm_omni/worker/omni_connector_model_runner_mixin.py Outdated
@@ -0,0 +1,380 @@
# SPDX-License-Identifier: Apache-2.0
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@hsliuustc0106 hsliuustc0106 self-requested a review March 30, 2026 07:45
@hsliuustc0106 hsliuustc0106 added the nightly-test label to trigger buildkite nightly test CI label Mar 30, 2026
@hsliuustc0106 hsliuustc0106 added the ready label to trigger buildkite CI label Apr 10, 2026
@hsliuustc0106
Copy link
Copy Markdown
Collaborator

merge ci passed but nightly ci failed

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
@natureofnature natureofnature force-pushed the refactor/communication_layer branch 2 times, most recently from eddf5d1 to 590a7d0 Compare April 12, 2026 07:52
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
@hsliuustc0106 hsliuustc0106 added the high priority high priority issue, needs to be done asap label Apr 12, 2026
@natureofnature
Copy link
Copy Markdown
Contributor Author

natureofnature commented Apr 12, 2026

merge ci passed but nightly ci failed

I fixed one issue that caused by this PR. And it seems current nightly CI problems (https://buildkite.com/vllm/vllm-omni/builds/6413/steps/canvas) are not related to this PR. @hsliuustc0106

@hsliuustc0106 hsliuustc0106 merged commit 0d4e975 into vllm-project:main Apr 13, 2026
7 of 8 checks passed
Celeste-jq pushed a commit to IsleOfDawnlight/vllm-omni-voxcpm that referenced this pull request Apr 14, 2026
…llm-project#1555)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
lengrongfu pushed a commit to lengrongfu/vllm-omni that referenced this pull request May 1, 2026
…llm-project#1555)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
clodaghwalsh17 pushed a commit to clodaghwalsh17/nm-vllm-omni-ent that referenced this pull request May 12, 2026
…llm-project#1555)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

high priority high priority issue, needs to be done asap nightly-test label to trigger buildkite nightly test CI ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants