Repository navigation
[Feature] Omni disagg: connector plumbing and downstream handlers - #6893
Draft
RishabhSaini wants to merge 4 commits into
Draft
RishabhSaini wants to merge 4 commits into
RishabhSaini wants to merge 4 commits into
Conversation
Add --standalone flag that boots any single pipeline stage as an independent HTTP server for disaggregated multi-stage serving. - extract_standalone_stage_config() utility extracts one stage from a multi-stage OmegaConf config, adjusting stage_id, final_output, clearing connectors and next-stage transforms - Engine fixes: postprocess() for latent-output stages, generate in supported_tasks for non-tokenizer stages - /v1/stage/run endpoint with entry mode (wraps speech handler, returns raw stage_output JSON) and downstream mode (accepts codec tokens, returns WAV audio) - Standalone gate: endpoint returns 404 on non-standalone servers - Client disconnect abort, codec size validation, sample rate from model Ref: vllm-project#6389 Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>
Server launch script (run_standalone_servers.sh) and reference coordinator client (standalone_disagg_client.py) that chains talker → code2wav over /v1/stage/run. Uses httpx, matching existing example patterns. No external deps beyond vllm-omni. Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>
Tests for extract_standalone_stage_config (stage extraction, connector clearing, transform stripping, input processor preservation, final_output_type inference) and CLI validation (--standalone requires --stage-id, mutually exclusive with --headless). 12 tests, all CPU-only (core_model marker). Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>
RishabhSaini
force-pushed
the
standalone-omni-disagg
branch
from
September 1, 2026 04:13
a9732e9 to
ddc8460
Compare
…_connectors - Add kv_transfer_params flow for standalone omni disagg: connector_info on OmniConnectorOutput → scheduler populates kv_transfer_params on finished requests → chat/speech response carries it → coordinator forwards to downstream stage → per-request sender resolution in recv thread via connector.get(metadata=...) - Extract downstream handlers to serving_stage.py (audio final + intermediate mode) - Add clear_connectors param to extract_standalone_stage_config - Initialize _per_request_sender in mixin __init__, add to prune scan - Tests for clear_connectors, serving_stage utilities, connector plumbing Ref vllm-project#6389 Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>
3 tasks done
hsliuustc0106
added
core
Collaborator
|
This PR touches vllm_omni/entrypoints/, vllm_omni/core/, examples/online_serving/, vllm_omni/worker/, tests/entrypoints/ (14 files). Based on CODEOWNERS coverage of the changed files, the most-related reviewers appear to be: Could one of you take a look when you get a chance? Thanks! |
Author
|
Follow-ons from #6883 review to pick up:
|
Collaborator
|
@RishabhSaini Are you planning to continue work on this PR? Please reply by 2026-10-20 04:14:48 UTC. A yes or a brief plan is enough; code changes do not need to be finished within these three days. If you are waiting on a review or other maintainer action, please let us know. If there is no reply, new commit, or maintainer takeover by then, we will recheck the PR's status before deciding whether to close it as inactive. It can be reopened when work resumes. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Extends standalone stage mode (#6883) to support omni multi-stage disaggregation across separate pods.
Ref #6389
Problem
PR 1 (#6883) enables TTS disagg with inline JSON transfer. Omni disagg (thinker → talker) transfers hidden state embeddings (float tensors, 1-15MB) which is impractical inline. The data needs to go through connectors (Mooncake/Mori RDMA) with the coordinator passing only a small transfer handle — the same pattern vLLM P/D uses for KV cache via
kv_transfer_params.Solution
Reuse
kv_transfer_paramsend-to-end. No new fields or protocols.The flow:
host,zmq_port) at init.kv_transfer_paramson finished requests using this connector info (fallback when the KV connector returns None).kv_transfer_paramsin the response — already wired, no handler changes.SamplingParams.extra_args(same as P/D).register_chunk_recvstashes per-request sender info. The recv thread passes it asmetadatatoconnector.get()Path 2, enabling multi-thinker routing.Other changes
Downstream handlers (
serving_stage.py): Two modes —run_downstream_audio(final stage → WAV) andrun_downstream_intermediate(middle stage → stage_output JSON). Coordinator signals viaoutput_mode.clear_connectorsparam:extract_standalone_stage_config(clear_connectors=False)preserves original stage_id (connector keys depend on it), connectors, andcustom_process_next_stage_input_func(controls embedding accumulation). DefaultTruepreserves PR 1 TTS behavior.Backward compatibility
All
Nonedefaults — orchestrator-managed deployments see no behavior change.SharedMemoryConnector(noget_connection_info) returnsNone, correctly preventing cross-pod handles for same-node-only connectors.Limitations
Test Plan
clear_connectors=Falsepreserves stage_id, connectors, producer hookclear_connectors=Trueclears (backward compat)metadatakwarg toconnector.get()Test Result
23/23 unit tests passed. TTS E2E verified on 2x NVIDIA H200.
Follow-on work
Async-chunk streaming for standalone stages, so the upstream stage streams frames directly to the downstream stage and recovers the ~92% TTFP reduction that standalone mode loses today.
Image gen disagg for BAGEL and HunyuanImage3, which needs CFG companion fan-out/fan-in.
This PR reuses
kv_transfer_paramsfor omni embedding transfer handles since the field is already wired end-to-end. If the community prefers a dedicatedstage_transfer_paramsfield, the change is mechanical.