[Core] Cleanup Omni connector runtime and transfer backends by spencerr221 · Pull Request #5096 · vllm-project/vllm-omni

spencerr221 · 2026-07-14T03:51:48Z

Purpose

Simplify SharedMemoryConnector by removing unused device, threshold, inline-payload paths, and related metrics.
Refactor MooncakeTransferEngineConnector initialization and receive flow into focused helpers.
Make Mooncake buffer TTL configurable and prevent cleanup of in-flight buffers.
Remove unused helpers, legacy wrappers, and the redundant vllm_config argument from OmniConnectorModelRunnerMixin.
Update affected callers and connector tests.

Test Plan

python3 -m pytest tests/distributed/omni_connectors/test_basic_connectors.py tests/distributed/omni_connectors/test_adapter_and_flow.py -q

vLLM Version: 0.25.0

vLLM-Omni Commit: 88d4c96

Test Result

17 passed.

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 · 2026-07-14T03:51:52Z

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.

lishunyang12

clean refactor overall — the mooncake except also wrapping the CUDA sync (miscounted as a timeout) is the only one i'd actually look at.

lishunyang12 · 2026-07-14T17:49:53Z

-        vllm_config: Any,
        model_config: Any,
        kv_transfer_manager: OmniKVTransferManager | None = None,
+        **_ignored_kwargs: Any,


Both callers dropped vllm_config, so **_ignored_kwargs now silently swallows typo'd kwargs. Drop it so bad calls fail loudly.

do not use 'any' type

lishunyang12 · 2026-07-14T17:49:53Z

+
+            sync_ms = self._sync_receive_buffer_if_needed()
+        except Exception as e:
+            self._metrics["timeouts"] += 1


This except also wraps _sync_receive_buffer_if_needed(), so a CUDA sync failure gets miscounted as a timeout and logged as RDMA Get error. Bump errors for the non-transfer path instead.

lishunyang12 · 2026-07-14T17:49:53Z

-        return self._get_by_key(get_key)
+        result = self._get_by_key(get_key)
+        if result is not None:
+            self._metrics["gets"] += 1


gets is bumped in 4 separate branches, easy to drift. Increment it once at a single exit instead.

hsliuustc0106

avoid using any type hint for input args

hsliuustc0106 · 2026-07-14T23:02:51Z

-        vllm_config: Any,
        model_config: Any,
        kv_transfer_manager: OmniKVTransferManager | None = None,
+        **_ignored_kwargs: Any,


do not use 'any' type

hsliuustc0106 · 2026-07-14T23:04:21Z


        _t2 = _time_mod.perf_counter()
-        _alloc_ms = (_t2 - _t1) * 1000
+        alloc_ms = (_t2 - _t1) * 1000


shall we add a logger level for time profiling? cc @ZacheryAU

hsliuustc0106 · 2026-07-14T23:05:05Z

-        zmq_addr = f"tcp://{src_host}:{src_port}"
-        req_socket = self._get_req_socket(zmq_addr, timeout_ms=_total_timeout_ms)
-
+        base_timeout_ms = 30000


avoid magic number inside funtions

Copilot

Pull request overview

This PR cleans up the Omni connector runtime APIs and connector backends by removing legacy/unused compatibility paths, refactoring the Mooncake transfer connector into smaller helpers, and simplifying the SHM connector to a single key-addressed SHM mechanism. It also updates the affected GPU model runners and connector tests to match the new behavior.

Changes:

Simplify OmniConnectorModelRunnerMixin initialization/signatures and remove legacy wrapper helpers.
Simplify SharedMemoryConnector by removing inline-payload/threshold logic and related metrics.
Refactor MooncakeTransferEngineConnector.get() into helper methods, add configurable buffer TTL, and prevent TTL cleanup of in-flight buffers.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
vllm_omni/worker/omni_connector_model_runner_mixin.py	Removes deprecated wrappers/helpers and updates connector initialization API.
vllm_omni/worker/gpu_generation_model_runner.py	Updates connector initialization callsite to match new mixin signature.
vllm_omni/worker/gpu_ar_model_runner.py	Updates connector initialization callsite to match new mixin signature (incl. KV manager passing).
vllm_omni/distributed/omni_connectors/connectors/shm_connector.py	Retires inline/threshold paths; always uses key-addressed SHM metadata.
vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py	Refactors init + get/recv flow; adds in-flight tracking and configurable TTL cleanup.
tests/distributed/omni_connectors/test_basic_connectors.py	Updates SHM connector tests to reflect key-addressed SHM-only behavior and ensures cleanup via `close()`.
tests/distributed/omni_connectors/test_adapter_and_flow.py	Updates flow test configuration/comments to reflect SHM metadata behavior change.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

    def init_omni_connectors(
        self,
-        vllm_config: Any,
        model_config: Any,
        kv_transfer_manager: OmniKVTransferManager | None = None,
+        **_ignored_kwargs: Any,
    ) -> None:


hsliuustc0106 · 2026-07-15T04:14:37Z

+                    self._metrics["gets"] += 1
+                return result

            if "shm" in metadata:


Please fix _get_data_with_lock before routing every payload through this path. Its finally block uses if obj, so deserializing a multi-element tensor raises RuntimeError: Boolean value of Tensor with more than one value is ambiguous before get() returns; falsey payloads also leave the lock file behind. I reproduced this with torch.ones(2, 2). Track successful deserialization explicitly and clean up independently of payload truthiness.

hsliuustc0106 · 2026-07-15T04:14:38Z

+                fcntl.flock(lockf, fcntl.LOCK_UN)
+
+            # meta contains {'name': ..., 'size': ...}
+            metadata = {"shm": meta, "size": size}


This now makes every payload use SHM, but shm_threshold_bytes is still exposed in the CLI/config and forwarded by connector initialization. Those existing settings are now silently ignored. Please remove/deprecate the stale plumbing or preserve the documented threshold behavior.

gcanlin · 2026-07-20T03:18:35Z

    # === Cross-stage Communication ===
-    shm_threshold_bytes: int = 65536
+    # Deprecated compatibility flag; SharedMemoryConnector always uses SHM.
+    shm_threshold_bytes: int | None = None


I think it's fine to remove it (and other deprecated warning) directly instead of keeping compatible. Sometimes we forget to remove it in next version, then it will be the new dead code. And this field doesn't seem to be used frequently by users.

Signed-off-by: spencerr221 <liubingyu62@gmail.com>

…ject#5096) Signed-off-by: spencerr221 <liubingyu62@gmail.com>

Since vllm-project#5096, SharedMemoryConnector no longer reads shm_threshold_bytes, all payloads use shared memory. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Nick Cao <ncao@redhat.com>

…ject#5096) Signed-off-by: spencerr221 <liubingyu62@gmail.com>

spencerr221 requested review from gcanlin, princepride, tzhouam, xuechendi, yenuo26 and yuanheng-zhao as code owners July 14, 2026 03:51

lishunyang12 reviewed Jul 14, 2026

View reviewed changes

hsliuustc0106 added the refactor refactoring for better code scalability and quality label Jul 14, 2026

hsliuustc0106 reviewed Jul 14, 2026

View reviewed changes

hsliuustc0106 added the core related to core module: cache, scheduler, engine, worker, modelrunner label Jul 14, 2026

hsliuustc0106 requested a review from Copilot July 14, 2026 23:06

Copilot started reviewing on behalf of hsliuustc0106 July 14, 2026 23:07

View session

Copilot AI reviewed Jul 14, 2026

View reviewed changes

Comment thread vllm_omni/worker/omni_connector_model_runner_mixin.py

Comment on lines 87 to 92

def init_omni_connectors(

self,

vllm_config: Any,

model_config: Any,

kv_transfer_manager: OmniKVTransferManager | None = None,

**_ignored_kwargs: Any,

) -> None:

hsliuustc0106 reviewed Jul 15, 2026

View reviewed changes

hsliuustc0106 added the omni code related to omni models label Jul 17, 2026 — with ChatGPT Codex Connector

Gaohan123 added this to the v0.26.0 milestone Jul 18, 2026

spencerr221 force-pushed the cleanup/connector_mixin branch from 88d4c96 to 1b47e9f Compare July 20, 2026 03:06

spencerr221 requested review from Gaohan123, Sy0307, alex-jw-brooks, david6666666 and ywang96 as code owners July 20, 2026 03:06

gcanlin reviewed Jul 20, 2026

View reviewed changes

npuichigo mentioned this pull request Jul 20, 2026

[RFC]: Placement-aware inter-stage transport — stop bouncing through host when stages share a GPU (status of #2379/#2507/#4147/#4476; prior art in sglang-omni & mstar) #5211

Open

hsliuustc0106 added the ready label to trigger buildkite CI label Jul 20, 2026

spencerr221 force-pushed the cleanup/connector_mixin branch from 1b47e9f to 3434bd3 Compare July 21, 2026 04:01

spencerr221 added 2 commits July 21, 2026 17:13

remove redundant code and refact transfer layer.

8b8ea9a

Signed-off-by: spencerr221 <liubingyu62@gmail.com>

[Core] Address Omni connector review feedback

4831e7f

Signed-off-by: spencerr221 <liubingyu62@gmail.com>

spencerr221 added 2 commits July 21, 2026 17:13

[Core] Remove deprecated SHM threshold option

12a97f8

Signed-off-by: spencerr221 <liubingyu62@gmail.com>

[Core] Fix tts bug.

96c1354

Signed-off-by: spencerr221 <liubingyu62@gmail.com>

spencerr221 force-pushed the cleanup/connector_mixin branch from fe66ff7 to 96c1354 Compare July 21, 2026 09:15

hsliuustc0106 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Jul 22, 2026

hsliuustc0106 merged commit dcf8b07 into vllm-project:main Jul 22, 2026
7 of 9 checks passed

bowieshi pushed a commit to bowieshi/vllm-omni that referenced this pull request Jul 22, 2026

[Core] Cleanup Omni connector runtime and transfer backends (vllm-pro…

b0166a3

…ject#5096) Signed-off-by: spencerr221 <liubingyu62@gmail.com>

pratapyash mentioned this pull request Jul 23, 2026

[Bug]: active_stream_window: 2 stalls Qwen3-Omni audio output after the first codec chunk — even with a single in-flight request — so the pipeline never completes any inference #5349

Closed

nagisa-kunhah pushed a commit to nagisa-kunhah/vllm-omni that referenced this pull request Jul 30, 2026

[Core] Cleanup Omni connector runtime and transfer backends (vllm-pro…

85031ee

…ject#5096) Signed-off-by: spencerr221 <liubingyu62@gmail.com>

NickCao mentioned this pull request Sep 14, 2026

[Misc] treewide: remove shm_threshold_bytes from deploy config #7522

Merged

khairulkabir1661 pushed a commit to khairulkabir1661/vllm-omni that referenced this pull request Sep 25, 2026

[Core] Cleanup Omni connector runtime and transfer backends (vllm-pro…

fb5ddcf

…ject#5096) Signed-off-by: spencerr221 <liubingyu62@gmail.com>

Conversation

spencerr221 commented Jul 14, 2026

Purpose

Test Plan

Test Result

Uh oh!

chatgpt-codex-connector Bot commented Jul 14, 2026

Uh oh!

lishunyang12 left a comment • 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!

lishunyang12 Jul 14, 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!

hsliuustc0106 Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

lishunyang12 Jul 14, 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!

lishunyang12 Jul 14, 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!

hsliuustc0106 left a comment

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

hsliuustc0106 Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

gcanlin Jul 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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

lishunyang12 left a comment •

edited

Loading

lishunyang12 Jul 14, 2026 •

edited

Loading

lishunyang12 Jul 14, 2026 •

edited

Loading

lishunyang12 Jul 14, 2026 •

edited

Loading

gcanlin Jul 20, 2026 •

edited

Loading