[Feature] Transfer AR-to-diffusion payloads over NIXL (#6264)

Signed-off-by: yuanwu <yuan.wu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/design/feature/omni_connectors/nixl_connector.md

@@ -1,5 +1,8 @@
 # NixlConnector
 
+For the reusable model-runner API, illustrated control flow, deadline semantics,
+and failure/ownership contracts, see [Cross-Stage Payload Transport](stage_payload.md).
+
 ## When to Use
 
 Multi-node or intra-node stage transfer over NIXL, which brokers RDMA/shared-memory
@@ -46,6 +49,15 @@ provides a matching torch-xpu, UCX with Level Zero support, and NIXL runtime.
 
 ## Configuration
 
+MiniMax-H3's default disaggregated and Turbo deploy configs use shared memory
+for same-host transfer. To opt into NIXL, select
+`vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml` or
+`vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml` with
+`--deploy-config`. These overlays inherit the corresponding model and
+parallel settings and switch both ends of the stage edge. They require NIXL
+and a working UCX backend; initialization failures do not automatically fall
+back to shared memory.
+
 ```yaml
 connectors:
   nixl_connector:
@@ -92,9 +104,44 @@ Parameters:
 - `transfer_timeout_s`: how long a `get()` waits for its `READ` to complete
   (default 300). NIXL 1.3 has no transfer cancellation API, so a timed-out transfer's
   buffers and registrations remain owned by the connector until NIXL reports a
-  terminal state; `close()` waits for that state before releasing them.
+  terminal state. `close()` polls remaining transfers once and returns without
+  releasing active DMA resources. A serialized background closer automatically
+  finishes cleanup after completion; no second call is required.
   `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` overrides it.
 
+### Diffusion worker integration
+
+Diffusion initializes `OmniConnectorModelRunnerMixin` in synchronous mode,
+borrowing its KV transfer manager's lazy connector. This mode creates no second
+connector, starts no receive/save threads, and leaves KV key builders and
+sharding callbacks unchanged. Connector lifetime remains with the manager.
+
+The shared `recv_stage_payload` entry point and asynchronous `_poll_single_request`
+use the same key and endpoint metadata builder. Diffusion passes its actual
+incoming edge and external request ID; explicit transfer handles take precedence.
+Sender endpoints are passed per request, without mutating KV sender state.
+Diffusion-specific prompt merging, output handles and device placement live in
+`diffusion/worker/stage_payload.py`.
+
+Diffusion supplies ordered TP and SP groups to the shared payload fanout.
+Only the rank leading both groups reads the connector. Each subsequent broadcast
+uses its own group-local rank zero, including ranks that received the payload
+from the preceding group. Tensor payloads use tensor-dictionary broadcasts rather
+than serializing the full payload as a Python object. Receive misses still reach
+the same collectives. Inline fallback is allowed only when all required payload
+keys remain present; otherwise each rank raises an explicit payload error.
+
+The synchronous path passes one monotonic deadline to `get_with_deadline()` for
+discovery and DMA completion waits, rather than calling the ordinary `get()`
+with its potentially much longer transfer timeout. Native calls, serialization,
+and distributed collectives are not preemptible by this deadline.
+
+`GroupCoordinator.broadcast_tensor_dict` preserves the group-local source for
+`broadcast_object` while passing the mapped global source to PyTorch tensor
+collectives. For subgroup `[1, 3]`, local source `0` maps to global source `1`;
+passing `1` back to `broadcast_object` would incorrectly select global rank `3`.
+The nonzero-subgroup regression covers both sender and receiver behavior.
+
 ### Source ownership and failure limits
 
 Claims and expiry use the same lock. Completion identifies both a payload
@@ -105,12 +152,24 @@ transfers retain the complete destination bundle and source claim until terminal
 Zero-element tensors retain their shape/dtype/skeleton positions but never create
 DMA descriptors or registrations.
 
-A lost metadata response, lost completion ACK, or abandoned consumer can retain a
-claim indefinitely. NIXL 1.3 cannot prove remote cancellation, so neither TTL nor
+After a lost metadata response, retries on the same consumer thread reuse the
+claim ID for that endpoint, key and requested generation until a reply arrives.
+This makes claim acquisition idempotent; subsequent independent reads still get
+distinct claims. Callers retry on the same thread and call `abandon_get()` when
+giving up. Synchronous stage payload receive does this automatically. Background
+recovery reacquires the same claim solely to acknowledge it, without issuing a
+READ. Lost completion ACKs are retried against the exact payload generation.
+A permanently unreachable or dead consumer can still retain a claim indefinitely.
+NIXL 1.3 cannot prove remote cancellation, so neither TTL nor
 `cleanup()` frees those allocations. Producer `close()` rejects new work and
-retains its agent, listener and claimed allocations; a later `close()` finishes
+retains its agent, listener and claimed allocations; the background closer finishes
 teardown after claims drain. Permanently abandoned claims remain until process
-exit. This favors memory safety over bounded shutdown/memory usage. Only trusted
+exit. A consumer with unfinished local transfers likewise retains a strong reference
+to its connector, agent, tensors and registrations after `close()` returns. Its
+background closer polls and releases resources only after a terminal state;
+permanently stuck transfers retain resources until process exit. Closing connectors
+reject new work and report unhealthy. This avoids an unbounded transfer-polling loop
+in shutdown without freeing DMA-owned memory. Only trusted
 peers may access this unauthenticated control plane. Both endpoints must use the
 claim-aware protocol; rolling compatibility with older GET_META clients is not
 provided. Legacy externally-owned metadata without a generation is accepted only

docs/design/feature/omni_connectors/stage_payload.md

@@ -0,0 +1,280 @@
+# Cross-Stage Payload Transport
+
+Stage payload transport moves model outputs such as encoder conditioning between
+workers without routing their tensor contents through the orchestrator. It is
+separate from native diffusion KV transfer: a deployment can use both, either, or
+neither. MiniMax-H3 is one consumer of this interface, not its implementation.
+
+## Architecture
+
+```mermaid
+flowchart TB
+    Config[Topology and deploy configuration] --> Keys[Input and output payload keys]
+    Keys --> Producer[Producer model runner]
+    Keys --> Consumer[Consumer model runner]
+    Producer --> Adapter[DiffusionStagePayloadMixin]
+    Consumer --> Adapter
+    Adapter --> Runtime[Shared OmniConnector runtime]
+    Runtime --> Descriptor[Request key and endpoint resolution]
+    Runtime --> Fanout[TP then SP or HSDP shard then replicate fanout]
+    Runtime --> Manager[OmniKVTransferManager owns lazy connector]
+    Manager --> SHM[SharedMemoryConnector]
+    Manager --> NIXL[NixlConnector]
+    Producer -. small transfer handle .-> Orchestrator[Orchestrator]
+    Orchestrator -. prompt and handle .-> Consumer
+    Consumer --> NativeKV[Separate native diffusion KV path]
+```
+
+The diffusion adapter owns model-facing selection and validation. The shared
+runtime owns request-key construction, endpoint metadata, retry policy, and
+distributed fanout. Backends own storage, transfer progress, and allocation
+lifetime. No pipeline should instantiate a NIXL agent or implement its own
+metadata-query loop.
+
+The synchronous runtime borrows the manager's connector. It does not create
+background AR receive/save loops or replace the manager's KV key callbacks.
+The asynchronous AR path continues using `get()` through its existing poller.
+Both paths share the descriptor and fanout helpers; their scheduling policies
+are deliberately different.
+
+## Configuration and Identity
+
+`stage_output_payload_keys` selects non-`None` entries in a producer's
+`DiffusionOutput.custom_output`. `stage_input_payload_keys` declares the required
+entries in the consumer's `prompt.additional_information`. Both normalize to
+tuples in `OmniDiffusionConfig`, including lists loaded from YAML.
+
+The model topology declares these keys; deploy configuration selects the
+connector for the edge. MiniMax-H3's standard and Turbo disaggregated profiles
+default to SharedMemory. Their `_nixl.yaml` overlays opt into NIXL with
+`--deploy-config`; choosing NIXL does not imply fallback to SharedMemory.
+
+Without an explicit handle, the key is:
+
+```text
+{external_req_id or request_id}_{from_stage}_{chunk_id}
+```
+
+The synchronous complete-payload path uses chunk zero. Source and destination
+IDs come from the configured edge, not from an assumption that stages are
+adjacent or that the producer is stage zero. An explicit handle can override
+the key, source, destination, and backend metadata:
+
+```python
+{
+    "key": "request-42_2_0",
+    "from_stage": "2",
+    "to_stage": "5",
+    "size_bytes": 4096,
+    "metadata": {...},
+    "payload_keys": ["text_encoder_output"],
+}
+```
+
+This descriptor is attached under `_stage_payload_transfer`. Backend metadata
+is opaque to the pipeline. If it is absent, request-scoped sender information
+can supply `source_host` and `source_port`. `payload_sender_info` takes
+precedence over legacy `kv_sender_info`. A flat endpoint is accepted directly;
+a stage-keyed map is resolved using the actual source stage, accepting string
+and integer keys. A single unambiguous candidate remains supported. Multiple
+unmatched candidates are not arbitrarily resolved to stage zero.
+
+## API Contracts
+
+| API | Owner | Contract |
+| --- | --- | --- |
+| `init_omni_connectors(..., synchronous=True)` | Shared runtime | Borrow an existing manager; no AR polling threads or KV callback mutation. |
+| `_stage_payload_recv_spec(...)` | Shared transport | Resolve explicit handle or conventional key plus request-scoped endpoint. |
+| `recv_stage_payload(..., retry_seconds=2.0)` | Shared transport | Leader retries with one monotonic deadline, retires discovery, then fans out success or absence. |
+| `get_with_deadline(..., deadline=...)` | Connector | Bound backend-controlled waits by the supplied absolute monotonic deadline. Return `(payload, bytes)` or `None`. |
+| `abandon_get(key)` | Connector | Retire unresolved discovery on the calling thread after the retry loop. Never free an active READ. Default is a no-op for stateless discovery. |
+| `_maybe_recv_stage_payload(request)` | Diffusion adapter | Skip warmup, receive, merge declared keys on the target device, and reject missing required values. |
+| `_maybe_send_stage_payload(requests, outputs)` | Diffusion adapter | Leader publishes selected complete outputs; all ranks attach the same handle and remove only transferred fields. |
+| `OmniKVTransferManager.close()` | Manager | Drain prefetch, close only the connector already created, and disable later lazy initialization. |
+
+`get_with_deadline()` is intentionally not a wrapper around an arbitrary
+blocking `get()`. The base implementation raises `NotImplementedError`.
+An unsupported backend must implement this capability before being used for
+synchronous remote-only payloads. Existing asynchronous `get()` users are
+unchanged.
+
+The deadline covers retry sleeps, metadata socket waits, and NIXL completion
+polling. It is **not a hard real-time bound** on Python execution, allocation,
+serialization, individual native driver calls, device copies, or payload-group
+collectives. A stuck native call cannot safely be preempted with a Python
+future timeout. The implementation does not start a detached READ and discard
+its ownership when the model-thread budget expires.
+
+## Successful Transfer
+
+The synchronous diffusion producer uses an explicit handle:
+
+```mermaid
+sequenceDiagram
+    participant P as Producer leader
+    participant B as Connector
+    participant O as Orchestrator
+    participant C as Consumer leader
+    participant R as Consumer peer ranks
+    P->>B: put(key, selected custom_output)
+    B-->>P: success, bytes, backend metadata
+    P->>P: Broadcast handle; remove transferred inline keys
+    P->>O: Output carrying small handle
+    O->>C: Prompt carrying handle and endpoint
+    C->>B: get_with_deadline(key, deadline)
+    B-->>C: Complete payload or absence
+    C->>B: abandon_get(key)
+    C->>R: Delivery status, then tensors if present
+    C->>C: Merge and validate required keys
+    R->>R: Merge and validate required keys
+```
+
+MiniMax-H3 instead uses the AR encoder's asynchronous sender. It does not
+forward the diffusion producer's explicit handle:
+
+```mermaid
+sequenceDiagram
+    participant A as AR encoder leader
+    participant S as Background sender
+    participant B as Connector
+    participant O as Orchestrator
+    participant C as Diffusion consumer leader
+    participant R as Consumer peer ranks
+    A->>S: Queue encoder conditioning for conventional key
+    A->>O: Output without inline conditioning
+    O->>C: Request identity and request-scoped sender endpoint
+    par Asynchronous publication
+        S->>B: put(conventional key, conditioning), with retries
+    and Bounded receive
+        C->>B: get_with_deadline(conventional key, endpoint)
+        B-->>C: Payload or absence
+    end
+    C->>B: abandon_get(key)
+    C->>R: Delivery status, then tensors if present
+    C->>C: Merge and validate encoder_output
+    R->>R: Merge and validate encoder_output
+```
+
+Within a TP-by-SP payload grid, exactly one rank, the leader in both groups,
+accesses the connector. The shared broadcaster traverses TP and then SP in
+the same order on every rank: TP first populates the source SP row, then SP
+propagates each TP column.
+
+When initialized HSDP groups span the stage, the adapter instead uses the
+HSDP shard and replicate dimensions. The shard broadcast populates the first
+replica row; the replicate broadcast propagates every shard position to the
+other rows. Both dimensions participate in leader election. This includes
+standalone HSDP, whose TP/SP groups are singletons, as well as pure sharding
+and pure replication. The two HSDP dimensions replace, rather than overlap,
+the TP/SP fanout.
+
+The single-consumer guarantee is scoped to these configured payload groups;
+it is not a global election across arbitrary independent DP/CFG/PP groups or
+stage replicas. Followers do not initialize or query a connector. Absence is
+broadcast too, so peers do not wait for tensor broadcasts that the leader will
+never issue.
+
+Pipelines must gather sharded producer outputs before publication. Selecting
+payload keys does not gather tensor shards automatically.
+
+## Execution and Failure Policy
+
+Full forward receives before pipeline execution and sends after output
+preparation. Stepwise execution receives once when a new cached state is
+created. That state retains the external request ID; final decoded output is
+published exactly once. Intermediate denoising steps and non-final streaming
+chunks are not published under the complete-payload key. Chunk-addressed
+streaming transfer remains the asynchronous path's separate contract.
+
+Dummy request IDs (`dummy_req_id` and its slash-prefixed variants) bypass both
+directions. A pure warmup batch does not lazily initialize the connector,
+publish buffers, wait for remote data, or enter payload fanout.
+
+For the synchronous diffusion producer, failed `put()` calls preserve inline
+fields; only successful publications replace those fields with a handle.
+This guarantee does not apply to H3's AR encoder asynchronous sender. That
+path omits inline conditioning before the background `put()` completes;
+exhausting its send retries does not restore conditioning to the forwarded
+output. It uses the conventional key and request-scoped sender endpoint,
+not the diffusion producer's handle-based handoff.
+
+If a receive fails, inline continuation is valid only when **all required
+keys have non-`None` values**. Missing required values raise an explicit
+stage-payload error on each rank after fanout. H3 with omitted inline
+conditioning therefore fails explicitly if remote publication or delivery
+fails; the original prompt is not equivalent conditioning. With no declared
+input keys, an explicit handle's `payload_keys` defines the required set.
+
+## NIXL Ownership and Recovery
+
+```mermaid
+stateDiagram-v2
+    [*] --> Published: put registers source
+    Published --> Released: unclaimed lease expires
+    Published --> Claimed: atomic metadata query
+    Claimed --> Reading: consumer submits READ
+    Claimed --> Recovering: reply lost and caller abandons
+    Recovering --> Recovering: query or ACK lost
+    Recovering --> Released: same claim recovered and ACKed
+    Reading --> Deferred: deadline expires while DMA may be active
+    Reading --> Released: terminal DMA and completion ACK
+    Deferred --> Deferred: active or unknown DMA state
+    Deferred --> Released: terminal state and completion ACK
+```
+
+Claims are scoped by payload generation and UUID. Query timeouts or lost
+metadata replies retain the cached UUID for retries, since the producer may
+already have granted ownership. An explicit `META_NOT_FOUND` reply clears
+the cached claim; a subsequent query can use a new UUID. At the end of the
+synchronous retry loop, `abandon_get()`
+moves unresolved queries to background recovery. Recovery asks for the same
+claim, obtains the generation, and acknowledges it **without submitting a
+READ**. A lost completion reply is retried against that exact generation,
+not by acquiring a fresh payload with the same key. Existing claims may be
+recovered while a producer is closing; new claims may not.
+
+If a READ was submitted and its state is active or unknown, timeout does not
+release source or destination allocations. The deferred-transfer reaper keeps
+tensors, registrations, descriptor lists, and transfer handles until NIXL
+reports a terminal state. Completion ACK failure also stays in background
+retry. Producer close can therefore remain deferred until ownership drains.
+
+Recovery is eventual, not guaranteed under permanent network partition or
+consumer process death. A producer cannot safely infer that an unreachable
+consumer has stopped DMA. Conservatively retained ownership in that case is
+an explicit safety limit, distinct from a live consumer silently abandoning
+a recoverable discovery claim.
+
+## Reusing the Interface
+
+| Backend | Synchronous support | Backend responsibility |
+| --- | --- | --- |
+| SharedMemory | Implemented | Nonblocking read-lock acquisition; return absence on contention; shared runtime retries. Copy and deserialize after acquiring the lock. |
+| NIXL | Implemented | Deadline-aware socket and completion waits, generation-scoped claims, deferred active-DMA cleanup, discovery recovery. |
+| Other existing connectors | Opt-in required | Implement `get_with_deadline`; implement `abandon_get` only when discovery creates ownership. Do not inherit an unbounded blocking call. |
+
+To add a model, declare input/output keys and use ordinary prompt and output
+structures. To add a backend, implement the connector capability and ownership
+tests, register it through the existing factory, and select it in deploy
+configuration. Neither requires a backend-specific diffusion model runner.
+
+Worker teardown disables offloading, closes the manager after draining
+prefetch, then tears down native KV and distributed resources in `finally`
+blocks. The synchronous mixin does not close the borrowed connector itself.
+Repeated manager close is safe, including for an unused manager.
+
+## Validation Map
+
+| Test module | Main invariants |
+| --- | --- |
+| `tests/diffusion/test_diffusion_stage_payload.py` | Warmup bypass, strict missing-payload errors, valid inline fallback, source-stage endpoints, TP/SP fanout, standalone HSDP election and SHM fanout, handle publication. |
+| `tests/diffusion/test_diffusion_model_runner.py` | New-request receive and final-only stepwise publication with cached external identity. |
+| `tests/distributed/omni_connectors/test_nixl_connector.py` | Lost replies, same-claim recovery, deadlines, generation isolation, active DMA retention, deferred close and ACKs. |
+| `tests/distributed/omni_connectors/test_shm_connector.py` | Key and metadata reads, lock contention, expired deadline, cleanup. |
+| `tests/worker/test_omni_connector_mixin.py` | Borrowed connector ownership and shared synchronous/asynchronous helpers. |
+| `tests/distributed/omni_connectors/test_kv_async_prefetch.py` | Prefetch ordering and idempotent manager close. |
+| `tests/config/test_omni_config.py` | Tuple normalization across topology, deploy, stage CLI, and serialization. |
+
+Mocked ownership tests do not establish native transport performance or
+accelerator correctness. Native NIXL and model E2E runs are separate checks;
+their results must name the tested revision and runtime.

recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md

@@ -2,10 +2,40 @@
 
 This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and
 audio WVAE encoder together in a vLLM stage. It sends their conditioning
-payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes
+payload over shared memory by default, or opt-in NIXL, to an encoder-free
+diffusion stage. The standard MiniMax H3 recipes
 remain single-stage and continue to run all components inside the diffusion
 pipeline.
 
+## Full encoder payload
+
+Both the standard and Turbo builtin deployments use `SharedMemoryConnector`
+for the same-host stage 0 → stage 1 edge. The corresponding `_nixl.yaml`
+overlays opt into `NixlConnector` with the UCX backend. Stage 0 is the sender and stage 1
+is the receiver. The producer hook `encoder2diffusion_full_payload` validates
+the complete conditioning and packages it under `encoder_output`; the diffusion
+stage declares `stage_input_payload_keys: [encoder_output]`.
+
+This is not a text-only handoff. The payload includes:
+
+- `hidden_states.output`: BF16 text embeddings.
+- `meta.token_role_ids`: INT64 token roles aligned with the text embeddings.
+- `embed.embedding`: FP32 visual conditioning latents when present.
+- `embed.speech_feat`: FP32 audio conditioning latents when present.
+- `kv_metadata.minimax_h3_encoder_layout`: packed layout metadata describing
+  the task, dimensions, condition shapes/lengths, keyframes, and reference order.
+
+Text-only requests still carry the full encoder schema, with empty optional
+media tensors. The H3 layout field is model conditioning metadata, **not**
+native paged-KV transfer. Stage 1 does not request native KV-cache reception.
+All participating encoder ranks complete their component collectives before
+the producer sends. The diffusion group leader receives the payload and the
+generic stage-payload path distributes it to the participating diffusion ranks.
+After a successful connector send, the orchestrator carries no duplicate inline
+conditioning; original image/video/audio inputs and temporary encoder media
+buffers are removed from the diffusion prompt. Stage 1 consumes the received
+`encoder_output` without re-encoding text or media.
+
 ## Start the server
 
 Choose the topology explicitly and load its deployment defaults:
@@ -32,18 +62,25 @@ parallelism, and `audio_vae` runs on the group leader. All three role entries
 are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit;
 the configuration does not create independent role schedulers or world sizes.
 
-For example, this five-GPU topology assigns one GPU to the encoder and four
-to the diffusion stage. `--stage-overrides` keeps placement and parallelism
-scoped to the owning stage rather than broadcasting an override to both:
+For a six-GPU TP2 encoder → TP4 diffusion topology, retain the builtin
+deployment and override only stage 1's TP and Ulysses degrees. VAE patch
+parallelism stays at 4. `--stage-overrides` keeps parallelism scoped to its
+owning stage rather than broadcasting an override to both:
 
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 \
 vllm-omni serve MiniMaxAI/MiniMax-H3 \
   --omni \
   --deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \
-  --stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}'
+  --stage-overrides '{"1":{"tensor_parallel_size":4,"ulysses_degree":1}}'
 ```
 
+The same override is supported with the Turbo builtin deployment. Device IDs
+in the deployment are relative to `CUDA_VISIBLE_DEVICES`; choose visible
+devices and per-stage `devices` overrides for the available hardware. These
+commands describe configuration, not a claim of hardware validation for this
+integrated encoder implementation.
+
 For memory-constrained deployments, start from the CPU-offload or distributed
 layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and
 quantization options only to Stage 1 with `--stage-overrides`; retain the
@@ -73,8 +110,35 @@ media conditioning from Stage 0, with no fallback to local encoding.
 
 This H3 topology explicitly keeps its single-replica
 diffusion stage inline, avoiding serialization of decoded video through a
-subprocess. It expects both stages to run in one deployment; cross-node payload
-transport is outside this configuration.
+subprocess. The builtin placement describes a single-host deployment. NIXL's
+transport design can support remote endpoints, but **cross-node H3 execution
+has not been validated** here; these examples are not a tested multi-node
+launch procedure. CPU wiring and payload-contract tests do not establish native
+NIXL, GPU, Turbo-weight, or cross-node execution success.
+
+## Optional NIXL configuration
+
+The base deployments above use `SharedMemoryConnector` and require same-host
+shared-memory access. To use NIXL instead, select the standard overlay:
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --omni \
+  --deploy-config vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml
+```
+
+For Turbo, use
+`--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml`
+with the LoRA options shown below. Both overlays require NIXL and a working
+UCX backend; see [NIXL installation](../../docs/design/feature/omni_connectors/nixl_connector.md#installation).
+Initialization failures do not automatically fall back to shared memory.
+
+Each overlay inherits its base deployment and switches **both** edge
+references to NIXL. The encoder roles, `model_loaded` flags, hooks, and full
+`encoder_output` schema stay unchanged. The processor also retains support for
+an inline full-conditioning handoff when supplied; stage 1 still does not fall
+back to local encoders. SharedMemory requires same-host shared-memory access
+and is not a cross-node transport.
 
 The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
 `fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
@@ -99,6 +163,6 @@ that artifact's own contract -- see the
 [Turbo LoRA table](MiniMax-H3.md#turbo-lora) for every published row.
 
 The base deployment intentionally retains 50 inference steps for non-LoRA
-quality. This disaggregated topology is validated on the `fl2v` artifacts,
-which serve T2VA and FL2VA; the `ref2v` artifacts need a `--task-type ref2va`
-server.
+quality. The `fl2v` artifacts serve T2VA and FL2VA; the `ref2v` artifacts need
+a `--task-type ref2va` server. Standard and Turbo configuration coverage does
+not imply that Turbo weights have been exercised on this integrated topology.

tests/config/test_config_factory.py

@@ -1588,19 +1588,32 @@ stages:
         assert "text_encoder_tp_size" not in stages[1].runtime_overrides
         assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
 
-    def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
+    @pytest.mark.parametrize(
+        "deploy_name",
+        ["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
+    )
+    def test_minimax_h3_disaggregated_defaults(self, deploy_name):
         pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
-        deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
+        deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
         stages = merge_pipeline_deploy(pipeline, deploy)
 
+        assert pipeline.model_arch == "MiniMaxH3Encoder"
+        assert stages[0].model_stage == "encoder"
         assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder"
+        processor = "vllm_omni.model_executor.stage_input_processors.minimax_h3"
+        assert stages[0].yaml_extras["prompt_transform_func"] == f"{processor}.prepare_encoder_prompt"
+        assert (
+            stages[0].yaml_engine_args["custom_process_next_stage_input_func"]
+            == f"{processor}.encoder2diffusion_full_payload"
+        )
+        assert stages[1].custom_process_input_func == f"{processor}.encoder2diffusion"
+        assert tuple(stages[1].yaml_engine_args["stage_input_payload_keys"]) == ("encoder_output",)
+        assert stages[1].yaml_engine_args["omni_kv_config"]["need_recv_cache"] is False
         assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline"
         assert stages[0].yaml_runtime["num_replicas"] == 1
         assert stages[1].yaml_runtime["num_replicas"] == 1
-        assert stages[1].yaml_engine_args["model_loaded"] == {
-            "text_encoder": False,
-            "vae_encoder": False,
-        }
+        assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False, "vae_encoder": False}
+        assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
         assert stages[0].yaml_engine_args["max_num_seqs"] == 1
         assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == {
             "text_encoder": {"parallel_mode": "tp"},
@@ -1614,7 +1627,12 @@ stages:
         assert parallel["ulysses_degree"] == 4
         assert parallel["vae_patch_parallel_size"] == 4
         assert stages[1].yaml_engine_args["inline_diffusion"] is True
-        assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
+        sampling = stages[1].yaml_extras["default_sampling_params"]
+        if "turbo" in deploy_name:
+            assert sampling["num_inference_steps"] == 5
+            assert sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
+        else:
+            assert sampling["num_inference_steps"] == 50
 
         turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
         turbo_stages = merge_pipeline_deploy(pipeline, turbo)
@@ -1625,6 +1643,73 @@ stages:
         assert turbo_sampling["num_inference_steps"] == 5
         assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
 
+    @pytest.mark.parametrize(
+        "deploy_name",
+        ["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
+    )
+    @pytest.mark.parametrize("shared_memory", [False, True], ids=["nixl", "shared_memory"])
+    def test_minimax_h3_deploy_connector_wiring_and_roles(self, deploy_name, shared_memory):
+        from vllm_omni.config.stage_config import resolve_deploy_yaml
+        from vllm_omni.distributed.omni_connectors.utils.initialization import load_omni_transfer_config
+        from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+        base_config = resolve_deploy_yaml(get_deploy_config_path(deploy_name))
+        if not shared_memory:
+            deploy_name = deploy_name.replace(".yaml", "_nixl.yaml")
+        deploy_path = get_deploy_config_path(deploy_name)
+        config = resolve_deploy_yaml(deploy_path)
+        assert config["async_chunk"] is False
+        connector_name = "shared_memory_connector" if shared_memory else "nixl_connector"
+        assert connector_name in config["connectors"]
+        assert config["stages"][0]["output_connectors"] == {"to_stage_1": connector_name}
+        assert config["stages"][1]["input_connectors"] == {"from_stage_0": connector_name}
+        load_deploy_config(deploy_path)
+        for base_stage, stage in zip(base_config["stages"], config["stages"], strict=True):
+            edge_keys = {"input_connectors", "output_connectors"}
+            assert {key: value for key, value in stage.items() if key not in edge_keys} == {
+                key: value for key, value in base_stage.items() if key not in edge_keys
+            }
+        transfer = load_omni_transfer_config(config_dict=config)
+        assert transfer is not None
+        assert set(transfer.connectors) == {("0", "1")}
+        edge = transfer.connectors[("0", "1")]
+        expected_name = "SharedMemoryConnector" if shared_memory else "NixlConnector"
+        assert edge.name == expected_name
+        assert "role" not in edge.extra
+        for stage_id, role in ((0, "sender"), (1, "receiver")):
+            spec = get_stage_connector_spec(transfer, stage_id, async_chunk=False)
+            assert spec["name"] == expected_name
+            assert spec["extra"]["role"] == role
+            if shared_memory:
+                assert "from_stage" not in spec["extra"]
+                assert "to_stage" not in spec["extra"]
+            else:
+                assert spec["extra"]["from_stage"] == 0
+                assert spec["extra"]["to_stage"] == 1
+                assert spec["extra"]["host"] == "auto"
+                assert spec["extra"]["zmq_port"] == 50071
+                assert spec["extra"]["backends"] == ["UCX"]
+        assert "role" not in edge.extra
+
+    @pytest.mark.parametrize(
+        "deploy_name",
+        ["minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated_turbo.yaml"],
+    )
+    def test_minimax_h3_builtin_stage_one_tp4_ulysses1_override(self, deploy_name):
+        stages, _ = StageConfigFactory._create_legacy_from_registry(
+            OMNI_PIPELINES["minimax_h3_disaggregated"],
+            {"stage_1_tensor_parallel_size": 4, "stage_1_ulysses_degree": 1},
+            deploy_config_path=get_deploy_config_path(deploy_name),
+        )
+        encoder_args = stages[0].to_omegaconf().engine_args
+        diffusion_args = stages[1].to_omegaconf().engine_args
+        assert encoder_args.tensor_parallel_size == 2
+        assert "ulysses_degree" not in stages[0].runtime_overrides
+        assert diffusion_args.parallel_config.tensor_parallel_size == 4
+        assert diffusion_args.parallel_config.ulysses_degree == 1
+        assert diffusion_args.parallel_config.vae_patch_parallel_size == 4
+        assert tuple(diffusion_args.stage_input_payload_keys) == ("encoder_output",)
+
     def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
         pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
 

tests/config/test_omni_config.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import warnings
-from dataclasses import fields
+from dataclasses import fields, replace
 from inspect import Parameter, signature
 from multiprocessing.reduction import ForkingPickler
 from pathlib import Path
@@ -246,6 +246,12 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
             assert omni_stage.diffusion_config is not None
             assert omni_stage.diffusion_config.stage_id == legacy_stage.stage_id
             assert omni_stage.diffusion_config.model_arch == engine_args.get("model_arch")
+            assert omni_stage.diffusion_config.stage_input_payload_keys == engine_args.get(
+                "stage_input_payload_keys", ()
+            )
+            assert omni_stage.diffusion_config.stage_output_payload_keys == engine_args.get(
+                "stage_output_payload_keys", ()
+            )
         elif omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR:
             assert isinstance(omni_stage, VllmOmniARStageConfig)
             assert not hasattr(omni_stage, "diffusion_config")
@@ -1725,6 +1731,58 @@ def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
         )
 
 
+@pytest.mark.parametrize("source", ["default", "topology", "deploy", "stage-cli"])
+@pytest.mark.parametrize("key_container", [list, tuple])
+def test_diffusion_stage_payload_keys_roundtrip(source, key_container):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+    topology_keys = {
+        "stage_input_payload_keys": ("conditioning", "metadata"),
+        "stage_output_payload_keys": ("latents",),
+    }
+    override_keys = {
+        "stage_input_payload_keys": (),
+        "stage_output_payload_keys": ("audio", "video"),
+    }
+    pipeline = _resolve_pipeline_or_skip("dreamzero")
+    topology = pipeline.get_stage(0)
+    if source != "default":
+        topology = replace(topology, **topology_keys)
+    pipeline = replace(pipeline, stages=(topology,))
+    deploy = _load_default_deploy(pipeline)
+    cli_overrides = {}
+    if source == "deploy":
+        deploy = replace(deploy, stages=[StageDeployConfig(stage_id=0, engine_extras=override_keys)])
+    elif source == "stage-cli":
+        cli_overrides = {f"stage_0_{name}": value for name, value in override_keys.items()}
+
+    stage = VllmOmniConfig.from_pipeline_config(
+        pipeline, user_deploy_config=deploy, cli_overrides=cli_overrides
+    ).stage_by_id(0)
+    expected = topology_keys if source == "topology" else override_keys
+    if source == "default":
+        expected = dict.fromkeys(topology_keys, ())
+    legacy_stage = merge_pipeline_deploy(pipeline, deploy)[0]
+    legacy_args = {**legacy_stage.yaml_engine_args, **(override_keys if source == "stage-cli" else {})}
+    restored_stage = ForkingPickler.loads(ForkingPickler.dumps(stage))
+    engine_args = build_engine_args_dict_from_omni_stage_config(restored_stage, model="test-model")
+    diffusion_kwargs = omni_config_module.extract_diffusion_stage_config_kwargs(
+        engine_args, stage_id=restored_stage.stage_id, include_engine_adapter_metadata=True
+    )
+    for name in topology_keys:
+        diffusion_kwargs[name] = key_container(diffusion_kwargs[name])
+    od_config = OmniDiffusionConfig.from_kwargs(**diffusion_kwargs)
+
+    for name, keys in expected.items():
+        assert legacy_args.get(name, ()) == keys
+        assert getattr(stage.diffusion_config, name) == keys
+        assert engine_args[name] == keys
+        assert getattr(od_config, name) == keys
+    for name in topology_keys:
+        assert getattr(topology, name) == (() if source == "default" else topology_keys[name])
+
+
 def test_diffusion_config_field_classification_covers_current_fields():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -1745,6 +1803,8 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "diffusion_kv_cache_dtype",
         "diffusion_kv_mode",
         "diffusion_kv_max_rows_per_request",
+        "stage_input_payload_keys",
+        "stage_output_payload_keys",
     } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
     assert {
         "revision",

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -146,3 +146,34 @@ def test_set_seq_parallel_pg_validates_sp_group_ranks(monkeypatch):
             world_size=4,
             sp_group_ranks=[[0, 2]],
         )
+
+
+@pytest.mark.parametrize("rank", [1, 3])
+def test_tensor_dict_broadcast_preserves_subgroup_source(rank, monkeypatch):
+    from vllm_omni.diffusion.distributed.group_coordinator import GroupCoordinator, TensorMetadata
+
+    coordinator = object.__new__(GroupCoordinator)
+    coordinator.rank = rank
+    coordinator.ranks = [1, 3]
+    coordinator.rank_in_group = [1, 3].index(rank)
+    coordinator.world_size = 2
+    coordinator.device_group = object()
+    coordinator.cpu_group = object()
+    metadata = [("payload", TensorMetadata("cpu", torch.float32, (2,)))]
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+    def broadcast_object(value, src):
+        assert src == 0
+        return metadata
+
+    def broadcast(tensor, src, group, async_op):
+        assert src == 1
+        assert group is coordinator.cpu_group
+        assert async_op
+        tensor.fill_(7)
+        return SimpleNamespace(wait=lambda: None)
+
+    monkeypatch.setattr(coordinator, "broadcast_object", broadcast_object)
+    monkeypatch.setattr(torch.distributed, "broadcast", broadcast)
+    result = coordinator.broadcast_tensor_dict({"payload": torch.full((2,), 7.0)} if rank == 1 else None)
+    assert torch.equal(result["payload"], torch.full((2,), 7.0))

tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py

@@ -89,6 +89,23 @@ def test_launch_failure_disables_runtime_key(monkeypatch):
     assert len(fused_rope._FAILED_KEYS) == 1
 
 
+def test_non_cuda_platform_does_not_launch_fused_kernel(monkeypatch):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    monkeypatch.setattr(fused_rope, "HAS_TRITON", True)
+    monkeypatch.setattr(fused_rope.current_omni_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+
+    def unexpected_launch(*args):
+        raise AssertionError("non-CUDA platforms must remain on the native path")
+
+    monkeypatch.setattr(fused_rope, "_launch_fused_qk_rotary_emb", unexpected_launch)
+    query = torch.zeros((1, 1, 1, 128), dtype=torch.bfloat16)
+    freqs = torch.zeros((1, 1, 64), dtype=torch.float32)
+    assert fused_rope.try_fused_qk_rotary_emb(query, query, freqs, freqs) is None
+    assert not fused_rope._FAILED_KEYS
+
+
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_failed_runtime_key_cache_is_bounded():
     from vllm_omni.diffusion.models.ernie_image import fused_rope

tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py

@@ -10,6 +10,7 @@ import torch
 from PIL import Image
 
 from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -42,8 +43,10 @@ def pipeline():
     return model
 
 
-def _request(task):
+def _request(task, *, payload=None):
     prompt: dict[str, Any] = {"prompt": "A person waves."}
+    if payload is not None:
+        prompt["additional_information"] = {"encoder_output": payload}
     if task != "t2va":
         prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))}
     sampling = OmniDiffusionSamplingParams(
@@ -64,7 +67,120 @@ def test_single_stage_encodes_without_stage_zero(pipeline, task):
     torch.testing.assert_close(audio, pipeline.decode.return_value[1])
 
 
-def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch):
+def test_single_stage_reuses_legacy_text_but_encodes_media(pipeline):
+    payload = {
+        "hidden_states": torch.full((2, 5120), 3.0, dtype=torch.bfloat16),
+        "token_tags": torch.ones(2, dtype=torch.long),
+    }
+    pipeline.forward(_request("fl2va", payload=payload))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_called_once()
+    torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["text_embeddings"], payload["hidden_states"])
+
+
+def test_text_disaggregation_keeps_local_media_encoding(pipeline):
+    payload = {
+        "hidden_states": torch.full((2, 5120), 3.0, dtype=torch.bfloat16),
+        "token_tags": torch.ones(2, dtype=torch.long),
+    }
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = True
+    pipeline.forward(_request("fl2va", payload=payload))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_called_once()
+
+
+@pytest.mark.parametrize("payload", [None, {}, {"hidden_states": torch.ones(2, 5120), "token_tags": torch.ones(2)}])
+def test_stage_one_rejects_missing_or_legacy_conditioning(pipeline, payload):
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    with pytest.raises(OmniClientError):
+        pipeline.forward(_request("t2va", payload=payload))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_not_called()
+    pipeline.diffuse.assert_not_called()
+
+
+@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"])
+@pytest.mark.parametrize("load_text_encoder", [True, False])
+def test_prepare_encode_uses_the_same_mode_and_conditioning(pipeline, task, load_text_encoder):
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
+    request = _request(task)
+    if not load_text_encoder:
+        conditioning, window_text = pipeline._prepare_local_conditioning(request.prompts[0], request.sampling_params)
+        assert window_text is None
+        request.prompts[0] = {"additional_information": {"encoder_output": conditioning.to_omni_payload()}}
+        pipeline.encode_prompt.reset_mock()
+        pipeline.video_vae.encode_image.reset_mock()
+    pipeline.load_text_encoder = load_text_encoder
+    pipeline.load_vae_encoder = load_text_encoder
+    state = StepRequestState(request_id="step", prompt=request.prompts[0], sampling=request.sampling_params)
+    assert pipeline.prepare_encode(state) is state
+    assert state.latents is not None
+    assert state.latents.shape[1] == 96
+    assert state.total_steps == 1
+    assert pipeline.encode_prompt.call_count == int(load_text_encoder)
+    assert pipeline.video_vae.encode_image.call_count == int(load_text_encoder and task != "t2va")
+
+
+def test_stage_one_step_execution_never_falls_back(pipeline):
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    request = _request("t2va")
+    with pytest.raises(OmniClientError, match="requires encoder conditioning"):
+        pipeline.prepare_encode(
+            StepRequestState(request_id="step", prompt=request.prompts[0], sampling=request.sampling_params)
+        )
+    pipeline.encode_prompt.assert_not_called()
+
+
+def test_multiple_outputs_encode_only_once(pipeline):
+    request = _request("fl2va")
+    request.sampling_params.num_outputs_per_prompt = 3
+    request.sampling_params.seed = 42
+    pipeline.forward(request)
+    pipeline.encode_prompt.assert_called_once()
+    pipeline.video_vae.encode_image.assert_called_once()
+    assert [call.kwargs["seed"] for call in pipeline.diffuse.call_args_list] == [42, 43, 44]
+
+
+def test_ref2va_partition_preserves_image_only_implicit_task(pipeline):
+    request = _request("fl2va")
+    del request.sampling_params.extra_args["task"]
+    pipeline.partition = "ref2va"
+    pipeline.supported_tasks = frozenset({"ref2va"})
+    pipeline.forward(request)
+    assert pipeline.diffuse.call_args.kwargs["task"] == "ref2va"
+
+
+def test_invalid_legacy_text_is_not_ignored(pipeline):
+    with pytest.raises(OmniClientError, match="encoder wire requires"):
+        pipeline.forward(_request("fl2va", payload={"hidden_states": "invalid"}))
+    pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_not_called()
+
+
+def test_text_weight_loading_preserves_encoder_specific_lifecycle(pipeline):
+    class Encoder(torch.nn.Module):
+        def load_weights(self, weights):
+            return {name for name, _ in weights}
+
+    pipeline.text_encoder = Encoder()
+    pipeline.video_vae = torch.nn.Identity()
+    pipeline.audio_vae = torch.nn.Identity()
+    assert pipeline.load_weights([("text_encoder.layer.weight", torch.ones(1))]) == {"text_encoder.layer.weight"}
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    pipeline.text_encoder = None
+    with pytest.raises(ValueError, match="disabled in this deployment"):
+        pipeline.load_weights([("text_encoder.layer.weight", torch.ones(1))])
+
+
+@pytest.mark.parametrize("connector_handoff", [False, True], ids=["inline", "full_payload"])
+def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch, connector_handoff):
     import numpy as np
 
     from vllm_omni.data_entry_keys import flatten_payload
@@ -74,7 +190,11 @@ def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monke
         MiniMaxH3TextConditioning,
     )
     from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
-    from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion, prepare_encoder_prompt
+    from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
+        encoder2diffusion,
+        encoder2diffusion_full_payload,
+        prepare_encoder_prompt,
+    )
 
     frames = np.zeros((48, 32, 32, 3), dtype=np.uint8)
     decode_frames = Mock(return_value=frames)
@@ -110,11 +230,33 @@ def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monke
     media = stage_zero._encode_media(stage_zero._media_input(transformed["additional_information"]))
     text = MiniMaxH3TextConditioning(*pipeline.encode_prompt.return_value)
     payload = MiniMaxH3EncoderConditioning.from_components(text, media).to_omni_payload()
-    source = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=flatten_payload(payload))])
+    wire = flatten_payload(payload)
+    source = SimpleNamespace(
+        finished=True, outputs=[SimpleNamespace(multimodal_output=None if connector_handoff else wire)]
+    )
     bridged = encoder2diffusion([source], request.prompts[0])
+    assert bridged["multi_modal_data"] is None
+    if connector_handoff:
+        assert "encoder_output" not in bridged["additional_information"]
+        transferred = encoder2diffusion_full_payload(pooling_output=wire)
+        assert set(transferred) == {"encoder_output"}
+        received = transferred["encoder_output"]
+        for section, key in (
+            ("hidden_states", "output"),
+            ("meta", "token_role_ids"),
+            ("embed", "embedding"),
+            ("embed", "speech_feat"),
+            ("kv_metadata", "minimax_h3_encoder_layout"),
+        ):
+            torch.testing.assert_close(received[section][key], payload[section][key])
+        # Simulate the generic diffusion runner merging the received payload;
+        # native NIXL transport is covered separately, not by this CPU test.
+        bridged["additional_information"].update(transferred)
     pipeline.load_text_encoder = False
     pipeline.load_vae_encoder = False
     pipeline.encode_prompt.reset_mock()
+    pipeline.video_vae.encode_image.reset_mock()
+    pipeline.video_vae.encode_video.reset_mock()
     pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params))
     external_kwargs = pipeline.diffuse.call_args.kwargs
     for key, value in local_kwargs.items():
@@ -123,6 +265,8 @@ def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monke
         else:
             assert value == external_kwargs[key]
     pipeline.encode_prompt.assert_not_called()
+    pipeline.video_vae.encode_image.assert_not_called()
+    pipeline.video_vae.encode_video.assert_not_called()
     assert decode_frames.call_count == sampler.call_count == 2
     assert all(call.kwargs["decoded_frames"] is frames for call in sampler.call_args_list)
     assert [block["kind"] for block in local_kwargs["ref_blocks"]] == ["image", "video_audio", "audio"]

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -10,9 +10,17 @@ from unittest.mock import Mock
 import pytest
 import torch
 
+from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def _media(**kwargs):
+    return MiniMaxH3EncoderMediaInput(
+        task="ref2va", height=768, width=1344, num_frames=124, latent_t=37, audio_t=207, **kwargs
+    )
+
+
 def test_h3_full_pipeline_profiler_includes_local_encoding():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
@@ -28,6 +36,9 @@ def test_h3_full_pipeline_profiler_includes_local_encoding():
         "denoise_step",
         "post_decode",
     ]
+    assert "_encode_video_conditions" not in targets
+    assert "_encode_video_audio_conditions" not in targets
+    assert "_encode_audio_conditions" not in targets
 
 
 def test_decoded_output_is_unchanged_without_active_model_offload():
@@ -316,46 +327,6 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text
     remove_offload.assert_called_once_with([*dits, *stages])
 
 
-@pytest.mark.parametrize("decode_fails", [False, True])
-def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    component = torch.nn.Linear(2, 2)
-    events: list[tuple[Any, ...]] = []
-
-    @contextmanager
-    def record_component(value):
-        events.append(("activate", value))
-        try:
-            yield
-        finally:
-            events.append(("offload", value))
-
-    monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    pipeline._model_cpu_offload_modules = [component]
-
-    def decode():
-        with pipeline._component_on_device(component):
-            events.append(("decode", component))
-            if decode_fails:
-                raise RuntimeError("decode failed")
-
-    if decode_fails:
-        with pytest.raises(RuntimeError, match="decode failed"):
-            decode()
-    else:
-        decode()
-
-    assert events == [
-        ("activate", component),
-        ("decode", component),
-        ("offload", component),
-    ]
-
-
 def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
@@ -416,6 +387,46 @@ def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
     apply_offload.assert_not_called()
 
 
+@pytest.mark.parametrize("decode_fails", [False, True])
+def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    component = torch.nn.Linear(2, 2)
+    events: list[tuple[Any, ...]] = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    pipeline._model_cpu_offload_modules = [component]
+
+    def decode():
+        with pipeline._component_on_device(component):
+            events.append(("decode", component))
+            if decode_fails:
+                raise RuntimeError("decode failed")
+
+    if decode_fails:
+        with pytest.raises(RuntimeError, match="decode failed"):
+            decode()
+    else:
+        decode()
+
+    assert events == [
+        ("activate", component),
+        ("decode", component),
+        ("offload", component),
+    ]
+
+
 def _encode_media(pipeline, **kwargs):
     from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
     from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -416,7 +416,7 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames):
     monkeypatch.setattr(
         mod.MiniMaxH3Pipeline,
         "_prepare_encoder_conditioning_inputs",
-        lambda self, value, sampling: context,
+        lambda self, value, sampling: context if value is conditioning else pytest.fail("wrong encoder handoff"),
     )
     monkeypatch.setattr(
         mod.MiniMaxH3Pipeline,

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py

@@ -15,6 +15,21 @@ from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import decode_tempora
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.mark.parametrize("chunked", [False, True])
+def test_encode_only_video_vae_rejects_full_and_chunked_decode(chunked):
+    from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    torch.nn.Module.__init__(vae)
+    vae.encode_only = True
+    latent = torch.empty(1, 1, 1, 1, 1)
+    with pytest.raises(RuntimeError, match="encode-only video VAE cannot decode"):
+        if chunked:
+            vae.decode_with_chunks(latent, on_chunk=lambda frames: None)
+        else:
+            vae.decode_latent(latent)
+
+
 @pytest.fixture(scope="module")
 def cpu_process_group():
     if dist.is_initialized():

tests/diffusion/test_diffusion_model_runner.py

@@ -547,6 +547,9 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     runner.od_config.step_execution = True
     req = _make_request()
     req.request_id = "req"
+    req.external_req_id = "external-req"
+    sent = []
+    runner._maybe_send_stage_payload = lambda requests, outputs: sent.append((requests[0].external_req_id, outputs))
 
     monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
     monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
@@ -560,6 +563,7 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
 
     first = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     assert first.get_request_output("req").result is None
+    assert sent == []
 
     scheduler_output = SimpleNamespace(
         finished_req_ids=set(),
@@ -569,11 +573,13 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     second = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     assert second.get_request_output("req").result == chunks[0]
     assert second.get_request_output("req").finished is False
+    assert sent == []
 
     DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     fourth = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
     assert fourth.get_request_output("req").result == chunks[1]
     assert fourth.get_request_output("req").finished is True
+    assert sent == [("external-req", [chunks[1]])]
 
 
 @pytest.mark.core_model
@@ -733,6 +739,7 @@ def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(mon
     req = _make_request()
 
     monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
 
     output = DiffusionModelRunner.execute_model(runner, req)
 
@@ -749,6 +756,7 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
     req = _make_request()
 
     monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
 
     output = DiffusionModelRunner.execute_model(runner, req)
 

tests/diffusion/test_diffusion_stage_payload.py

@@ -0,0 +1,785 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+HANDLE_KEY = DiffusionModelRunner._STAGE_PAYLOAD_HANDLE_KEY
+
+
+class _FakeConnector:
+    def abandon_get(self, get_key):
+        pass
+
+    def get_with_deadline(self, from_stage, to_stage, get_key, metadata=None, *, deadline):
+        return self.get(from_stage, to_stage, get_key, metadata)
+
+    def __init__(self, payload=None, *, raises=False, put_result=(True, 128, {"schema_version": 1})):
+        self._payload = payload
+        self._raises = raises
+        self._put_result = put_result
+        self.calls: list[tuple[str, str, str, object]] = []
+        self.put_calls: list[tuple[str, str, str, object]] = []
+
+    def get(self, from_stage, to_stage, get_key, metadata=None):
+        self.calls.append((from_stage, to_stage, get_key, metadata))
+        if self._raises:
+            raise RuntimeError("transfer failed")
+        if self._payload is None:
+            return None
+        return self._payload, 0
+
+    def put(self, from_stage, to_stage, put_key, data):
+        self.put_calls.append((from_stage, to_stage, put_key, data))
+        if self._raises:
+            raise RuntimeError("transfer failed")
+        return self._put_result
+
+
+class _FakeKVTransferManager:
+    from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
+
+    _resolve_sender_info = OmniKVTransferManager._resolve_sender_info
+
+    def __init__(self, connector, recv_stages=("0", "1"), send_stages=("0", "1")):
+        self.connector = connector
+        self.recv_stages = recv_stages
+        self.send_stages = send_stages
+        self.sender_info_calls: list[tuple[dict, str | None]] = []
+
+    def update_sender_info(self, sender_info, sender_stage_id=None):
+        self.sender_info_calls.append((sender_info, sender_stage_id))
+
+
+def _make_runner(connector, *, payload_keys=("text_encoder_output",), recv_stages=("0", "1")):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(stage_input_payload_keys=payload_keys, stage_id=1)
+    runner.device = torch.device("cpu")
+    runner._local_rank = 0
+    runner.pipeline = None
+    runner.kv_transfer_manager = _FakeKVTransferManager(connector, recv_stages=recv_stages)
+    runner.init_omni_connectors(runner.od_config, runner.kv_transfer_manager, synchronous=True)
+    return runner
+
+
+def _make_sender(connector, *, payload_keys=("prompt_embeds",), send_stages=("0", "1")):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(stage_output_payload_keys=payload_keys, stage_id=0)
+    runner.device = torch.device("cpu")
+    runner.pipeline = None
+    runner.kv_transfer_manager = _FakeKVTransferManager(connector, send_stages=send_stages)
+    runner.init_omni_connectors(runner.od_config, runner.kv_transfer_manager, synchronous=True)
+    return runner
+
+
+def _make_output(**custom):
+    return SimpleNamespace(custom_output=dict(custom))
+
+
+def _make_request(prompt, *, request_id="req-7", kv_sender_info=None, payload_sender_info=None):
+    return SimpleNamespace(
+        prompt=prompt,
+        request_id=request_id,
+        kv_sender_info=kv_sender_info,
+        payload_sender_info=payload_sender_info,
+    )
+
+
+def _conditioning():
+    return {"text_encoder_output": {"hidden_states": torch.zeros(4, 8), "token_tags": torch.zeros(4)}}
+
+
+@pytest.mark.parametrize("shard_size,replicate_size", [(4, 1), (1, 4), (2, 2)])
+def test_standalone_hsdp_elects_one_payload_consumer(monkeypatch, shard_size, replicate_size):
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    leaders = []
+    for rank in range(shard_size * replicate_size):
+        shard_group = SimpleNamespace(world_size=shard_size, rank_in_group=rank % shard_size)
+        replicate_group = SimpleNamespace(world_size=replicate_size, rank_in_group=rank // shard_size)
+        monkeypatch.setattr(parallel_state, "get_fs_group", lambda: shard_group)
+        monkeypatch.setattr(parallel_state, "get_hsdp_replicate_group", lambda: replicate_group)
+        monkeypatch.setattr(parallel_state, "get_sp_group", lambda: None)
+        runner = _make_runner(_FakeConnector())
+        monkeypatch.setattr(runner, "_get_local_tp_group", lambda: None)
+        if runner.is_data_transfer_rank():
+            leaders.append(rank)
+
+    assert leaders == [0]
+
+
+@pytest.mark.parametrize("shard_size,replicate_size", [(4, 1), (1, 4), (2, 2)])
+@pytest.mark.parametrize("delivered", [True, False])
+def test_standalone_hsdp_payload_fanout(monkeypatch, shard_size, replicate_size, delivered):
+    from uuid import uuid4
+
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+
+    packets = {}
+
+    class MeshGroup:
+        def __init__(self, dimension, peer_rank, rank, size):
+            self.key = (dimension, peer_rank)
+            self.rank_in_group = rank
+            self.world_size = size
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "delivered"] = value
+            return packets[self.key, "delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "payload"] = value
+            return packets[self.key, "payload"]
+
+    connector = SharedMemoryConnector({})
+    request_id = f"hsdp-{uuid4().hex}"
+    key = f"{request_id}_0_0"
+    expected = torch.arange(8).reshape(2, 4)
+    calls = []
+    original_get = connector.get_with_deadline
+
+    def tracked_get(*args, **kwargs):
+        calls.append(rank)
+        return original_get(*args, **kwargs)
+
+    monkeypatch.setattr(connector, "get_with_deadline", tracked_get)
+    monkeypatch.setattr(parallel_state, "get_sp_group", lambda: None)
+    try:
+        if delivered:
+            assert connector.put("0", "1", key, {"encoder_output": expected})[0]
+        for rank in range(shard_size * replicate_size):
+            shard_rank = rank % shard_size
+            replicate_rank = rank // shard_size
+            shard_group = MeshGroup("shard", replicate_rank, shard_rank, shard_size)
+            replicate_group = MeshGroup("replicate", shard_rank, replicate_rank, replicate_size)
+            monkeypatch.setattr(parallel_state, "get_fs_group", lambda: shard_group)
+            monkeypatch.setattr(parallel_state, "get_hsdp_replicate_group", lambda: replicate_group)
+            runner = _make_runner(connector, payload_keys=("encoder_output",))
+            monkeypatch.setattr(runner, "_get_local_tp_group", lambda: None)
+            request = _make_request({"prompt": "test"}, request_id=request_id)
+            if delivered:
+                runner._maybe_recv_stage_payload(request)
+                torch.testing.assert_close(request.prompt["additional_information"]["encoder_output"], expected)
+            else:
+                with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+                    runner._maybe_recv_stage_payload(request)
+        assert set(calls) == {0}
+        if delivered:
+            assert calls == [0]
+            assert connector.get("0", "1", key) is None
+    finally:
+        connector.close()
+
+
+def test_key_convention_fetch_merges_into_additional_information():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == [("0", "1", "req-7_0_0", None)]
+    assert set(req.prompt["additional_information"]) == {"text_encoder_output"}
+    assert req.prompt["additional_information"]["text_encoder_output"]["hidden_states"].shape == (4, 8)
+
+
+def test_handle_path_uses_its_own_key_and_metadata():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, payload_keys=())
+    handle = {
+        "key": "custom-key",
+        "from_stage": "3",
+        "to_stage": "4",
+        "metadata": {"schema_version": 1},
+    }
+    req = _make_request({"prompt": "a cat", HANDLE_KEY: handle})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == [("3", "4", "custom-key", {"schema_version": 1})]
+    # The handle is transport bookkeeping and must not reach the pipeline.
+    assert HANDLE_KEY not in req.prompt
+    assert "text_encoder_output" in req.prompt["additional_information"]
+
+
+def test_receive_retries_until_producer_publishes(monkeypatch):
+    connector = _FakeConnector(_conditioning())
+    get = connector.get
+    attempts = []
+
+    def delayed_get(*args, **kwargs):
+        attempts.append(args)
+        return get(*args, **kwargs) if len(attempts) == 3 else None
+
+    monkeypatch.setattr(connector, "get", delayed_get)
+    runner = _make_runner(connector)
+    request = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(request)
+
+    assert len(attempts) == 3
+    assert "text_encoder_output" in request.prompt["additional_information"]
+
+
+def test_receive_retry_budget_preserves_inline_payload(monkeypatch):
+    from vllm_omni.distributed.omni_connectors.model_runner import omni_connector_payload_transport
+
+    clock = [0.0]
+
+    def advance(duration):
+        clock[0] += duration
+
+    monkeypatch.setattr(
+        omni_connector_payload_transport, "time", SimpleNamespace(monotonic=lambda: clock[0], sleep=advance)
+    )
+    connector = _FakeConnector(None)
+    runner = _make_runner(connector)
+    inline = _conditioning()
+    request = _make_request({"additional_information": inline})
+
+    runner._maybe_recv_stage_payload(request)
+
+    assert 1 < len(connector.calls) <= 41
+    assert clock[0] == pytest.approx(2.0)
+    assert request.prompt["additional_information"] is inline
+
+
+def test_stage_without_declared_keys_never_touches_the_connector():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, payload_keys=())
+    req = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == []
+    assert "additional_information" not in req.prompt
+
+
+def test_undeclared_payload_keys_are_dropped():
+    connector = _FakeConnector({**_conditioning(), "debug_blob": torch.zeros(2)})
+    runner = _make_runner(connector)
+    req = _make_request({"prompt": "a cat"})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert set(req.prompt["additional_information"]) == {"text_encoder_output"}
+
+
+@pytest.mark.parametrize(
+    "connector",
+    [
+        _FakeConnector(None),
+        _FakeConnector(raises=True),
+        _FakeConnector("not-a-dict"),
+    ],
+    ids=["missing", "raises", "wrong_type"],
+)
+def test_failed_transfer_falls_back_to_the_inline_prompt(connector):
+    runner = _make_runner(connector)
+    inline = {"text_encoder_output": {"hidden_states": torch.ones(2, 8)}}
+    req = _make_request({"prompt": "a cat", "additional_information": dict(inline)})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert torch.equal(
+        req.prompt["additional_information"]["text_encoder_output"]["hidden_states"],
+        inline["text_encoder_output"]["hidden_states"],
+    )
+
+
+def test_legacy_sender_info_is_passed_as_request_metadata():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request({"prompt": "a cat"}, kv_sender_info={0: {"host": "10.0.0.1", "zmq_port": 50171}})
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert runner.kv_transfer_manager.sender_info_calls == []
+    assert connector.calls[0][3] == {"source_host": "10.0.0.1", "source_port": 50171}
+
+
+def test_payload_sender_info_overrides_kv_sender_info_for_full_payloads():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request(
+        {"prompt": "a cat"},
+        kv_sender_info={0: {"host": "10.0.0.1", "zmq_port": 50171}},
+        payload_sender_info={"host": "10.0.0.1", "zmq_port": 50071},
+    )
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert runner.kv_transfer_manager.sender_info_calls == []
+    assert connector.calls[0][3] == {"source_host": "10.0.0.1", "source_port": 50071}
+
+
+def test_synchronous_receive_keeps_endpoints_request_scoped_and_uses_external_ids():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, recv_stages=("2", "5"))
+    for index in range(2):
+        request = _make_request(
+            {},
+            request_id=f"internal-{index}",
+            payload_sender_info={"host": f"10.0.0.{index + 1}", "zmq_port": 50071 + index},
+        )
+        request.external_req_id = f"external-{index}"
+        runner._maybe_recv_stage_payload(request)
+        assert "text_encoder_output" in request.prompt["additional_information"]
+    assert connector.calls == [
+        ("2", "5", "external-0_2_0", {"source_host": "10.0.0.1", "source_port": 50071}),
+        ("2", "5", "external-1_2_0", {"source_host": "10.0.0.2", "source_port": 50072}),
+    ]
+    assert runner.kv_transfer_manager.sender_info_calls == []
+    assert runner._pending_load_reqs == {}
+    assert runner._get_req_chunk == {}
+
+
+@pytest.mark.parametrize("stage_key", [2, "2"])
+@pytest.mark.parametrize("use_handle", [False, True])
+def test_sender_map_selects_actual_source_stage(stage_key, use_handle):
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, recv_stages=("0" if use_handle else "2", "5"))
+    prompt = {}
+    if use_handle:
+        prompt[HANDLE_KEY] = {"key": "explicit", "from_stage": "2", "to_stage": "5"}
+    request = _make_request(
+        prompt,
+        payload_sender_info={
+            "0": {"host": "wrong", "zmq_port": 50000},
+            stage_key: {"host": "right", "zmq_port": 50002},
+        },
+    )
+    runner._maybe_recv_stage_payload(request)
+    assert connector.calls[0][0:2] == ("2", "5")
+    assert connector.calls[0][3] == {"source_host": "right", "source_port": 50002}
+
+
+def test_tp_payload_is_fetched_once_by_leader_and_broadcast_to_followers():
+    connector = _FakeConnector(_conditioning())
+    state = {}
+
+    class _FakeTPGroup:
+        world_size = 4
+
+        def __init__(self, rank):
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                state["delivered"] = value
+            return state["delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                state["payload"] = value
+            return state["payload"]
+
+    leader = _make_runner(connector)
+    leader._local_rank = 0
+    leader._get_local_tp_group = lambda: _FakeTPGroup(0)
+    follower = _make_runner(connector)
+    follower._local_rank = 1
+    follower._get_local_tp_group = lambda: _FakeTPGroup(1)
+
+    leader_req = _make_request({"prompt": "a cat"})
+    follower_req = _make_request({"prompt": "a cat"})
+    leader._maybe_recv_stage_payload(leader_req)
+    follower._maybe_recv_stage_payload(follower_req)
+
+    assert len(connector.calls) == 1
+    leader_output = leader_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    follower_output = follower_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    assert torch.equal(leader_output, follower_output)
+
+
+def test_tp_payload_miss_is_broadcast_without_follower_connector_access():
+    connector = _FakeConnector(None)
+    state = {}
+
+    class _FakeTPGroup:
+        world_size = 2
+
+        def __init__(self, rank):
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                state["delivered"] = value
+            return state["delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            raise AssertionError("missing payload must not be broadcast")
+
+    leader = _make_runner(connector)
+    leader._get_local_tp_group = lambda: _FakeTPGroup(0)
+    follower = _make_runner(connector)
+    follower._local_rank = 1
+    follower._get_local_tp_group = lambda: _FakeTPGroup(1)
+
+    with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+        leader._maybe_recv_stage_payload(_make_request({"prompt": "a cat"}))
+    with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+        follower._maybe_recv_stage_payload(_make_request({"prompt": "a cat"}))
+
+    assert 1 < len(connector.calls) <= 41
+
+
+def test_sp_payload_is_fetched_once_by_leader_and_broadcast_to_followers():
+    connector = _FakeConnector(_conditioning())
+    state = {}
+
+    class _FakeSPGroup:
+        world_size = 4
+
+        def __init__(self, rank):
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                state["delivered"] = value
+            return state["delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                state["payload"] = value
+            return state["payload"]
+
+    leader = _make_runner(connector)
+    leader._get_local_tp_group = lambda: None
+    leader._stage_payload_broadcast_groups = lambda: (_FakeSPGroup(0),)
+    follower = _make_runner(connector)
+    follower._local_rank = 1
+    follower._get_local_tp_group = lambda: None
+    follower._stage_payload_broadcast_groups = lambda: (_FakeSPGroup(1),)
+
+    leader_req = _make_request({"prompt": "a cat"})
+    follower_req = _make_request({"prompt": "a cat"})
+    leader._maybe_recv_stage_payload(leader_req)
+    follower._maybe_recv_stage_payload(follower_req)
+
+    assert len(connector.calls) == 1
+    leader_output = leader_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    follower_output = follower_req.prompt["additional_information"]["text_encoder_output"]["hidden_states"]
+    assert torch.equal(leader_output, follower_output)
+
+
+def test_missing_incoming_edge_is_reported_not_fetched():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector, recv_stages=(None, None))
+    req = _make_request({"prompt": "a cat"})
+
+    with pytest.raises(RuntimeError, match="no incoming edge"):
+        runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == []
+    assert "additional_information" not in req.prompt
+
+
+@pytest.mark.parametrize("delivered", [True, False])
+def test_payload_reaches_full_tp_sp_grid_once(monkeypatch, delivered):
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    connector = _FakeConnector(_conditioning() if delivered else None)
+    packets = {}
+
+    class GridGroup:
+        world_size = 2
+
+        def __init__(self, dimension, peer_rank, rank):
+            self.key = (dimension, peer_rank)
+            self.rank_in_group = rank
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "delivered"] = value
+            return packets[self.key, "delivered"]
+
+        def broadcast_tensor_dict(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key, "payload"] = value
+            return packets[self.key, "payload"]
+
+    for sp_rank in range(2):
+        for tp_rank in range(2):
+            runner = _make_runner(connector)
+            tp_group = GridGroup("tp", sp_rank, tp_rank)
+            sp_group = GridGroup("sp", tp_rank, sp_rank)
+            monkeypatch.setattr(runner, "_get_local_tp_group", lambda: tp_group)
+            monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
+            request = _make_request({"prompt": "a cat"})
+
+            if delivered:
+                runner._maybe_recv_stage_payload(request)
+                output = request.prompt["additional_information"]["text_encoder_output"]
+                assert torch.equal(output["hidden_states"], connector._payload["text_encoder_output"]["hidden_states"])
+            else:
+                with pytest.raises(RuntimeError, match="Stage payload unavailable"):
+                    runner._maybe_recv_stage_payload(request)
+    if delivered:
+        assert len(connector.calls) == 1
+    else:
+        assert 1 < len(connector.calls) <= 41
+
+
+def test_native_kv_runner_keeps_synchronous_payload_transport(monkeypatch):
+    from unittest.mock import Mock
+
+    from vllm_omni.diffusion.worker import diffusion_model_runner as runner_module
+
+    connector = _FakeConnector(_conditioning())
+    manager = _FakeKVTransferManager(connector)
+    config = SimpleNamespace(stage_input_payload_keys=("text_encoder_output",), stage_id=1, kv_transfer_config=object())
+    monkeypatch.setattr(runner_module.OmniKVTransferManager, "from_od_config", lambda config: manager)
+    monkeypatch.setattr(runner_module, "DiffusionKVModelRunnerBackend", Mock())
+    runner = DiffusionModelRunner(SimpleNamespace(), config, torch.device("cpu"))
+    runner._initialize_generator = Mock()
+    request = _make_request({"prompt": "a cat"})
+    request.sampling_params = SimpleNamespace()
+
+    runner._prepare_request_for_forward(request, od_config=config)
+
+    assert runner.kv_transfer_manager is None
+    assert runner._kv_transfer_manager is manager
+    assert runner._kv_connector is None
+    assert not runner._kv_prefetch_enabled
+    assert len(connector.calls) == 1
+    assert "text_encoder_output" in request.prompt["additional_information"]
+    runner._initialize_generator.assert_called_once_with(request.sampling_params)
+
+
+def test_step_mode_receives_payload_before_kv_and_reuses_cached_state():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    runner.state_cache = {}
+    request = _make_request({"prompt": "a cat"})
+    request.sampling_params = SimpleNamespace()
+    kv_calls = []
+
+    def receive_kv(state_request, **kwargs):
+        assert "text_encoder_output" in state_request.prompt["additional_information"]
+        kv_calls.append(state_request.request_id)
+
+    runner.kv_transfer_manager.receive_multi_kv_cache_distributed = receive_kv
+    scheduled = SimpleNamespace(
+        scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request, diffusion_kv_metadata=None)],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+    states, new_ids = runner._update_states(scheduled)
+    assert new_ids == [request.request_id]
+    assert states[0].prompt is request.prompt
+    assert len(connector.calls) == 1
+    assert kv_calls == [request.request_id]
+
+    scheduled.scheduled_new_reqs = []
+    scheduled.scheduled_cached_reqs.request_ids = [request.request_id]
+    cached_states, new_ids = runner._update_states(scheduled)
+    assert cached_states == states
+    assert new_ids == []
+    assert len(connector.calls) == 1
+    assert kv_calls == [request.request_id]
+
+
+def test_non_dict_prompt_is_left_alone():
+    connector = _FakeConnector(_conditioning())
+    runner = _make_runner(connector)
+    req = _make_request("a cat")
+
+    runner._maybe_recv_stage_payload(req)
+
+    assert connector.calls == []
+    assert req.prompt == "a cat"
+
+
+@pytest.mark.parametrize("request_id", ["dummy_req_id", "dummy_req_id/profile"])
+def test_warmup_never_receives_or_publishes_payload(request_id):
+    connector = _FakeConnector(raises=True)
+    request = _make_request({}, request_id=request_id)
+    receiver = _make_runner(connector)
+    receiver._maybe_recv_stage_payload(request)
+    sender = _make_sender(connector)
+    sender._maybe_send_stage_payload([request], [_make_output(prompt_embeds=torch.ones(1))])
+    assert connector.calls == []
+    assert connector.put_calls == []
+
+
+@pytest.mark.parametrize("inline", [{}, {"text_encoder_output": None}])
+def test_failed_transfer_requires_complete_inline_payload(inline):
+    runner = _make_runner(_FakeConnector(raises=True))
+    with pytest.raises(RuntimeError, match="missing keys.*text_encoder_output"):
+        runner._maybe_recv_stage_payload(_make_request({"additional_information": inline}))
+
+
+def test_send_puts_declared_keys_and_attaches_a_handle():
+    connector = _FakeConnector()
+    runner = _make_sender(connector)
+    req = _make_request({"prompt": "a cat"})
+    output = _make_output(prompt_embeds=torch.zeros(2, 8), latents=torch.zeros(1, 4))
+
+    runner._maybe_send_stage_payload([req], [output])
+
+    from_stage, to_stage, put_key, data = connector.put_calls[0]
+    assert (from_stage, to_stage, put_key) == ("0", "1", "req-7_0_0")
+    # Only declared keys travel over the connector.
+    assert isinstance(data, dict)
+    assert set(data) == {"prompt_embeds"}
+    handle = output.custom_output[HANDLE_KEY]
+    assert handle["key"] == "req-7_0_0"
+    assert handle["metadata"] == {"schema_version": 1}
+    assert handle["payload_keys"] == ["prompt_embeds"]
+    assert handle["size_bytes"] == 128
+    # Transferred keys must not also travel inline through the orchestrator.
+    assert "prompt_embeds" not in output.custom_output
+    # Undeclared keys are untouched.
+    assert "latents" in output.custom_output
+
+
+def test_send_key_matches_the_receive_key_convention():
+    sender_connector = _FakeConnector()
+    sender = _make_sender(sender_connector)
+    req = _make_request({"prompt": "a cat"})
+    sender._maybe_send_stage_payload([req], [_make_output(prompt_embeds=torch.zeros(2, 8))])
+
+    recv_connector = _FakeConnector(_conditioning())
+    receiver = _make_runner(recv_connector)
+    receiver._maybe_recv_stage_payload(_make_request({"prompt": "a cat"}))
+
+    assert sender_connector.put_calls[0][2] == recv_connector.calls[0][2]
+
+
+def test_stage_without_declared_output_keys_never_puts():
+    connector = _FakeConnector()
+    runner = _make_sender(connector, payload_keys=())
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert connector.put_calls == []
+    assert HANDLE_KEY not in output.custom_output
+
+
+def test_missing_outgoing_edge_is_reported_not_sent():
+    connector = _FakeConnector()
+    runner = _make_sender(connector, send_stages=(None, None))
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert connector.put_calls == []
+    assert HANDLE_KEY not in output.custom_output
+
+
+def test_output_without_the_declared_keys_is_skipped():
+    connector = _FakeConnector()
+    runner = _make_sender(connector)
+    output = _make_output(latents=torch.zeros(1, 4))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert connector.put_calls == []
+    assert HANDLE_KEY not in output.custom_output
+
+
+@pytest.mark.parametrize(
+    "connector",
+    [
+        _FakeConnector(raises=True),
+        _FakeConnector(put_result=(False, 0, None)),
+    ],
+    ids=["raises", "rejected"],
+)
+def test_failed_send_leaves_the_inline_payload_unannotated(connector):
+    runner = _make_sender(connector)
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert HANDLE_KEY not in output.custom_output
+    assert "prompt_embeds" in output.custom_output
+
+
+class _FakeTPGroup:
+    def __init__(self, rank_in_group, world_size=2):
+        self.rank_in_group = rank_in_group
+        self.world_size = world_size
+        self.broadcast_calls: list[object] = []
+        self.leader_packet: object = None
+
+    def broadcast_object(self, obj, src=0):
+        self.broadcast_calls.append(obj)
+        return obj if self.rank_in_group == src else self.leader_packet
+
+
+@pytest.mark.parametrize("rank_in_group", [0, 1], ids=["leader", "follower"])
+def test_send_puts_only_on_the_tp_leader(monkeypatch, rank_in_group):
+    connector = _FakeConnector()
+    runner = _make_sender(connector)
+    tp_group = _FakeTPGroup(rank_in_group)
+    expected_handle = {
+        "key": "req-7_0_0",
+        "from_stage": "0",
+        "to_stage": "1",
+        "size_bytes": 128,
+        "metadata": {"schema_version": 1},
+        "payload_keys": ["prompt_embeds"],
+    }
+    tp_group.leader_packet = {"req-7": expected_handle}
+    monkeypatch.setattr(DiffusionModelRunner, "_get_local_tp_group", staticmethod(lambda: tp_group))
+    output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+    runner._maybe_send_stage_payload([_make_request({"prompt": "a cat"})], [output])
+
+    assert len(connector.put_calls) == (1 if rank_in_group == 0 else 0)
+    # Every rank ends up reporting the same handle, leader or not.
+    assert output.custom_output[HANDLE_KEY] == expected_handle
+    # ...and every rank drops the inline copy, not just the sender.
+    assert "prompt_embeds" not in output.custom_output
+    assert len(tp_group.broadcast_calls) == 1
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+@pytest.mark.parametrize("accepted", [True, False])
+def test_send_elects_one_tp_sp_leader(monkeypatch, tp_size, accepted):
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    connector = _FakeConnector(put_result=(accepted, 128, {"schema_version": 1}))
+    packets = {}
+    outputs = []
+
+    class GridGroup:
+        def __init__(self, dimension, peer_rank, rank, world_size):
+            self.key = (dimension, peer_rank)
+            self.rank_in_group = rank
+            self.world_size = world_size
+
+        def broadcast_object(self, value, src=0):
+            if self.rank_in_group == src:
+                packets[self.key] = value
+            return packets[self.key]
+
+    for sp_rank in range(2):
+        for tp_rank in range(tp_size):
+            runner = _make_sender(connector)
+            tp_group = GridGroup("tp", sp_rank, tp_rank, tp_size)
+            sp_group = GridGroup("sp", tp_rank, sp_rank, 2)
+            monkeypatch.setattr(runner, "_get_local_tp_group", lambda: tp_group)
+            monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
+            output = _make_output(prompt_embeds=torch.zeros(2, 8))
+
+            runner._maybe_send_stage_payload([_make_request({})], [output])
+
+            outputs.append(output.custom_output)
+            assert (HANDLE_KEY in output.custom_output) == accepted
+            assert ("prompt_embeds" in output.custom_output) != accepted
+
+    assert len(connector.put_calls) == 1
+    if accepted:
+        assert all(output[HANDLE_KEY] == outputs[0][HANDLE_KEY] for output in outputs)

tests/diffusion/test_diffusion_step_pipeline.py

@@ -616,7 +616,8 @@ def test_make_batch_explicit_mapping_still_selects_and_orders_states():
 
 
 @pytest.mark.cpu
-def test_step_profiler_reports_denoise_step_as_diffuse():
+def test_step_profiler_reports_denoise_step_as_diffuse(monkeypatch):
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
     pipeline = _AutoDenoiseProfilerPipeline()
 
     assert pipeline.denoise_step() == "ok"
@@ -629,6 +630,13 @@ def test_step_profiler_reports_denoise_step_as_diffuse():
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
 
+    @pytest.fixture(autouse=True)
+    def mock_platform_memory(self, monkeypatch):
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "is_available", lambda: True)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+
     @pytest.mark.parametrize(
         ("diffusion_kv_mode", "should_cleanup"),
         [

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -87,6 +87,38 @@ async def test_inline_dispatch_request_success(client, mock_engine):
     assert output.request_id == "req-1"
 
 
+@pytest.mark.parametrize("streaming", [False, True])
+def test_inline_dispatch_preserves_payload_sender_info(client, mock_engine, streaming):
+    async def _run_test():
+        requests = []
+
+        async def _step_streaming(request):
+            requests.append(request)
+            yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[MagicMock()])]
+
+        mock_engine.step_streaming = _step_streaming
+        client.od_config.streaming_output = streaming
+        payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50071}
+        kv_transfer_params = {"remote_engine_id": "mooncake-producer", "remote_block_ids": [1, 2]}
+        await client.add_request_async(
+            "req-payload",
+            "A test prompt",
+            OmniDiffusionSamplingParams(),
+            kv_transfer_params=kv_transfer_params,
+            payload_sender_info=payload_sender_info,
+        )
+
+        for _ in range(10):
+            if requests:
+                break
+            await asyncio.sleep(0.01)
+
+        assert requests[0].payload_sender_info == payload_sender_info
+        assert requests[0].kv_transfer_params == kv_transfer_params
+
+    asyncio.run(_run_test())
+
+
 @pytest.mark.asyncio
 async def test_inline_non_streaming_dispatches_lifecycle_before_final(client, mock_engine):
     lifecycle = OmniRequestOutput.from_diffusion(

tests/diffusion/test_stage_payload_kv_coexistence.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_subprocess_client_and_worker_preserve_both_transfer_contracts(streaming):
+    """Exercise the IPC message and both process request builders without GPUs."""
+
+    async def run():
+        client = object.__new__(StageDiffusionClient)
+        client._engine_dead = False
+        client.stage_id = 1
+        client.replica_id = 0
+        client._encoder = MagicMock()
+        client._encoder.encode.side_effect = lambda message: message
+        client._request_socket = MagicMock()
+        kv_params = {"remote_engine_id": "mooncake-producer", "remote_block_ids": [1, 2]}
+        payload_sender = {"host": "10.0.0.1", "zmq_port": 50071}
+        kv_sender = {0: {"host": "10.0.0.2", "zmq_port": 50171}}
+        # Keep main's existing fifth positional argument bound to native KV.
+        await client.add_request_async(
+            "req-both",
+            "prompt",
+            OmniDiffusionSamplingParams(),
+            kv_sender,
+            kv_params,
+            payload_sender_info=payload_sender,
+        )
+        message = client._request_socket.send.call_args.args[0]
+        assert message["kv_transfer_params"] == kv_params
+        assert message["payload_sender_info"] == payload_sender
+        captured = []
+
+        class Engine:
+            async def step_streaming(self, request):
+                captured.append(request)
+                yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[])]
+
+        proc = object.__new__(StageDiffusionProc)
+        proc._engine = Engine()
+        args = (message["request_id"], message["prompt"], message["sampling_params"])
+        transfer_args = {
+            "kv_sender_info": kv_sender,
+            "kv_transfer_params": kv_params,
+            "payload_sender_info": message["payload_sender_info"],
+        }
+        if streaming:
+            outputs = [output async for output in proc._process_streaming_request(*args, **transfer_args)]
+        else:
+            outputs = [await proc._process_request(*args, **transfer_args)]
+        assert outputs[0].request_id == "req-both"
+        assert captured[0].kv_transfer_params == kv_params
+        assert captured[0].kv_sender_info == kv_sender
+        assert captured[0].payload_sender_info == payload_sender
+
+    asyncio.run(run())

tests/diffusion/test_worker_wrapper_base.py

@@ -12,10 +12,13 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+import gc
+import weakref
 from types import SimpleNamespace
 from typing import Any
 
 import pytest
+import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.worker.diffusion_worker import (
@@ -227,12 +230,13 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
-    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
+    @pytest.mark.parametrize("native_kv", [False, True])
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture, native_kv):
         events: list[str] = []
         offload_backend = mocker.Mock()
         offload_backend.disable.side_effect = lambda: events.append("offload")
         kv_manager = mocker.Mock()
-        kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
+        kv_manager.close.side_effect = lambda: events.append("kv")
         destroy = mocker.patch(
             "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
             side_effect=lambda: events.append("distributed"),
@@ -240,14 +244,71 @@ class TestWorkerWrapperBaseDelegation:
         worker = DiffusionWorker.__new__(DiffusionWorker)
         worker.model_runner = SimpleNamespace(
             offload_backend=offload_backend,
-            kv_transfer_manager=kv_manager,
+            kv_transfer_manager=None if native_kv else kv_manager,
+            _kv_transfer_manager=kv_manager,
         )
 
         worker.shutdown()
 
         assert events == ["offload", "kv", "distributed"]
+        kv_manager.close.assert_called_once_with()
         destroy.assert_called_once_with()
 
+    @pytest.mark.parametrize("cleanup_fails", [False, True])
+    def test_worker_shutdown_releases_model_references(self, mocker: MockerFixture, cleanup_fails):
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env")
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector")
+        offload_backend = mocker.Mock()
+        if cleanup_fails:
+            offload_backend.disable.side_effect = RuntimeError("offload cleanup failed")
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        weights = torch.ones(1)
+        weights_ref = weakref.ref(weights)
+        pipeline = SimpleNamespace(weights=weights)
+        worker.model_runner = SimpleNamespace(
+            pipeline=pipeline, offload_backend=offload_backend, kv_transfer_manager=None
+        )
+        worker.lora_manager = SimpleNamespace(pipeline=pipeline)
+        worker._sleep_saved_buffers = {"buffer": weights}
+        del weights, pipeline
+
+        if cleanup_fails:
+            with pytest.raises(RuntimeError, match="offload cleanup failed"):
+                worker.shutdown()
+        else:
+            worker.shutdown()
+        gc.collect()
+
+        assert worker.model_runner is None
+        assert worker.lora_manager is None
+        assert worker._sleep_saved_buffers == {}
+        assert weights_ref() is None
+
+    @pytest.mark.parametrize("owns_pool", [False, True])
+    @pytest.mark.parametrize("cleanup_fails", [False, True])
+    def test_worker_shutdown_releases_only_owned_sleep_pool(self, mocker: MockerFixture, owns_pool, cleanup_fails):
+        destroy = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env")
+        shutdown_connector = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector")
+        allocator_class = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class")
+        release = allocator_class.return_value.get_instance.return_value.release_pools
+        if cleanup_fails:
+            release.side_effect = RuntimeError("pool release failed")
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        worker.model_runner = None
+        worker._owns_sleep_pool = owns_pool
+
+        if owns_pool and cleanup_fails:
+            with pytest.raises(RuntimeError, match="pool release failed"):
+                worker.shutdown()
+        else:
+            worker.shutdown()
+
+        assert release.call_count == int(owns_pool)
+        shutdown_connector.assert_called_once_with()
+        destroy.assert_called_once_with()
+        if not cleanup_fails:
+            assert not worker._owns_sleep_pool
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method
@@ -507,6 +568,26 @@ class TestCustomPipelineWorkerExtension:
         # Verify re_init_pipeline was called with custom_pipeline_args
         mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)
 
+    @pytest.mark.parametrize("cleanup_fails", [False, True])
+    def test_custom_pipeline_initialization_failure_cleans_worker(
+        self, mocker: MockerFixture, mock_od_config, cleanup_fails
+    ):
+        worker = mocker.Mock()
+        worker.re_init_pipeline.side_effect = RuntimeError("custom pipeline load failed")
+        if cleanup_fails:
+            worker.shutdown.side_effect = RuntimeError("cleanup failed")
+        worker_class = mocker.Mock(return_value=worker)
+        mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class", return_value=worker_class)
+
+        with pytest.raises(RuntimeError, match="custom pipeline load failed"):
+            WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                custom_pipeline_args={"pipeline_class": "unused"},
+            )
+
+        worker.shutdown.assert_called_once_with()
+
     def test_custom_pipeline_with_explicit_extension(self, mocker: MockerFixture, mock_od_config):
         """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided."""
 

tests/distributed/omni_connectors/test_kv_async_prefetch.py

@@ -205,6 +205,25 @@ def test_start_prefetch_sweeps_orphan():
     assert "rid-next" in receiver._prefetch_futures
 
 
+def test_manager_close_stops_prefetch_before_connector_and_is_idempotent():
+    receiver = object.__new__(OmniKVTransferManager)
+    events = []
+    receiver.shutdown_prefetch = lambda: events.append("prefetch")
+    receiver._connector = SimpleNamespace(close=lambda: events.append("connector"))
+    receiver.close()
+    receiver.close()
+    assert events == ["prefetch", "connector", "prefetch"]
+    assert receiver.connector is None
+
+
+def test_manager_close_does_not_initialize_unused_connector():
+    receiver = object.__new__(OmniKVTransferManager)
+    receiver.shutdown_prefetch = lambda: None
+    receiver._connector = None
+    receiver.close()
+    assert receiver.connector is None
+
+
 def test_shutdown_prefetch_clears_state():
     sender, receiver, _ = _make_sender_receiver()
     _seed_payload(sender, "rid-sd")

tests/distributed/omni_connectors/test_nixl_connector.py

@@ -132,6 +132,7 @@ def test_manager_receiver_does_not_bind_shared_producer_port(producer, need_recv
                 "role": "receiver",
                 "host": "127.0.0.1",
                 "zmq_port": PORT,
+                "metadata_query_timeout_ms": 5000,
                 "backends": ["UCX"],
             },
             from_stage="0",
@@ -280,14 +281,15 @@ def test_abandoned_claim_survives_close_and_late_completion(producer, consumer,
     consumer._notify_transfer_done("abandoned", claimed)
     assert not producer._pending
     assert not producer._agent.registered
-    assert producer._listener_thread.is_alive()
-    producer.close()
+    producer._close_thread.join(timeout=2.0)
+    assert not producer._close_thread.is_alive()
+    assert producer._listener_thread is None
     assert not producer._agent.registered
     assert producer._closed
 
 
 @pytest.fixture
-def copying_native_agent(consumer, monkeypatch):
+def ownership_copying_agent(consumer, monkeypatch):
     """Strict fake descriptors and actual CPU copies; not native NIXL evidence."""
     calls = []
 
@@ -320,7 +322,7 @@ def copying_native_agent(consumer, monkeypatch):
 
 @pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"])
 @pytest.mark.usefixtures("reliable_claim_queries")
-def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consumer, copying_native_agent, case):
+def test_zero_byte_leaves_complete_source_claim(producer, consumer, ownership_copying_agent, case):
     empty = torch.empty((2, 0, 3), dtype=torch.float64)
     other = torch.empty((0,), dtype=torch.int64)
     scalar = torch.tensor(7)
@@ -345,7 +347,7 @@ def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consume
         torch.testing.assert_close(actual, payload, rtol=0, atol=0)
     assert size == received_size
     if case in ("empty", "all_empty"):
-        assert not copying_native_agent
+        assert not ownership_copying_agent
     assert not producer._pending and not producer._agent.registered
     assert not consumer._agent.registered
 
@@ -354,7 +356,7 @@ def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consume
 @pytest.mark.parametrize("outcome", ["done", "error", "timeout", "unknown"])
 @pytest.mark.usefixtures("reliable_claim_queries")
 def test_read_ownership_through_terminal_and_deferred_paths(
-    producer, consumer, copying_native_agent, monkeypatch, direct, outcome
+    producer, consumer, ownership_copying_agent, monkeypatch, direct, outcome
 ):
     import threading
 
@@ -411,6 +413,7 @@ class _FakeNixlAgent:
         self.registered = []
 
     def get_reg_descs(self, regions, memory_type):
+        assert regions and all(region[0] > 0 and region[1] > 0 for region in regions)
         return ("reg", tuple(regions), memory_type)
 
     def register_memory(self, descs, backends=None):
@@ -548,21 +551,150 @@ def test_source_endpoint_metadata_overrides_configured_sender(consumer, monkeypa
     assert queried == [("req-tp4-rank3", "10.0.0.3", PORT + 3 * 16)]
 
 
-def test_unknown_key_is_queried_once(producer, consumer, monkeypatch):
+@pytest.mark.parametrize("force_timeout", [False, True], ids=["real-query", "recv-timeout"])
+def test_unknown_key_is_queried_once(producer, consumer, monkeypatch, force_timeout):
     socket = consumer._get_req_socket(f"tcp://127.0.0.1:{PORT}")
     send_count = 0
     original_send = socket.send
 
     def count_send(message):
         nonlocal send_count
+        assert socket.getsockopt(zmq.RCVTIMEO) == 10
+        assert socket.getsockopt(zmq.SNDTIMEO) == 10
         send_count += 1
         return original_send(message)
 
     monkeypatch.setattr(socket, "send", count_send)
+    if force_timeout:
+
+        def timeout_recv():
+            raise zmq.Again()
+
+        monkeypatch.setattr(socket, "recv", timeout_recv)
 
     assert consumer._resolve_metadata("never-published", None) is None
     assert send_count == 1
-    assert socket.getsockopt(zmq.RCVTIMEO) == 10
+    if force_timeout:
+        assert socket.closed
+
+
+@pytest.mark.usefixtures("reliable_claim_queries")
+def test_deadline_get_retains_active_dma_then_releases(producer, consumer, ownership_copying_agent, monkeypatch):
+    producer.put("0", "1", "deadline-read", torch.ones(2))
+    consumer._stop_event.set()
+    consumer._transfer_wakeup.set()
+    consumer._transfer_thread.join(timeout=2)
+    state = ["PROC"]
+    monkeypatch.setattr(consumer._agent, "check_xfer_state", lambda handle: state[0])
+    started = time.monotonic()
+    assert consumer.get_with_deadline("0", "1", "deadline-read", deadline=started + 0.03) is None
+    assert time.monotonic() - started < 0.5
+    assert producer._pending["deadline-read"].claims
+    assert producer._agent.registered
+    assert consumer._agent.registered
+    assert consumer._deferred_transfers
+    state[0] = "DONE"
+    consumer._reap_deferred_transfers()
+    assert not producer._pending
+    assert not producer._agent.registered
+    assert not consumer._agent.registered
+    assert not consumer._deferred_transfers
+
+
+def test_receive_deadline_overrides_dma_timeout(consumer, monkeypatch):
+    monkeypatch.setattr(consumer._agent, "check_xfer_state", lambda handle: "PROC", raising=False)
+    consumer._transfer_timeout_s = 300
+    consumer._req_local.deadline = time.monotonic() + 0.02
+    started = time.monotonic()
+    try:
+        with pytest.raises(TimeoutError):
+            consumer._wait_for_transfer(1, "bounded")
+        assert time.monotonic() - started < 0.5
+    finally:
+        consumer._req_local.deadline = None
+
+
+@pytest.mark.parametrize("closing", [False, True])
+def test_abandoned_metadata_query_recovers_same_claim_without_read(producer, consumer, monkeypatch, closing):
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _GET_META_MSG
+
+    consumer._stop_event.set()
+    consumer._transfer_wakeup.set()
+    consumer._transfer_thread.join(timeout=2)
+    producer.put("0", "1", "abandoned-query", torch.ones(1))
+    messages = []
+
+    def send(message):
+        messages.append(message)
+
+    def receive():
+        reply = producer._handle_handshake_message(messages[-1])
+        if len(messages) == 1:
+            raise zmq.Again()
+        return reply
+
+    socket = types.SimpleNamespace(send=send, recv=receive)
+    monkeypatch.setattr(consumer, "_get_req_socket", lambda *args: socket)
+    monkeypatch.setattr(consumer, "_invalidate_req_socket", lambda *args: None)
+    assert consumer.get_with_deadline("0", "1", "abandoned-query", deadline=time.monotonic() + 1) is None
+    assert not consumer._abandoned_queries
+    consumer.abandon_get("abandoned-query")
+    claim = next(iter(producer._pending["abandoned-query"].claims))
+    if closing:
+        producer._closing = True
+    consumer._reap_deferred_transfers()
+    assert not consumer._abandoned_queries
+    assert not producer._pending
+    claims = [
+        msgspec.msgpack.decode(message[len(_GET_META_MSG) :])["claim_id"]
+        for message in messages
+        if message.startswith(_GET_META_MSG)
+    ]
+    assert claims == [claim, claim]
+
+
+def test_expired_ack_is_deferred_instead_of_raising(consumer):
+    consumer._req_local.deadline = time.monotonic() - 1
+    try:
+        assert consumer._notify_transfer_done("expired", {"source_host": "127.0.0.1", "source_port": PORT}) is False
+    finally:
+        consumer._req_local.deadline = None
+
+
+@pytest.mark.parametrize("direct", [False, True])
+@pytest.mark.parametrize("lost_replies", [1, 3])
+def test_lost_metadata_reply_retry_releases_source(producer, consumer, monkeypatch, direct, lost_replies):
+    _, _, metadata = producer.put("0", "1", "lost-reply", torch.ones(1))
+    replies = []
+    requests = []
+
+    def send(message):
+        requests.append(message)
+        replies.append(producer._handle_handshake_message(message))
+
+    def recv():
+        reply = replies.pop(0)
+        if len(requests) <= lost_replies:
+            raise zmq.Again()
+        return reply
+
+    socket = types.SimpleNamespace(send=send, recv=recv)
+    monkeypatch.setattr(consumer, "_get_req_socket", lambda *args: socket)
+    forwarded = metadata if direct else None
+    for attempt in range(lost_replies):
+        assert consumer._resolve_metadata("lost-reply", forwarded) is None
+    resolved = consumer._resolve_metadata("lost-reply", forwarded)
+    assert resolved is not None
+    assert producer._pending["lost-reply"].claims == {resolved["claim_id"]}
+    independent = consumer._resolve_metadata("lost-reply", forwarded)
+    assert independent["claim_id"] != resolved["claim_id"]
+    consumer._notify_transfer_done("lost-reply", resolved)
+    consumer._notify_transfer_done("lost-reply", resolved)
+    assert producer._pending["lost-reply"].claims == {independent["claim_id"]}
+    consumer._notify_transfer_done("lost-reply", independent)
+    assert producer._pending == {}
+    assert producer._registered_descs == []
+    assert consumer._req_local.claims == {}
 
 
 @pytest.mark.usefixtures("reliable_claim_queries")
@@ -681,6 +813,77 @@ def test_deferred_transfer_releases_exactly_once_after_done(nixl_connector_cls):
         connector.close()
 
 
+@pytest.mark.parametrize("poll_raises", [False, True], ids=["stuck", "poll-error"])
+def test_close_returns_without_releasing_active_dma(nixl_connector_cls, poll_raises):
+    import threading
+
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import (
+        _RETAINED_PRODUCERS,
+        _DeferredTransfer,
+    )
+
+    connector = nixl_connector_cls({"role": "receiver"})
+    released = []
+
+    def check_state(handle):
+        if poll_raises:
+            raise RuntimeError("device unavailable")
+        return "PROC"
+
+    connector._agent.check_xfer_state = check_state
+    connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle))
+    connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle))
+    connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent))
+    connector._agent.deregister_memory = lambda descs: released.append(("registration", descs))
+    transfer = _DeferredTransfer(
+        tensors=[torch.zeros(1)],
+        registrations=["registration"],
+        dlists=["local", "remote"],
+        handles=["transfer"],
+        remote_agent="producer",
+    )
+    connector._defer_transfer(transfer)
+    closing = threading.Thread(target=connector.close, daemon=True)
+    try:
+        closing.start()
+        closing.join(timeout=2.0)
+        assert not closing.is_alive(), "close must not wait forever for DMA completion"
+        assert released == []
+        assert transfer in connector._deferred_transfers
+        assert transfer.tensors
+        assert connector._closing
+        assert connector in _RETAINED_PRODUCERS
+        assert connector.health()["status"] == "unhealthy"
+        cleanup_thread = connector._close_thread
+        assert cleanup_thread is not None
+        connector.close()
+        assert connector._close_thread is cleanup_thread
+        connector._agent.check_xfer_state = lambda handle: "DONE"
+        cleanup_thread.join(timeout=2.0)
+        assert not cleanup_thread.is_alive()
+        assert connector._closed
+        assert not connector._deferred_transfers
+        assert not transfer.tensors
+        assert connector not in _RETAINED_PRODUCERS
+    finally:
+        connector._agent.check_xfer_state = lambda handle: "DONE"
+        closing.join(timeout=2.0)
+        connector.close()
+
+    assert released == [
+        ("handle", "transfer"),
+        ("dlist", "local"),
+        ("dlist", "remote"),
+        ("agent", "producer"),
+        ("registration", "registration"),
+    ]
+    assert not connector._deferred_transfers
+    assert not transfer.tensors
+    assert connector not in _RETAINED_PRODUCERS
+    connector.close()
+    assert len(released) == 5
+
+
 def test_active_sibling_transfer_requires_deferred_ownership(nixl_connector_cls):
     connector = nixl_connector_cls({"role": "receiver"})
     connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle]
@@ -691,7 +894,8 @@ def test_active_sibling_transfer_requires_deferred_ownership(nixl_connector_cls)
         connector.close()
 
 
-def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_connector_cls):
+@pytest.mark.parametrize("with_empty", [False, True])
+def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_connector_cls, with_empty):
     connector = nixl_connector_cls({"role": "receiver"})
     connector._agent.add_remote_agent = lambda metadata: "producer"
     connector._agent.get_xfer_descs = lambda regions, memory_type: (regions, memory_type)
@@ -719,6 +923,9 @@ def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_conn
         ],
         "size": 8,
     }
+    if with_empty:
+        metadata["tensor_specs"].insert(1, {"shape": [2, 0, 3], "dtype": "torch.int64", "device": "cpu", "size": 0})
+        metadata["descriptor_groups"][1]["tensor_indices"] = [2]
 
     try:
         assert connector.get("0", "1", "req-active-sibling", metadata) is None
@@ -729,7 +936,9 @@ def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_conn
         assert len(deferred.registrations) == 2
         assert len(deferred.dlists) == 4
         assert deferred.remote_agent == "producer"
-        assert len(deferred.tensors) == 2
+        assert len(deferred.tensors) == 2 + with_empty
+        if with_empty:
+            assert deferred.tensors[1].shape == (2, 0, 3)
     finally:
         connector._agent.check_xfer_state = lambda handle: "DONE"
         connector.close()
@@ -893,3 +1102,198 @@ def test_receive_device_config_wins(nixl_connector_cls):
         assert connector._resolve_receive_device("cuda:3") == torch.device("cpu")
     finally:
         connector.close()
+
+
+@pytest.fixture
+def copying_native_agent(consumer, monkeypatch):
+    """Enforce native nonzero descriptors and copy only the submitted regions."""
+    agent = consumer._agent
+    calls = {"agents": [], "descs": [], "transfers": [], "released": []}
+
+    def add_remote(metadata):
+        calls["agents"].append(metadata)
+        return "producer"
+
+    def descriptors(regions, memory_type):
+        assert regions and all(len(region) == 3 and region[0] > 0 and region[1] > 0 for region in regions)
+        calls["descs"].append((regions, memory_type))
+        return regions
+
+    def prepare(operation, local, local_ids, remote, remote_ids):
+        assert operation == "READ"
+        assert local_ids == remote_ids == list(range(len(local)))
+        assert len(local) == len(remote)
+        pairs = list(zip(local, remote, strict=True))
+        for destination, source in pairs:
+            assert destination[1] == source[1] > 0
+        calls["transfers"].append(pairs)
+        return pairs
+
+    def transfer(pairs):
+        for destination, source in pairs:
+            ctypes.memmove(destination[0], source[0], source[1])
+
+    monkeypatch.setattr(agent, "add_remote_agent", add_remote, raising=False)
+    monkeypatch.setattr(agent, "get_xfer_descs", descriptors, raising=False)
+    monkeypatch.setattr(agent, "prep_xfer_dlist", lambda agent, descs: descs, raising=False)
+    monkeypatch.setattr(agent, "make_prepped_xfer", prepare, raising=False)
+    monkeypatch.setattr(agent, "transfer", transfer, raising=False)
+    monkeypatch.setattr(agent, "check_xfer_state", lambda handle: "DONE", raising=False)
+    for method in ("release_xfer_handle", "release_dlist_handle", "remove_remote_agent"):
+        monkeypatch.setattr(
+            agent, method, lambda handle, method=method: calls["released"].append((method, handle)), raising=False
+        )
+    return calls
+
+
+@pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"])
+@pytest.mark.usefixtures("reliable_claim_queries")
+def test_zero_byte_leaves_roundtrip_without_native_descriptors(
+    producer, consumer, copying_native_agent, monkeypatch, case
+):
+    empty = torch.empty((2, 0, 3), dtype=torch.float64)
+    other_empty = torch.empty((0,), dtype=torch.int64)
+    scalar = torch.tensor(7, dtype=torch.int64)
+    vector = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    payloads = {
+        "empty": empty,
+        "all_empty": [empty, other_empty],
+        "mixed": [empty, scalar, other_empty, vector, empty],
+        "scalar": scalar,
+        "structured_empty": {"leaves": (empty, [other_empty]), "label": "empty slots"},
+        "structured_mixed": {"leaves": (empty, [scalar, other_empty, vector, empty]), "label": "mixed slots"},
+    }
+    payload = payloads[case]
+    # Simulate mixed source memory groups without requiring a GPU in unit tests.
+    monkeypatch.setattr(producer, "_resolve_memory_type", lambda t: "VRAM" if t.dtype == torch.int64 else "DRAM")
+    # Splitting the producer's DRAM group on receipt also exercises local IDs:
+    # IDs are group-relative, whereas spec/skeleton indices remain global.
+    monkeypatch.setattr(consumer, "_resolve_memory_type", lambda t: "VRAM" if t.dtype == torch.float32 else "DRAM")
+    consumer._receive_device = torch.device("cpu")
+    success, size, metadata = producer.put("0", "1", "req-empty-slots", payload)
+    assert success
+    tensors = producer._pending["req-empty-slots"].tensors
+    specs = metadata["tensor_specs"]
+    assert len(specs) == len(tensors)
+    assert size == sum(t.numel() * t.element_size() for t in tensors)
+    for tensor, spec in zip(tensors, specs, strict=True):
+        assert spec["shape"] == list(tensor.shape)
+        assert spec["dtype"] == str(tensor.dtype)
+        assert spec["device"] == str(tensor.device)
+    indices = [i for group in metadata["descriptor_groups"] for i in group["tensor_indices"]]
+    assert sorted(indices) == [i for i, tensor in enumerate(tensors) if tensor.numel()]
+    # Verify receiver override applies to empty specs too, without trusting a
+    # producer's device index. Real CUDA placement is covered by the native suite.
+    for spec in specs:
+        if spec["size"] == 0:
+            spec["device"] = "cuda:7"
+    received = consumer.get("0", "1", "req-empty-slots", metadata)
+    assert received is not None
+    actual, received_size = received
+    assert received_size == size
+    if isinstance(payload, dict):
+        assert actual.keys() == payload.keys()
+        assert actual["label"] == payload["label"]
+        assert isinstance(actual["leaves"], tuple)
+        assert isinstance(actual["leaves"][1], list)
+        torch.testing.assert_close(actual["leaves"], payload["leaves"], rtol=0, atol=0)
+    else:
+        torch.testing.assert_close(actual, payload, rtol=0, atol=0)
+    calls = copying_native_agent
+    assert sum(len(pairs) for pairs in calls["transfers"]) == len(indices)
+    if case in ("empty", "all_empty"):
+        assert metadata["descriptor_groups"] == []
+        assert calls == {"agents": [], "descs": [], "transfers": [], "released": []}
+        assert size == 0
+    else:
+        assert len(calls["agents"]) == 1
+        assert len(calls["released"]) == 3 * len(calls["transfers"]) + 1
+    assert producer._pending == producer._published == {}
+    assert producer._registered_descs == producer._agent.registered == []
+    assert consumer._agent.registered == consumer._remote_agents == consumer._deferred_transfers == []
+
+
+@pytest.mark.parametrize(
+    "defect",
+    [
+        "missing_groups",
+        "missing_nonempty",
+        "duplicate",
+        "empty_index",
+        "zero_region",
+        "wrong_region_size",
+        "null_pointer",
+        "negative_device",
+        "short_region",
+        "bad_region_type",
+        "bool_index",
+        "float_index",
+        "empty_group",
+        "unequal_lengths",
+        "missing_memory_type",
+        "fake_empty_size",
+        "fake_nonempty_size",
+        "negative_shape",
+        "float_shape",
+        "bad_dtype",
+        "bad_spec",
+    ],
+)
+@pytest.mark.usefixtures("reliable_claim_queries")
+def test_invalid_empty_tensor_metadata_rejected_before_native_calls(producer, consumer, copying_native_agent, defect):
+    success, _, metadata = producer.put(
+        "0", "1", "req-invalid-empty", [torch.empty((0, 2)), torch.tensor(3.0), torch.empty((1, 0))]
+    )
+    assert success
+    group = metadata["descriptor_groups"][0]
+    region = group["regions"][0]
+    if defect == "missing_groups":
+        metadata.pop("descriptor_groups")
+    elif defect == "missing_nonempty":
+        metadata["descriptor_groups"] = []
+    elif defect == "duplicate":
+        group["tensor_indices"] *= 2
+        group["regions"] *= 2
+    elif defect == "empty_index":
+        group["tensor_indices"] = [0]
+    elif defect in ("zero_region", "wrong_region_size"):
+        group["regions"] = [(region[0], 0 if defect == "zero_region" else 8, 0, "")]
+    elif defect == "null_pointer":
+        group["regions"] = [(0, 4, 0, "")]
+    elif defect == "negative_device":
+        group["regions"] = [(region[0], 4, -1, "")]
+    elif defect == "short_region":
+        group["regions"] = [(region[0], 4)]
+    elif defect == "bad_region_type":
+        group["regions"] = [None]
+    elif defect in ("bool_index", "float_index"):
+        group["tensor_indices"] = [True if defect == "bool_index" else 1.0]
+    elif defect == "empty_group":
+        metadata["descriptor_groups"].append({"memory_type": "DRAM", "tensor_indices": [], "regions": []})
+    elif defect == "unequal_lengths":
+        group["regions"] = []
+    elif defect == "missing_memory_type":
+        group.pop("memory_type")
+    elif defect == "fake_empty_size":
+        metadata["tensor_specs"][1]["size"] = 0
+        metadata["descriptor_groups"] = []
+    elif defect == "fake_nonempty_size":
+        metadata["tensor_specs"][0]["size"] = 4
+    elif defect == "negative_shape":
+        metadata["tensor_specs"][0]["shape"] = [-1, 0]
+    elif defect == "float_shape":
+        metadata["tensor_specs"][0]["shape"] = [0.0, 2]
+    elif defect == "bad_dtype":
+        metadata["tensor_specs"][0]["dtype"] = "torch.not_a_dtype"
+    elif defect == "bad_spec":
+        metadata["tensor_specs"][0] = None
+    assert consumer.get("0", "1", "req-invalid-empty", metadata) is None
+    assert copying_native_agent == {"agents": [], "descs": [], "transfers": [], "released": []}
+    assert consumer._agent.registered == []
+    # Metadata acquisition now claims the source before validation. With no
+    # submitted DMA, failure safely completes that claim instead of leaking it;
+    # it must still report no successful payload receive.
+    assert producer._pending == producer._published == {}
+    assert producer._registered_descs == producer._agent.registered == []
+    assert consumer._metrics["gets"] == 0
+    assert consumer._metrics["errors"] == 1

tests/distributed/omni_connectors/test_nixl_connector_native.py

@@ -4,12 +4,13 @@
 import multiprocessing as mp
 import socket
 import time
+from dataclasses import fields
 from typing import Any
 
 import pytest
 import torch
 
-from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct
+from vllm_omni.data_entry_keys import EmbeddingsStruct, HiddenStatesStruct, MetaStruct, OmniPayloadStruct
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.parallel]
@@ -29,7 +30,83 @@ def _free_port() -> int:
         return int(sock.getsockname()[1])
 
 
-def _producer(port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any) -> None:
+def _h3_conditioning(payload_kind: str, device: str = "cpu"):
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderConditioning
+
+    media = {}
+    if payload_kind == "h3_ref2va":
+        media = {
+            "visual_condition": torch.arange(5 * 96, dtype=torch.float32).reshape(5, 96) / 8,
+            "visual_condition_shapes": ((1, 2, 2), (2, 2, 4)),
+            "audio_condition": torch.arange(10 * 32, dtype=torch.float32, device=device).reshape(10, 32) / 16,
+            "audio_condition_lengths": (2, 3),
+            "ref_blocks": (
+                {"kind": "image", "latent_t": 1, "latent_h": 2, "latent_w": 2},
+                {"kind": "video_audio", "ref_audio_t": 3, "latent_t": 2, "latent_h": 2, "latent_w": 4},
+            ),
+        }
+    elif payload_kind == "h3_fl2va":
+        media = {
+            "visual_condition": torch.arange(2 * 96, dtype=torch.float32).reshape(2, 96) / 8,
+            "visual_condition_shapes": ((1, 2, 2), (1, 2, 2)),
+            "keyframe_frame_indices": (0, 16),
+        }
+    return MiniMaxH3EncoderConditioning(
+        hidden_states=(torch.arange(3 * 5120, device=device).reshape(3, 5120) % 127).to(torch.bfloat16),
+        token_tags=torch.tensor([1, 0, 1], dtype=torch.int64),
+        task=payload_kind.removeprefix("h3_"),
+        height=256,
+        width=448,
+        num_frames=17,
+        latent_t=5,
+        audio_t=10,
+        **media,
+    )
+
+
+def _h3_structured_payload(payload_kind: str) -> OmniPayloadStruct:
+    # Only adapt to the existing public Struct; NIXL owns all serialization.
+    wire = _h3_conditioning(payload_kind, "cuda:0").to_omni_payload()
+    # Exercise both CPU and CUDA empty leaves (T2VA has two empty slots).
+    wire["embed"]["speech_feat"] = wire["embed"]["speech_feat"].to("cuda:0")
+    return OmniPayloadStruct(
+        hidden_states=HiddenStatesStruct(**wire["hidden_states"]),
+        meta=MetaStruct(**wire["meta"]),
+        embed=EmbeddingsStruct(**wire["embed"]),
+        kv_metadata=wire["kv_metadata"],
+        request_id="native-smoke",
+    )
+
+
+def _assert_h3_roundtrip(payload: dict[str, Any], payload_kind: str) -> tuple:
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderConditioning
+
+    expected = _h3_conditioning(payload_kind)
+    expected_wire = expected.to_omni_payload()
+    tensors = []
+    for section, values in expected_wire.items():
+        for key, reference in values.items():
+            actual = payload[section][key]
+            assert actual.device == torch.device("cuda:1"), (section, key, actual.device)
+            # Exact values, shapes and dtypes, including empty wire slots and layout.
+            torch.testing.assert_close(actual.cpu(), reference, rtol=0, atol=0)
+            tensors.append((f"{section}.{key}", str(actual.dtype), list(actual.shape), str(actual.device)))
+    actual = MiniMaxH3EncoderConditioning.from_omni_payload(payload)
+    for field in fields(expected):
+        reference = getattr(expected, field.name)
+        value = getattr(actual, field.name)
+        if isinstance(reference, torch.Tensor):
+            assert value.device == torch.device("cuda:1"), field.name
+            torch.testing.assert_close(value.cpu(), reference, rtol=0, atol=0)
+        else:
+            assert value == reference, field.name
+    assert payload["request_id"] == "native-smoke"
+    return (payload_kind, tensors)
+
+
+def _producer(
+    port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any, payload_kind: str
+) -> None:
     from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector
 
     current_omni_platform.set_device(0)
@@ -48,6 +125,8 @@ def _producer(port: int, ready: Any, consumed: Any, result: Any, claimed: Any, p
             kv_metadata={"empty": torch.empty((2, 0, 3), device="cuda:0")},
             request_id="native-smoke",
         )
+        if payload_kind != "generic":
+            payload = _h3_structured_payload(payload_kind)
         success, size, metadata = connector.put("0", "1", "native-smoke", payload)
         result.put(("put", success, size, metadata["kind"] if metadata else None))
         wire.put(metadata)
@@ -76,7 +155,15 @@ def _producer(port: int, ready: Any, consumed: Any, result: Any, claimed: Any, p
 
 
 def _consumer(
-    port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any, direct: bool
+    port: int,
+    ready: Any,
+    consumed: Any,
+    result: Any,
+    claimed: Any,
+    proceed: Any,
+    wire: Any,
+    direct: bool,
+    payload_kind: str,
 ) -> None:
     from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
         OmniKVCacheConfig,
@@ -133,11 +220,15 @@ def _consumer(
         if received is None:
             raise TimeoutError("native NIXL transfer did not complete")
         payload, size = received
+        result.put(("native_states", states))
+        if payload_kind != "generic":
+            result.put(("get", size, *_assert_h3_roundtrip(payload, payload_kind)))
+            consumed.set()
+            return
         hidden = payload["hidden_states"]["output"]
         token_roles = payload["meta"]["token_role_ids"]
         empty = payload["kv_metadata"]["empty"]
         assert empty.shape == (2, 0, 3) and str(empty.device) == "cuda:1"
-        result.put(("native_states", states))
         result.put(
             (
                 "get",
@@ -161,8 +252,9 @@ def _consumer(
 
 
 @pytest.mark.skipif(not _native_nixl_available(), reason="requires NIXL and at least two CUDA devices")
+@pytest.mark.parametrize("payload_kind", ["generic", "h3_ref2va", "h3_fl2va", "h3_t2va"])
 @pytest.mark.parametrize("direct", [False, True])
-def test_native_two_process_structured_mixed_device_transfer(direct):
+def test_native_two_process_structured_mixed_device_transfer(direct, payload_kind):
     context = mp.get_context("spawn")
     ready = context.Event()
     consumed = context.Event()
@@ -171,8 +263,12 @@ def test_native_two_process_structured_mixed_device_transfer(direct):
     wire = context.Queue()
     result = context.Queue()
     port = _free_port()
-    producer = context.Process(target=_producer, args=(port, ready, consumed, result, claimed, proceed, wire))
-    consumer = context.Process(target=_consumer, args=(port, ready, consumed, result, claimed, proceed, wire, direct))
+    producer = context.Process(
+        target=_producer, args=(port, ready, consumed, result, claimed, proceed, wire, payload_kind)
+    )
+    consumer = context.Process(
+        target=_consumer, args=(port, ready, consumed, result, claimed, proceed, wire, direct, payload_kind)
+    )
 
     producer.start()
     consumer.start()
@@ -180,8 +276,10 @@ def test_native_two_process_structured_mixed_device_transfer(direct):
     consumer.join(timeout=60)
     if producer.is_alive():
         producer.terminate()
+        producer.join(timeout=10)
     if consumer.is_alive():
         consumer.terminate()
+        consumer.join(timeout=10)
 
     records = [result.get(timeout=5) for _ in range(5)]
     print("Native ownership records:", records)
@@ -193,12 +291,17 @@ def test_native_two_process_structured_mixed_device_transfer(direct):
     assert records[0][3] == "structured"
     get_record = next(record for record in records if record[0] == "get")
     assert get_record[1] == records[0][2]
-    assert get_record[2:6] == (
-        list(range(8)),
-        "torch.float32",
-        [8],
-        "cuda:1",
-    )
-    assert get_record[6:] == ([1, 2, 3], "torch.int64", [3], "cuda:1")
+    if payload_kind == "generic":
+        assert get_record[2:6] == (
+            list(range(8)),
+            "torch.float32",
+            [8],
+            "cuda:1",
+        )
+        assert get_record[6:] == ([1, 2, 3], "torch.int64", [3], "cuda:1")
+    else:
+        assert get_record[2] == payload_kind
+        assert len(get_record[3]) == 5
+        print(f"Native H3 exact wire/semantic roundtrip: {get_record}")
     assert next(record for record in records if record[0] == "cleanup") == ("cleanup", 0, 0)
     assert next(record for record in records if record[0] == "source_retained") == ("source_retained", 1)

tests/distributed/omni_connectors/test_shm_connector.py

@@ -2,7 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for SharedMemoryConnector focusing on TP / CFG / metadata fallback."""
 
+import fcntl
 import os
+import time
 
 import pytest
 import torch
@@ -25,6 +27,15 @@ def connector():
 
 
 class TestKeyBasedReadWrite:
+    def test_deadline_receive_does_not_wait_for_writer_lock(self, connector):
+        key = "deadline_locked_payload"
+        connector.put("0", "1", key, {"value": 7})
+        with open(f"/dev/shm/shm_{key}_lockfile.lock", "rb+") as lock_file:
+            fcntl.flock(lock_file, fcntl.LOCK_EX)
+            assert connector.get_with_deadline("0", "1", key, deadline=time.monotonic() + 1) is None
+        assert connector.get_with_deadline("0", "1", key, deadline=time.monotonic() - 1) is None
+        assert connector.get_with_deadline("0", "1", key, deadline=time.monotonic() + 1)[0] == {"value": 7}
+
     def test_put_then_get_by_key(self, connector):
         data = {"hello": "world", "n": 42}
         ok, size, meta = connector.put("s0", "s1", "test_key_1", data)

tests/e2e/online_serving/test_wan22_t2v.py

@@ -5,7 +5,8 @@
 Online serving smoke for ``Wan-AI/Wan2.2-T2V-A14B-Diffusers`` (text-to-video via ``/v1/videos``).
 
 Uses a single ``default`` ``OmniServerParams`` row via ``_get_diffusion_feature_cases`` (no extra
-``server_args``). Multi-variant / parallel coverage lives in ``test_wan22_expansion.py`` (L4).
+``server_args``), with explicit startup budgets for loading both experts from slow storage.
+Multi-variant / parallel coverage lives in ``test_wan22_expansion.py`` (L4).
 
 From ``tests/``::
 
@@ -33,7 +34,7 @@ def _get_diffusion_feature_cases(model: str):
     """Return a single default ``OmniServerParams`` row (no extra ``server_args``)."""
     return [
         pytest.param(
-            OmniServerParams(model=model),
+            OmniServerParams(model=model, init_timeout=1800, stage_init_timeout=1800, startup_timeout=2100),
             id="default",
             marks=SINGLE_CARD_FEATURE_MARKS,
         ),

tests/engine/test_orchestrator_kv_sender_info.py

@@ -43,13 +43,16 @@ class _DummyDiffusionStage:
         self.engine_input_source = engine_input_source or [0]
         self.calls = []
 
-    async def add_request_async(self, request_id, prompt, sampling_params, kv_sender_info=None):
+    async def add_request_async(
+        self, request_id, prompt, sampling_params, kv_sender_info=None, payload_sender_info=None
+    ):
         self.calls.append(
             {
                 "request_id": request_id,
                 "prompt": prompt,
                 "sampling_params": sampling_params,
                 "kv_sender_info": kv_sender_info,
+                "payload_sender_info": payload_sender_info,
             }
         )
 
@@ -154,6 +157,31 @@ def test_payload_sender_endpoint_matches_resolver_with_unequal_replicas(outgoing
     assert sender_info["zmq_port"] != consumer.extra["sender_zmq_port"]
 
 
+@pytest.mark.parametrize("initialized", [False, True])
+def test_intermediate_payload_sender_ignores_incoming_kv_endpoint(initialized):
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 1
+    client.replica_id = 3
+    client.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            stage_connector_config={
+                "name": "NixlConnector",
+                "extra": {
+                    "role": "receiver",
+                    "host": "10.0.0.1",
+                    "zmq_port": 47000,
+                    "from_stage": 0,
+                    "outgoing": {"host": "10.0.0.2", "zmq_port": 48000, "from_stage": 1},
+                },
+            }
+        )
+    )
+    client._omni_kv_config = {"connector_config": {"type": "NixlConnector", "role": "receiver", "zmq_port": 47100}}
+    client._kv_sender_info = {"host": "10.0.0.1", "zmq_port": 47100} if initialized else None
+
+    assert client._build_payload_sender_info() == {"host": "10.0.0.2", "zmq_port": 51073}
+
+
 def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
     client = object.__new__(StageEngineCoreClient)
     client.stage_id = 0
@@ -235,6 +263,90 @@ def test_stage_engine_core_client_preserves_explicit_loopback_sender_host():
     }
 
 
+def test_stage_engine_core_client_payload_sender_uses_unresolved_base_port():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.replica_id = 0
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            stage_connector_config={
+                "name": "NixlConnector",
+                "extra": {"role": "sender", "from_stage": 0, "zmq_port": 50171},
+            },
+            omni_kv_config={"connector_config": {"type": "NixlConnector", "role": "sender", "zmq_port": 50171}},
+        )
+    )
+    client._omni_kv_config = client.vllm_config.model_config.omni_kv_config
+
+    assert client._build_payload_sender_info() == {
+        "host": "10.20.30.40",
+        "zmq_port": 50071,
+    }
+
+
+def test_stage_engine_core_client_payload_sender_preserves_stage_port_without_kv_config():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.replica_id = 0
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            stage_connector_config={
+                "name": "NixlConnector",
+                "extra": {"role": "sender", "from_stage": 0, "zmq_port": 50071},
+            },
+            omni_kv_config=None,
+        )
+    )
+    client._omni_kv_config = None
+
+    assert client._build_payload_sender_info() == {
+        "host": "10.20.30.40",
+        "zmq_port": 50071,
+    }
+
+
+def test_stage_engine_core_client_payload_sender_falls_back_to_kv_connector_config():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.replica_id = 0
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            stage_connector_config=None,
+            omni_kv_config={
+                "connector_config": {
+                    "type": "NixlConnector",
+                    "role": "sender",
+                    "zmq_port": 50171,
+                }
+            },
+        )
+    )
+    client._omni_kv_config = client.vllm_config.model_config.omni_kv_config
+
+    assert client._build_payload_sender_info() == {
+        "host": "10.20.30.40",
+        "zmq_port": 50071,
+    }
+
+
+def test_stage_engine_core_client_payload_sender_falls_back_to_initialized_kv_sender():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.replica_id = 0
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_connector_config=None))
+    client._omni_kv_config = None
+    client._kv_sender_info = {"host": "10.20.30.40", "zmq_port": 50171}
+
+    assert client._build_payload_sender_info() == {
+        "host": "10.20.30.40",
+        "zmq_port": 50071,
+    }
+
+
 def test_forward_to_diffusion_attaches_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
@@ -263,6 +375,36 @@ def test_forward_to_diffusion_attaches_kv_sender_info():
     assert req_state.stage_submit_ts[1] > 0
 
 
+def test_diffusion_resubmission_preserves_sender_endpoints():
+    orchestrator = object.__new__(Orchestrator)
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+    sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
+    payload_endpoint = {"host": "10.0.0.2", "zmq_port": 50071}
+    sender_pool.stage_client.get_payload_sender_info = lambda: payload_endpoint
+    diffusion_pool = StagePool(1, diffusion_stage)
+    orchestrator.num_stages = 2
+    orchestrator.stage_pools = [sender_pool, diffusion_pool]
+    orchestrator._cfg_tracker = CfgCompanionTracker()
+    state = OrchestratorRequestState(
+        request_id="streaming",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()],
+        final_stage_id=1,
+    )
+    output = SimpleNamespace(request_id="streaming", finished=True)
+
+    async def forward_twice():
+        await orchestrator._forward_to_next_stage("streaming", 0, output, state)
+        await orchestrator._forward_to_next_stage("streaming", 0, output, state)
+
+    asyncio.run(forward_twice())
+
+    assert len(diffusion_stage.calls) == 2
+    for call in diffusion_stage.calls:
+        assert call["kv_sender_info"] == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
+        assert call["payload_sender_info"] == payload_endpoint
+
+
 def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
@@ -293,7 +435,7 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
 def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
     orchestrator = object.__new__(Orchestrator)
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
-    diffusion_stage.custom_process_input_func = lambda *_args, **_kwargs: []
+    setattr(diffusion_stage, "custom_process_input_func", lambda *_args, **_kwargs: [])
     sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
     diffusion_pool = StagePool(1, diffusion_stage)
 

tests/helpers/runtime.py

@@ -150,6 +150,7 @@ class OmniServerParams(NamedTuple):
     use_stage_cli: bool = False
     init_timeout: int | None = None
     stage_init_timeout: int | None = None  # None: fixture supplies default (600 s)
+    startup_timeout: int = SERVER_STARTUP_TIMEOUT_S
 
 
 class OmniServer:
@@ -163,8 +164,10 @@ class OmniServer:
         port: int | None = None,
         env_dict: dict[str, str] | None = None,
         use_omni: bool = True,
+        startup_timeout: int = SERVER_STARTUP_TIMEOUT_S,
     ) -> None:
         cleanup_test_environment()
+        self.startup_timeout = startup_timeout
         self.model = model
         args = list(serve_args)
         self.serve_args = args
@@ -248,7 +251,7 @@ class OmniServer:
             cwd=_omni_subprocess_cwd(),
         )
 
-        max_wait = SERVER_STARTUP_TIMEOUT_S
+        max_wait = self.startup_timeout
         # System clock corrections must not shorten or extend startup waits.
         start_time = time.monotonic()
         while time.monotonic() - start_time < max_wait:
@@ -442,8 +445,11 @@ class OmniServerStageCli(OmniServer):
         stage_ids: list[int] | None = None,
         port: int | None = None,
         env_dict: dict[str, str] | None = None,
+        startup_timeout: int = SERVER_STARTUP_TIMEOUT_S,
     ) -> None:
-        super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True)
+        super().__init__(
+            model, serve_args or [], port=port, env_dict=env_dict, use_omni=True, startup_timeout=startup_timeout
+        )
         self.stage_config_path = stage_config_path
         self.master_port = get_open_port()
         resolved_cfg = resolve_deploy_yaml(stage_config_path)
@@ -553,7 +559,7 @@ class OmniServerStageCli(OmniServer):
             for replica_id in range(self.stage_replica_counts.get(stage_id, 1)):
                 self._launch_stage(stage_id, headless=True, replica_id=replica_id)
 
-        max_wait = SERVER_STARTUP_TIMEOUT_S
+        max_wait = self.startup_timeout
         # System clock corrections must not shorten or extend startup waits.
         start_time = time.monotonic()
         while time.monotonic() - start_time < max_wait:
@@ -1033,6 +1039,7 @@ def iter_omni_server(
                 server_args,
                 port=port,
                 env_dict=params.env_dict,
+                startup_timeout=params.startup_timeout,
             ) as server:
                 if model != original_model:
                     server.model = original_model
@@ -1050,6 +1057,7 @@ def iter_omni_server(
                     port=port,
                     env_dict=params.env_dict,
                     use_omni=params.use_omni,
+                    startup_timeout=params.startup_timeout,
                 )
                 if port
                 else OmniServer(
@@ -1057,6 +1065,7 @@ def iter_omni_server(
                     server_args,
                     env_dict=params.env_dict,
                     use_omni=params.use_omni,
+                    startup_timeout=params.startup_timeout,
                 )
             ) as server:
                 if model != original_model:

tests/helpers/tests/test_runtime_startup_timeout.py

@@ -38,7 +38,11 @@ def startup_server(request, tmp_path, monkeypatch):
 
 @pytest.mark.parametrize("wall_clock_jump", [3600, -3600])
 @pytest.mark.parametrize("becomes_ready", [True, False], ids=["ready", "timeout"])
-def test_startup_deadline_ignores_wall_clock_changes(startup_server, monkeypatch, wall_clock_jump, becomes_ready):
+@pytest.mark.parametrize("startup_timeout", [runtime.SERVER_STARTUP_TIMEOUT_S, 2100])
+def test_startup_deadline_ignores_wall_clock_changes(
+    startup_server, monkeypatch, wall_clock_jump, becomes_ready, startup_timeout
+):
+    startup_server.startup_timeout = startup_timeout
     elapsed = 0.0
     wall_offset = 0.0
     first_probe_at = None
@@ -66,7 +70,7 @@ def test_startup_deadline_ignores_wall_clock_changes(startup_server, monkeypatch
         if first_probe_at is None:
             first_probe_at = elapsed
         # Even a backwards clock correction must not extend the deadline.
-        assert elapsed - first_probe_at < runtime.SERVER_STARTUP_TIMEOUT_S, "startup exceeded its elapsed-time budget"
+        assert elapsed - first_probe_at < startup_timeout, "startup exceeded its elapsed-time budget"
         probes += 1
         if probes == 1:
             wall_offset += wall_clock_jump
@@ -84,7 +88,7 @@ def test_startup_deadline_ignores_wall_clock_changes(startup_server, monkeypatch
         assert first_probe_at is not None
         assert elapsed - first_probe_at == 2
     else:
-        with pytest.raises(RuntimeError, match=f"failed to start within {runtime.SERVER_STARTUP_TIMEOUT_S} seconds"):
+        with pytest.raises(RuntimeError, match=f"failed to start within {startup_timeout} seconds"):
             startup_server._start_server()
         assert first_probe_at is not None
-        assert elapsed - first_probe_at == runtime.SERVER_STARTUP_TIMEOUT_S
+        assert elapsed - first_probe_at == startup_timeout

tests/model_executor/models/test_minimax_h3_encoder_payload.py

@@ -91,8 +91,11 @@ def test_pipeline_owns_three_encoders_then_dit_and_decoders() -> None:
 
     encoder, diffusion = MINIMAX_H3_PIPELINE.stages
     assert encoder.model_arch == "MiniMaxH3Encoder"
-    assert encoder.custom_process_next_stage_input_func is None
+    assert encoder.custom_process_next_stage_input_func == (
+        "vllm_omni.model_executor.stage_input_processors.minimax_h3.encoder2diffusion_full_payload"
+    )
     assert diffusion.model_arch == "MiniMaxH3Pipeline"
+    assert diffusion.stage_input_payload_keys == ("encoder_output",)
     assert diffusion.requires_multimodal_data is False
     assert diffusion.requires_full_payload_input is False
 
@@ -441,7 +444,7 @@ def test_encode_media_keeps_audio_budgets_and_component_residency_separate(
     else:
         with pytest.raises(ValueError, match="at most 15 seconds"):
             encode_media(media, **kwargs)
-    assert scopes == [video_vae, audio_vae]
+    assert scopes == ([video_vae, audio_vae] if valid else [])
 
 
 def test_encoder_output_reuses_encoder_handoff_and_round_trips(monkeypatch) -> None:

tests/model_executor/stage_input_processors/test_minimax_h3.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Regression tests for MiniMax H3's disaggregated text-encoder contract."""
+"""Regression tests for MiniMax H3's disaggregated encoder contract."""
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -10,7 +10,7 @@ import pytest
 import torch
 from PIL import Image
 
-from vllm_omni.data_entry_keys import flatten_payload
+from vllm_omni.data_entry_keys import flatten_payload, to_struct
 from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import resolve_minimax_h3_diffusion_model_path
 from vllm_omni.engine.serialization import (
     deserialize_additional_information,
@@ -24,8 +24,10 @@ from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
 )
 from vllm_omni.model_executor.models.minimax_h3.conditioning import (
     MINIMAX_H3_CONDITION_LABELS_KEY,
+    MINIMAX_H3_ENCODER_LAYOUT_KEY,
     MINIMAX_H3_ENCODER_REQUEST_KEY,
     MINIMAX_H3_PRESENTATION_TASK_KEY,
+    MINIMAX_H3_TEXT_CONDITIONING_SCHEMA,
     MiniMaxH3EncoderConditioning,
 )
 from vllm_omni.model_executor.models.minimax_h3.encoder_processing import _audio_items, _load_audio
@@ -40,6 +42,7 @@ from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
 from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
     _diffusion_sampling_params,
     encoder2diffusion,
+    encoder2diffusion_full_payload,
     prepare_encoder_prompt,
 )
 
@@ -208,7 +211,8 @@ def test_prepare_ref2va_rejects_short_embedded_video_audio(monkeypatch):
         prepare_encoder_prompt(prompt, [sampling])
 
 
-def test_prepare_ref2va_keeps_embedded_and_standalone_audio_budgets_separate(monkeypatch):
+@pytest.mark.parametrize("standalone_count", [1, 2])
+def test_prepare_ref2va_uses_separate_embedded_and_standalone_audio_budgets(monkeypatch, standalone_count):
     duration_seconds = 8.0
     sample_rate = 16_000
     _mock_ref2va_video_with_audio(monkeypatch, duration_seconds=duration_seconds)
@@ -222,16 +226,21 @@ def test_prepare_ref2va_keeps_embedded_and_standalone_audio_budgets_separate(mon
         "prompt": "hello",
         "multi_modal_data": {
             "video": "original.mp4",
-            "audio": (torch.zeros(int(duration_seconds * sample_rate)), sample_rate),
+            "audio": [(torch.zeros(int(duration_seconds * sample_rate)), sample_rate) for _ in range(standalone_count)],
         },
     }
 
-    transformed = prepare_encoder_prompt(prompt, [sampling])
-    assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
-        ("audio", 1),
-        ("video", 1),
-        ("audio", 2),
-    ]
+    if standalone_count == 2:
+        with pytest.raises(OmniClientError, match="at most 15 seconds in total"):
+            prepare_encoder_prompt(prompt, [sampling])
+    else:
+        transformed = prepare_encoder_prompt(prompt, [sampling])
+        from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
+
+        media = MiniMaxH3Encoder._media_input(transformed["additional_information"])
+        assert len(media.video_audios) == 1 and len(media.audios) == 1
+        assert media.video_audios[0][0].shape[-1] == 8 * sample_rate
+        assert media.audios[0][0].shape[-1] == 8 * sample_rate
 
 
 def _encoder_output() -> dict:
@@ -283,6 +292,60 @@ def test_encoder2diffusion_waits_for_one_finished_source() -> None:
         encoder2diffusion([SimpleNamespace(), SimpleNamespace()], {"prompt": "hello"})
 
 
+@pytest.mark.parametrize(
+    ("case", "message"),
+    [
+        ("hidden_width", r"hidden_states must have shape \[tokens, 5120\]"),
+        ("hidden_dtype", "hidden_states must have dtype torch.bfloat16"),
+        ("hidden_noncontiguous", "hidden_states must use contiguous strided layout"),
+        ("hidden_non_strided", "hidden_states must use contiguous strided layout"),
+        ("token_dtype", "token_tags must have dtype torch.int64"),
+        ("token_layout", "token_tags must use contiguous strided layout"),
+        ("token_count", "token_tags must align with hidden_states"),
+        ("token_value", "token_tags must contain only 0 and 1"),
+    ],
+)
+@pytest.mark.parametrize("via_connector", [False, True])
+def test_encoder2diffusion_rejects_h3_v1_contract_mismatch(case: str, message: str, via_connector: bool) -> None:
+    payload = _encoder_output()
+    payload["hidden_states"]["output"] = torch.zeros(4, 5120, dtype=torch.bfloat16)
+    payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [0], [0]], dtype=torch.int64)
+
+    if case == "hidden_width":
+        payload["hidden_states"]["output"] = torch.zeros(4, 5119, dtype=torch.bfloat16)
+    elif case == "hidden_dtype":
+        payload["hidden_states"]["output"] = torch.zeros(4, 5120, dtype=torch.float32)
+    elif case == "hidden_noncontiguous":
+        payload["hidden_states"]["output"] = torch.empty(5120, 4, dtype=torch.bfloat16).t()
+    elif case == "hidden_non_strided":
+        payload["hidden_states"]["output"] = torch.empty(
+            (4, 5120),
+            dtype=torch.bfloat16,
+            layout=torch.sparse_coo,
+        )
+    elif case == "token_dtype":
+        payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [0], [0]], dtype=torch.int32)
+    elif case == "token_layout":
+        payload["meta"]["token_role_ids"] = torch.tensor(
+            [[1], [0], [1], [0], [1], [0], [1], [0]],
+            dtype=torch.int64,
+        )[::2]
+    elif case == "token_count":
+        payload["meta"]["token_role_ids"] = torch.tensor([[1], [0], [0]], dtype=torch.int64)
+    elif case == "token_value":
+        payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [2], [0]], dtype=torch.int64)
+    else:  # pragma: no cover - parameterization is exhaustive
+        raise AssertionError(f"unknown test case: {case}")
+
+    with pytest.raises(RuntimeError, match=message) as exc_info:
+        if via_connector:
+            encoder2diffusion_full_payload(pooling_output=payload)
+        else:
+            encoder2diffusion([_source_output(payload)], {"prompt": "hello"})
+
+    assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA in str(exc_info.value)
+
+
 def test_ref2va_one_image_tokens_and_tags_match_fused_presentation():
     tokenizer = _SegmentTokenizer()
     labels = [("image", 1), ("audio", 1)]
@@ -393,3 +456,236 @@ def test_diffusion_resolver_normalizes_partial_partition_directory(tmp_path):
     (ref2va / "text_encoder").mkdir(parents=True)
 
     assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, "ref2va") == str(ref2va)
+
+
+def _source_output(payload) -> SimpleNamespace:
+    return SimpleNamespace(
+        finished=True,
+        request_id="request-1",
+        outputs=[SimpleNamespace(multimodal_output=payload)],
+    )
+
+
+def _assert_same_tensor_view(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0, check_stride=True)
+    assert actual.data_ptr() == expected.data_ptr()
+    assert actual.untyped_storage().data_ptr() == expected.untyped_storage().data_ptr()
+
+
+def _full_encoder_output() -> dict:
+    return MiniMaxH3EncoderConditioning(
+        hidden_states=torch.randn(4, 5120, dtype=torch.bfloat16),
+        token_tags=torch.tensor([1, 1, 0, 0], dtype=torch.int64),
+        task="ref2va",
+        height=256,
+        width=448,
+        num_frames=17,
+        latent_t=5,
+        audio_t=10,
+        visual_condition=torch.arange(12 * 96, dtype=torch.float32).reshape(12, 96),
+        visual_condition_shapes=((1, 4, 4), (2, 4, 4)),
+        audio_condition=torch.arange(10 * 32, dtype=torch.float32).reshape(10, 32),
+        audio_condition_lengths=(3, 2),
+        ref_blocks=(
+            {"kind": "image", "latent_t": 1, "latent_h": 4, "latent_w": 4},
+            {"kind": "video_audio", "ref_audio_t": 3, "latent_t": 2, "latent_h": 4, "latent_w": 4},
+            {"kind": "audio", "ref_audio_t": 2},
+        ),
+        keyframe_frame_indices=(0, 16),
+    ).to_omni_payload()
+
+
+def _runner_payload(payload: dict) -> dict:
+    # Encoder.make_omni_output emits dotted private layout metadata. The runner
+    # selects each request's media tensors and flattens its nested categories.
+    payload = dict(payload)
+    payload[f"kv_metadata.{MINIMAX_H3_ENCODER_LAYOUT_KEY}"] = payload.pop("kv_metadata")[MINIMAX_H3_ENCODER_LAYOUT_KEY]
+    payload["meta"] = {"token_role_ids": payload["meta"]["token_role_ids"].reshape(-1, 1)}
+    return flatten_payload(payload)
+
+
+@pytest.mark.parametrize("representation", ["nested", "struct", "flat", "runner"])
+@pytest.mark.parametrize("with_media", [False, True])
+@pytest.mark.parametrize("via_connector", [False, True])
+def test_encoder_handoff_preserves_all_components(representation, with_media, via_connector):
+    payload = _full_encoder_output() if with_media else _encoder_output()
+    inputs = {
+        "nested": payload,
+        "struct": to_struct(payload),
+        "flat": flatten_payload(payload),
+        "runner": _runner_payload(payload),
+    }
+    if via_connector:
+        result = encoder2diffusion_full_payload(pooling_output=inputs[representation], request_id="request-1")
+        assert set(result) == {"encoder_output"}
+    else:
+        result = encoder2diffusion([_source_output(inputs[representation])], {"prompt": "hello"})[
+            "additional_information"
+        ]
+    output = result["encoder_output"]
+    for group in ("hidden_states", "embed", "meta", "kv_metadata"):
+        for key, tensor in payload[group].items():
+            torch.testing.assert_close(output[group][key], tensor, rtol=0, atol=0)
+    # Nonempty tensors must not be stripped, copied to CPU, or coerced to a
+    # different dtype by the adapter. Empty FP32 optional slots remain on wire.
+    assert output["hidden_states"]["output"] is payload["hidden_states"]["output"]
+    if with_media:
+        for key in ("embedding", "speech_feat"):
+            _assert_same_tensor_view(output["embed"][key], payload["embed"][key])
+    else:
+        for tensor in output["embed"].values():
+            assert tensor.shape == (0,)
+            assert tensor.dtype == torch.float32
+    for _ in range(2):
+        result = deserialize_additional_information(serialize_additional_information(result))
+    restored = MiniMaxH3EncoderConditioning.from_omni_payload(result["encoder_output"])
+    original = MiniMaxH3EncoderConditioning.from_omni_payload(payload)
+    for name in (
+        "task",
+        "height",
+        "width",
+        "num_frames",
+        "latent_t",
+        "audio_t",
+        "visual_condition_shapes",
+        "audio_condition_lengths",
+        "ref_blocks",
+        "keyframe_frame_indices",
+    ):
+        assert getattr(restored, name) == getattr(original, name)
+    for group, fields in output.items():
+        for key, tensor in fields.items():
+            torch.testing.assert_close(result["encoder_output"][group][key], tensor, rtol=0, atol=0)
+
+
+def test_full_payload_hook_skips_only_absent_output():
+    assert encoder2diffusion_full_payload(pooling_output=None) is None
+
+
+@pytest.mark.parametrize("via_connector", [False, True])
+@pytest.mark.parametrize(
+    ("case", "message"),
+    [
+        ("not_mapping", "no conditioning payload"),
+        ("text_only", "requires text, visual and audio tensors"),
+        ("missing_hidden", "requires text, visual and audio tensors"),
+        ("missing_tags", "requires text, visual and audio tensors"),
+        ("missing_layout", "requires private layout metadata"),
+        ("truncated_layout", "header is truncated"),
+        ("layout_dtype", "one-dimensional integer tensor"),
+        ("schema", "unsupported MiniMax H3 encoder wire schema"),
+        ("visual_shape", "visual condition must be FP32"),
+        ("audio_dtype", "audio condition must be FP32"),
+        ("empty_visual", "empty visual condition slot"),
+        ("empty_audio", "empty audio condition slot"),
+    ],
+)
+def test_encoder_handoff_rejects_invalid_full_payload(case, message, via_connector):
+    payload = _full_encoder_output()
+    layout = payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY]
+    if case == "not_mapping":
+        payload = []
+    elif case == "text_only":
+        payload.pop("embed")
+    elif case == "missing_hidden":
+        payload.pop("hidden_states")
+    elif case == "missing_tags":
+        payload["meta"] = {}
+    elif case == "missing_layout":
+        payload.pop("kv_metadata")
+    elif case == "truncated_layout":
+        payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY] = layout[:5]
+    elif case == "layout_dtype":
+        payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY] = layout.float()
+    elif case == "schema":
+        layout[1] = 99
+    elif case == "visual_shape":
+        payload["embed"]["embedding"] = torch.zeros(11, 96)
+    elif case == "audio_dtype":
+        payload["embed"]["speech_feat"] = payload["embed"]["speech_feat"].to(torch.bfloat16)
+    elif case == "empty_visual":
+        payload = _encoder_output()
+        payload["embed"]["embedding"] = torch.empty(0, 96)
+    elif case == "empty_audio":
+        payload = _encoder_output()
+        payload["embed"]["speech_feat"] = torch.empty(0, dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError, match=message):
+        if via_connector:
+            encoder2diffusion_full_payload(pooling_output=payload)
+        else:
+            encoder2diffusion([_source_output(payload)], {"prompt": "hello"})
+
+
+@pytest.mark.parametrize("inline", [False, True])
+def test_encoder_handoff_cleans_media_without_mutating_prompt_or_stripping_outputs(inline):
+    incoming = _full_encoder_output()
+    retained_hidden = torch.ones(1)
+    media = torch.zeros(2)
+    prompt = {
+        "prompt": "hello",
+        "negative_prompt": "blur",
+        "multi_modal_data": {"image": object(), "audio": object(), "video": object()},
+        "model_intermediate_buffer": {"media": media},
+        "additional_information": {
+            "global_request_id": ["request-1"],
+            "private": "preserved",
+            "hidden_states": {"layers": {0: media}, "output": retained_hidden},
+            "meta": {MINIMAX_H3_ENCODER_REQUEST_KEY: {"task": "ref2va"}, "private": 42},
+            "encoder_output": incoming,
+        },
+    }
+    source = _source_output(_runner_payload(incoming) if inline else None)
+    if not inline:
+        del source.outputs[0].multimodal_output
+    result = encoder2diffusion([source], [prompt])
+    info = result["additional_information"]
+    assert result["multi_modal_data"] is None
+    assert "model_intermediate_buffer" not in result
+    assert result["negative_prompt"] == "blur"
+    assert info["global_request_id"] == ["request-1"]
+    assert info["private"] == "preserved"
+    assert info["meta"] == {"private": 42}
+    assert set(info["hidden_states"]) == {"output"}
+    assert info["hidden_states"]["output"] is retained_hidden
+    for key in ("embedding", "speech_feat"):
+        _assert_same_tensor_view(info["encoder_output"]["embed"][key], incoming["embed"][key])
+    if not inline:
+        assert info["encoder_output"] is incoming
+    assert prompt["multi_modal_data"] is not None
+    assert prompt["model_intermediate_buffer"]["media"] is media
+    assert prompt["additional_information"]["hidden_states"]["layers"][0] is media
+    assert MINIMAX_H3_ENCODER_REQUEST_KEY in prompt["additional_information"]["meta"]
+
+
+def test_connector_handoff_without_inline_output_cleans_encoder_only_information():
+    prompt = {
+        "prompt": "hello",
+        "multi_modal_data": {"image": object()},
+        "model_intermediate_buffer": {"private": object()},
+        "additional_information": {
+            "hidden_states": {"layers": {0: torch.zeros(1)}},
+            "meta": {MINIMAX_H3_ENCODER_REQUEST_KEY: {"task": "t2va"}},
+        },
+    }
+    result = encoder2diffusion([_source_output(None)], prompt)
+    assert result == {"prompt": "hello", "multi_modal_data": None, "additional_information": {}}
+
+
+@pytest.mark.parametrize("payload", [None, _encoder_output()])
+@pytest.mark.parametrize("request_id", ["other", ["other"], ("other",)])
+def test_stage_wire_rejects_request_id_mismatch(payload, request_id):
+    prompt = {"prompt": "hello", "additional_information": {"global_request_id": request_id}}
+    with pytest.raises(RuntimeError, match="request ID does not match"):
+        encoder2diffusion([_source_output(payload)], prompt)
+
+
+@pytest.mark.parametrize("outputs", [None, [], (), [SimpleNamespace(), SimpleNamespace()]])
+def test_stage_wire_rejects_invalid_completion_count(outputs):
+    with pytest.raises(RuntimeError, match="exactly one completion"):
+        encoder2diffusion([SimpleNamespace(outputs=outputs)], {"prompt": "hello"})
+
+
+def test_stage_wire_empty_sources_and_invalid_prompt():
+    assert encoder2diffusion([], {"prompt": "hello"}) is None
+    with pytest.raises(TypeError, match="invalid MiniMax H3 prompt type"):
+        encoder2diffusion([_source_output(None)], object())

tests/worker/test_omni_connector_mixin.py

@@ -21,6 +21,7 @@ import torch
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
     OmniKVTransferManager,
 )
+from vllm_omni.distributed.omni_connectors.model_runner import omni_connector_runtime
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
     resolve_connector_spec,
@@ -59,6 +60,13 @@ class MockConnector:
     def close(self):
         pass
 
+    def get_with_deadline(self, from_stage, to_stage, get_key, metadata=None, *, deadline):
+        self.last_deadline = deadline
+        return self.get(from_stage, to_stage, get_key, metadata)
+
+    def abandon_get(self, get_key):
+        self.abandoned_key = get_key
+
 
 def _make_model_config(
     stage_id: int = 0,
@@ -91,6 +99,45 @@ class MixinHost(OmniConnectorModelRunnerMixin):
     pass
 
 
+def test_stage_payload_recv_spec_preserves_external_id_edge_and_handle():
+    sender = {"host": "10.0.0.1", "zmq_port": "50071"}
+    assert MixinHost._stage_payload_recv_spec("external", "2", "5", 3, sender) == (
+        "2",
+        "5",
+        "external_2_3",
+        {"source_host": "10.0.0.1", "source_port": 50071},
+    )
+    metadata = {"schema_version": 1}
+    handle = {"key": "explicit", "from_stage": 7, "to_stage": 9, "metadata": metadata}
+    assert MixinHost._stage_payload_recv_spec("external", "2", "5", sender_info=sender, handle=handle) == (
+        "7",
+        "9",
+        "explicit",
+        metadata,
+    )
+
+
+def test_synchronous_payload_init_borrows_connector_without_threads_or_kv_mutation():
+    host = MixinHost()
+    connector = MockConnector()
+    manager = SimpleNamespace(connector=connector, kv_recv_key_builder=object())
+    original_state = vars(manager).copy()
+    with patch.object(host, "_create_connector", side_effect=AssertionError("duplicate connector")):
+        host.init_omni_connectors(_make_model_config(), manager, synchronous=True)
+    assert host._omni_connector_initialized
+    assert host._recv_thread is None
+    assert host._save_thread is None
+    assert host._pending_load_reqs == {}
+    assert vars(manager) == original_state
+    connector.put("2", "5", "external_2_0", {"value": 7})
+    assert host.recv_stage_payload("external", "2", "5") == {"value": 7}
+    assert isinstance(connector.last_deadline, float)
+    assert connector.abandoned_key == "external_2_0"
+    with patch.object(connector, "close") as close:
+        host.shutdown_omni_connectors()
+        close.assert_not_called()
+
+
 class _FakeTPGroup:
     def __init__(self, *, world_size: int, rank_in_group: int, follower_result: Any = None):
         self.world_size = world_size
@@ -138,6 +185,121 @@ def test_init_payload_connector_ownership(role, custom_func, expected):
     host.shutdown_omni_connectors()
 
 
+@pytest.mark.parametrize("async_chunk", [False, True])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank):
+    host = MixinHost()
+    host.init_omni_connectors(_make_model_config(async_chunk=async_chunk))
+    # No background threads: deterministic interleaved polling.
+    host._stage_id = 1
+    host._omni_connector = MockConnector(stage_id=1)
+    connector = host._omni_connector
+    connector.sender_host = "configured-host"
+    connector.sender_zmq_port = 50051
+    connector.get = MagicMock(return_value=None)
+    first = _make_request("r1", "external1")
+    first.payload_sender_info = {"host": "producer-a", "zmq_port": "50101"}
+    second = _make_request("r2", "external2")
+    second.payload_sender_info = {"host": "producer-b", "zmq_port": 51101}
+    tp_group = _FakeTPGroup(world_size=2, rank_in_group=rank)
+    try:
+        with patch.object(host, "_get_local_tp_group", return_value=tp_group):
+            host.register_chunk_recv(first)
+            host.register_chunk_recv(second)
+            for req_id in ("r2", "r1", "r2"):
+                assert not host._poll_single_request(req_id)
+            if rank == 1:
+                connector.get.assert_not_called()
+            else:
+                assert connector.get.call_args_list == [
+                    unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
+                    unittest.mock.call("0", "1", "external1_0_0", {"source_host": "producer-a", "source_port": 50101}),
+                    unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}),
+                ]
+                connector.get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1)
+                assert host._poll_single_request("r2")
+                assert "r2" not in host._pending_load_reqs
+                assert "r1" in host._pending_load_reqs
+            host.cleanup_finished_request("r1")
+            host.cleanup_finished_request("r2")
+            assert not host._pending_load_reqs
+            # Reused internal ID without sender info cannot inherit either endpoint.
+            replacement = _make_request("r1", "external-new")
+            connector.get.reset_mock(return_value=True)
+            connector.get.return_value = None
+            host.register_chunk_recv(replacement)
+            assert not host._poll_single_request("r1")
+            if rank == 0:
+                connector.get.assert_called_once_with("0", "1", "external-new_0_0")
+            else:
+                connector.get.assert_not_called()
+            assert connector.sender_host == "configured-host"
+            assert connector.sender_zmq_port == 50051
+            assert first.payload_sender_info == {"host": "producer-a", "zmq_port": "50101"}
+    finally:
+        host.cleanup_finished_request("r1")
+        host.shutdown_omni_connectors()
+
+
+@pytest.mark.parametrize("role", ["sender", "receiver"])
+def test_split_factory_resolves_rank_replica_without_mutating_config(role):
+    extra = {"role": role, "zmq_port": 50071, "from_stage": 2, "host": "producer"}
+    config = _make_model_config()
+    config.stage_id = 2 if role == "sender" else 3
+    config.stage_connector_config = {"name": "NixlConnector", "extra": extra}
+    with (
+        patch.object(omni_connector_runtime, "get_local_tp_rank", return_value=3),
+        patch.object(omni_connector_runtime, "get_omni_replica_id", return_value=2),
+        patch.object(omni_connector_runtime.OmniConnectorFactory, "create_connector") as create,
+    ):
+        assert MixinHost._create_connector(config) is create.return_value
+    spec = create.call_args.args[0]
+    assert spec.name == "NixlConnector"
+    assert spec.extra["stage_id"] == config.stage_id
+    assert spec.extra["role"] == role
+    assert spec.extra["zmq_port" if role == "sender" else "sender_zmq_port"] == 50071 + 2 + 3 * 16 + 2 * 1024
+    if role == "receiver":
+        assert "zmq_port" not in spec.extra
+        assert spec.extra["sender_host"] == "producer"
+    assert extra == {"role": role, "zmq_port": 50071, "from_stage": 2, "host": "producer"}
+
+
+@pytest.mark.parametrize("rank", [None, 0, 1])
+def test_split_transfer_rank_without_runtime_initialization(rank):
+    host = MixinHost()
+    group = None if rank is None else _FakeTPGroup(world_size=2, rank_in_group=rank)
+    with patch.object(host, "_get_local_tp_group", return_value=group):
+        assert host.is_data_transfer_rank() is (rank != 1)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_split_full_payload_hook_and_send_are_leader_only(rank):
+    host = MixinHost()
+    host.init_omni_connectors(_make_model_config())
+    host._omni_connector = MockConnector()
+    output = {"encoder_output": {"video": torch.ones(1), "audio": torch.zeros(1)}}
+    host._custom_process_func = MagicMock(return_value=output)
+    group = _FakeTPGroup(world_size=2, rank_in_group=rank)
+    try:
+        with patch.object(host, "_get_local_tp_group", return_value=group):
+            assert host.send_full_payload_outputs(None, {"r1": (output, _make_request("r1"))}) == ["r1"]
+        if rank == 1:
+            host._custom_process_func.assert_not_called()
+            assert not host._pending_save_reqs
+        else:
+            host._custom_process_func.assert_called_once()
+            task = host._pending_save_reqs["r1"].popleft()
+            assert task["data"] is output
+            host.cleanup_finished_request("r1")
+            assert "r1" in host._deferred_send_cleanup
+            assert host._send_single_request(task)
+            assert not host._pending_save_counts
+            assert not host._deferred_send_cleanup
+            assert host._omni_connector.get("0", "1", "r1_0_0")[0] is output
+    finally:
+        host.shutdown_omni_connectors()
+
+
 class TestMixinAsyncChunkSendRecv(unittest.TestCase):
     """Test 2: Async chunk send/recv + bg threads."""
 

vllm_omni/config/omni_config.py

@@ -887,6 +887,9 @@ class _DiffusionConfigProjection:
     custom_pipeline_args: dict[str, Any] | None = None
     additional_config: dict[str, Any] = field(default_factory=dict)
     kv_transfer_config: KVTransferConfig | None = None
+    # Full stage-payload transport, independent of native paged KV transfer.
+    stage_input_payload_keys: tuple[str, ...] = ()
+    stage_output_payload_keys: tuple[str, ...] = ()
     enable_stage_verification: bool = True
     prompt_file_path: str | None = None
     quantization_config: _QuantizationConfigType = None
@@ -2158,6 +2161,10 @@ def _build_diffusion_config_projection(
     quantization_config: _QuantizationConfigType,
 ) -> _DiffusionConfigProjection:
     diffusion_kwargs = engine.to_kwargs()
+    # Match the legacy builder: topology supplies defaults, while explicit
+    # deploy/CLI values (including empty tuples) retain precedence.
+    diffusion_kwargs.setdefault("stage_input_payload_keys", tuple(topology.stage_input_payload_keys))
+    diffusion_kwargs.setdefault("stage_output_payload_keys", tuple(topology.stage_output_payload_keys))
     diffusion_kwargs["stage_id"] = topology.stage_id
     diffusion_kwargs["model_arch"] = _first_defined(
         diffusion_kwargs.get("model_arch"),

vllm_omni/config/stage_config.py

@@ -257,6 +257,15 @@ class StagePipelineConfig:
     prompt_transform_func: str | None = None
     prompt_expand_func: str | None = None
     cfg_kv_collect_func: str | None = None
+    # Payload keys this stage expects to receive from its upstream stage over
+    # the omni connector instead of through the orchestrator IPC hop. Declaring
+    # them is what enables the worker-side connector receive path, so an empty
+    # tuple leaves existing deployments untouched.
+    stage_input_payload_keys: tuple[str, ...] = ()
+    # Mirror of the above on the producing side: keys this stage hands to the
+    # next stage over the connector. Only diffusion producers need this; AR
+    # stages already send through ``send_full_payload_outputs``.
+    stage_output_payload_keys: tuple[str, ...] = ()
     omni_kv_config: dict[str, Any] | None = None
     scheduler_cls: str | None = None
     # Model subdirectory indirections: for multi-component HF repos where the
@@ -1010,6 +1019,10 @@ def _build_engine_args(
     if ps.model_path_resolver:
         engine_args["model_path_resolver"] = ps.model_path_resolver
     engine_args["inline_diffusion"] = ps.inline_diffusion
+    if ps.stage_input_payload_keys:
+        engine_args["stage_input_payload_keys"] = tuple(ps.stage_input_payload_keys)
+    if ps.stage_output_payload_keys:
+        engine_args["stage_output_payload_keys"] = tuple(ps.stage_output_payload_keys)
 
     # Pipeline-wide top-level DeployConfig settings, applied to every stage.
     for name in _PIPELINE_WIDE_ENGINE_FIELDS:

vllm_omni/deploy/minimax_h3_disaggregated.yaml

@@ -4,12 +4,18 @@ async_chunk: false
 trust_remote_code: true
 distributed_executor_backend: mp
 
+connectors:
+  shared_memory_connector:
+    name: SharedMemoryConnector
+
 stages:
   - stage_id: 0
     devices: "0,1"
     max_num_seqs: 1
     gpu_memory_utilization: 0.6
     tensor_parallel_size: 2
+    output_connectors:
+      to_stage_1: shared_memory_connector
     hf_overrides:
       minimax_h3_encoder_components:
         text_encoder:
@@ -33,6 +39,8 @@ stages:
       tensor_parallel_size: 1
       ulysses_degree: 4
       vae_patch_parallel_size: 4
+    input_connectors:
+      from_stage_0: shared_memory_connector
     default_sampling_params:
       seed: 42
       num_inference_steps: 50

vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml

@@ -0,0 +1,17 @@
+base_config: minimax_h3_disaggregated.yaml
+
+connectors:
+  nixl_connector:
+    name: NixlConnector
+    extra:
+      host: "auto"
+      zmq_port: 50071
+      backends: ["UCX"]
+
+stages:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: nixl_connector
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: nixl_connector

vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml

@@ -4,12 +4,18 @@ async_chunk: false
 trust_remote_code: true
 distributed_executor_backend: mp
 
+connectors:
+  shared_memory_connector:
+    name: SharedMemoryConnector
+
 stages:
   - stage_id: 0
     devices: "0,1"
     max_num_seqs: 1
     gpu_memory_utilization: 0.6
     tensor_parallel_size: 2
+    output_connectors:
+      to_stage_1: shared_memory_connector
     hf_overrides:
       minimax_h3_encoder_components:
         text_encoder:
@@ -33,6 +39,8 @@ stages:
       tensor_parallel_size: 1
       ulysses_degree: 4
       vae_patch_parallel_size: 4
+    input_connectors:
+      from_stage_0: shared_memory_connector
     default_sampling_params:
       seed: 42
       num_inference_steps: 5

vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml

@@ -0,0 +1,17 @@
+base_config: minimax_h3_disaggregated_turbo.yaml
+
+connectors:
+  nixl_connector:
+    name: NixlConnector
+    extra:
+      host: "auto"
+      zmq_port: 50071
+      backends: ["UCX"]
+
+stages:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: nixl_connector
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: nixl_connector

vllm_omni/diffusion/data.py

@@ -1060,6 +1060,15 @@ class OmniDiffusionConfig:
     # Model-specific function for collecting CFG KV caches (set at runtime)
     cfg_kv_collect_func: Any | None = None
 
+    # Conditioning keys fetched from the upstream stage over the omni connector
+    # rather than carried inline through the orchestrator. Empty disables the
+    # worker-side connector receive path.
+    stage_input_payload_keys: tuple[str, ...] = ()
+
+    # Keys handed to the next stage over the omni connector. Empty disables the
+    # worker-side connector send path.
+    stage_output_payload_keys: tuple[str, ...] = ()
+
     # Quantization: str method name, dict config, QuantizationConfig, or None.
     # str is resolved to {"method": <str>} internally.
     # Per-component: {"transformer": {"method": "fp8"}, "vae": None}
@@ -1183,6 +1192,8 @@ class OmniDiffusionConfig:
             materialize_legacy_offload_flags,
         )
 
+        self.stage_input_payload_keys = tuple(self.stage_input_payload_keys)
+        self.stage_output_payload_keys = tuple(self.stage_output_payload_keys)
         if self.vae_fast_path not in VAE_FAST_PATH_LEVELS:
             raise ValueError(f"vae_fast_path must be one of {list(VAE_FAST_PATH_LEVELS)}, got {self.vae_fast_path!r}")
         if self.diffusion_compile_granularity not in {"regional", "full"}:

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -384,6 +384,7 @@ class GroupCoordinator:
         group = self.device_group
         metadata_group = self.cpu_group
         assert src < self.world_size, f"Invalid src rank ({src})"
+        source_rank_in_group = src
         src = self.ranks[src]
 
         rank = self.rank
@@ -394,7 +395,7 @@ class GroupCoordinator:
             # `metadata_list` lives in CPU memory.
             # `broadcast_object_list` has serialization & deserialization,
             # all happening on CPU. Therefore, we can use the CPU group.
-            self.broadcast_object(metadata_list, src=src)
+            self.broadcast_object(metadata_list, src=source_rank_in_group)
             async_handles = []
             for tensor in tensor_list:
                 if tensor.numel() == 0:
@@ -411,7 +412,7 @@ class GroupCoordinator:
                 async_handle.wait()
 
         else:
-            metadata_list = self.broadcast_object(None, src=src)
+            metadata_list = self.broadcast_object(None, src=source_rank_in_group)
             tensor_dict = {}
             async_handles = []
             for key, value in metadata_list:

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -111,6 +111,7 @@ class InlineStageDiffusionClient(StageClientBase):
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
         kv_transfer_params: dict[str, Any] | None = None,
+        payload_sender_info: dict[str, Any] | None = None,
     ) -> None:
         # Each request mutates its sampling state while it is normalized and
         # executed. Callers commonly reuse one params object for concurrent
@@ -128,7 +129,8 @@ class InlineStageDiffusionClient(StageClientBase):
                 prompt,
                 sampling_params,
                 kv_sender_info,
-                kv_transfer_params,
+                kv_transfer_params=kv_transfer_params,
+                payload_sender_info=payload_sender_info,
             )
         )
         self._tasks[request_id] = task
@@ -140,6 +142,7 @@ class InlineStageDiffusionClient(StageClientBase):
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[str, Any] | None = None,
         kv_transfer_params: dict[str, Any] | None = None,
+        payload_sender_info: dict[str, Any] | None = None,
     ) -> None:
         try:
             request = OmniDiffusionRequest(
@@ -148,6 +151,7 @@ class InlineStageDiffusionClient(StageClientBase):
                 request_id=request_id,
                 kv_sender_info=kv_sender_info,
                 kv_transfer_params=kv_transfer_params,
+                payload_sender_info=payload_sender_info,
             )
 
             if self.od_config.streaming_output:

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -2011,6 +2011,10 @@ class MiniMaxH3Pipeline(
             return None
         additional_information = raw_prompt.get("additional_information") or {}
         encoder_output = additional_information.get("encoder_output")
+        if encoder_output is None:
+            # Preserve main's text-only handoff for deployments that still
+            # encode media locally; a supplied unified payload takes priority.
+            encoder_output = additional_information.get("text_encoder_output")
         if encoder_output is None:
             return None
         if not isinstance(encoder_output, Mapping):
@@ -2210,6 +2214,12 @@ class MiniMaxH3Pipeline(
                     self._validate_turbo_sampling(sampling, turbo_spec)
                 if has_native_lora:
                     self._validate_native_sampling(sampling, task=task)
+                if self._fasth3 is not None:
+                    self._fasth3.check_request(
+                        sampling,
+                        video_shift=self.default_video_shift,
+                        audio_shift=self.default_audio_shift,
+                    )
                 text_conditioning = self._extract_text_conditioning(raw_prompt)
                 if require_external_text and text_conditioning is None:
                     raise OmniClientError(

vllm_omni/diffusion/request.py

@@ -55,6 +55,7 @@ class OmniDiffusionRequest:
     sampling_params: OmniDiffusionSamplingParams
     request_id: str
     kv_sender_info: dict[str, Any] | None = None
+    payload_sender_info: dict[str, Any] | None = None
     # Optional opaque, model-owned input prepared before Scheduler admission.
     # Model code validates its concrete type when consuming it on the Worker.
     prepared_layout: Any | None = None

vllm_omni/diffusion/stage_diffusion_client.py

@@ -340,6 +340,7 @@ class StageDiffusionClient(StageClientBase):
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
         kv_transfer_params: dict[str, Any] | None = None,
+        payload_sender_info: dict[str, Any] | None = None,
     ) -> None:
         if self._engine_dead:
             raise EngineDeadError()
@@ -358,6 +359,7 @@ class StageDiffusionClient(StageClientBase):
                     "sampling_params": self._sampling_params_to_dict(sampling_params),
                     "kv_sender_info": kv_sender_info,
                     "kv_transfer_params": kv_transfer_params,
+                    "payload_sender_info": payload_sender_info,
                 }
             )
         )

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -168,6 +168,7 @@ class StageDiffusionProc:
         kv_sender_info: dict[str, Any] | None = None,
         on_request_started: Callable[[OmniRequestOutput], Awaitable[None]] | None = None,
         kv_transfer_params: dict[str, Any] | None = None,
+        payload_sender_info: dict[str, Any] | None = None,
     ) -> OmniRequestOutput:
         """Build a diffusion request and consume DiffusionEngine.step_streaming() to completion."""
         sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
@@ -178,6 +179,7 @@ class StageDiffusionProc:
             request_id=request_id,
             kv_sender_info=kv_sender_info,
             kv_transfer_params=kv_transfer_params,
+            payload_sender_info=payload_sender_info,
         )
 
         # Non-streaming callers share the streaming engine path but only
@@ -204,6 +206,7 @@ class StageDiffusionProc:
         sampling_params_dict: dict,
         kv_sender_info: dict[str, Any] | None = None,
         kv_transfer_params: dict[str, Any] | None = None,
+        payload_sender_info: dict[str, Any] | None = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         """Process a streaming diffusion request and yield the results from DiffusionEngine.step_streaming()."""
         sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
@@ -214,6 +217,7 @@ class StageDiffusionProc:
             request_id=request_id,
             kv_sender_info=kv_sender_info,
             kv_transfer_params=kv_transfer_params,
+            payload_sender_info=payload_sender_info,
         )
 
         async for results in self._engine.step_streaming(request):  # pyright: ignore[reportOptionalMemberAccess]
@@ -360,6 +364,7 @@ class StageDiffusionProc:
             sampling_params_dict: dict,
             kv_sender_info: dict[str, Any] | None = None,
             kv_transfer_params: dict[str, Any] | None = None,
+            payload_sender_info: dict[str, Any] | None = None,
         ) -> None:
             """Process a single diffusion request and send the response."""
             try:
@@ -375,6 +380,7 @@ class StageDiffusionProc:
                         kv_sender_info=kv_sender_info,
                         on_request_started=_send_request_started,
                         kv_transfer_params=kv_transfer_params,
+                        payload_sender_info=payload_sender_info,
                     )
                     await response_socket.send(encoder.encode({"type": "result", "output": result}))
                 else:
@@ -384,6 +390,7 @@ class StageDiffusionProc:
                         sampling_params_dict,
                         kv_sender_info=kv_sender_info,
                         kv_transfer_params=kv_transfer_params,
+                        payload_sender_info=payload_sender_info,
                     ):
                         await response_socket.send(encoder.encode({"type": "result", "output": result}))
             except DiffusionRequestAbortedError as e:
@@ -462,7 +469,8 @@ class StageDiffusionProc:
                             msg["prompt"],
                             msg["sampling_params"],
                             msg.get("kv_sender_info"),
-                            msg.get("kv_transfer_params"),
+                            kv_transfer_params=msg.get("kv_transfer_params"),
+                            payload_sender_info=msg.get("payload_sender_info"),
                         )
                     )
                     tasks[request_id] = task

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -66,6 +66,7 @@ from vllm_omni.diffusion.sched.interface import (
 )
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.diffusion.worker.stage_payload import DiffusionStagePayloadMixin
 from vllm_omni.diffusion.worker.utils import (
     BatchRunnerOutput,
     RunnerOutput,
@@ -78,7 +79,6 @@ from vllm_omni.diffusion.worker.utils import (
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
-from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 if TYPE_CHECKING:
     from vllm_omni.inputs.data import OmniInteractionPrompt
@@ -148,7 +148,7 @@ def _normalize_pipeline_outputs(
     return outputs
 
 
-class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
+class DiffusionModelRunner(DiffusionStagePayloadMixin):
     """
     Model runner that handles model loading and execution for diffusion models.
 
@@ -194,11 +194,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self.state_cache: dict[str, StepRequestState] = {}
 
         # Initialize KV cache manager for connector management.
+        payload_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
         self.kv_transfer_manager = (
-            OmniKVTransferManager.from_od_config(od_config)
-            if getattr(od_config, "kv_transfer_config", None) is None
-            else None
+            payload_transfer_manager if getattr(od_config, "kv_transfer_config", None) is None else None
         )
+        self.init_omni_connectors(od_config, payload_transfer_manager, synchronous=True)
         self._kv_connector = None
         from vllm_omni.diffusion.diffusion_kv.kv_connector import KVReceiveProgress, native_prefetch_enabled
 
@@ -629,6 +629,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         kv_prefetch_job: KVPrefetchJob | None = None,
         use_prefetch: bool = False,
     ) -> None:
+        # Fetch upstream conditioning before anything else: the pipeline reads
+        # it out of the prompt during the forward below.
+        self._maybe_recv_stage_payload(req)
+
         if self.kv_transfer_manager is None:
             self._initialize_generator(req.sampling_params)
             return
@@ -843,6 +847,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             ):
                 cache_summary(self.pipeline, details=True)
 
+        self._maybe_send_stage_payload(reqs, outputs)
+
         return self._runner_output_from_outputs(reqs, outputs)
 
     def _attach_stepwise_metadata(
@@ -1039,12 +1045,14 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 new_request_ids.append(request_id)
                 if request_id in self.state_cache:
                     raise ValueError(f"Received duplicate new-request payload for cached request {request_id}.")
+                self._maybe_recv_stage_payload(sched_new_req.req)
                 new_state = StepRequestState(
                     request_id=request_id,
                     sampling=copy.deepcopy(sched_new_req.req.sampling_params),
                     prompt=sched_new_req.req.prompt,
                     kv_sender_info=sched_new_req.req.kv_sender_info,
                     prepared_layout=getattr(sched_new_req.req, "prepared_layout", None),
+                    external_req_id=getattr(sched_new_req.req, "external_req_id", None),
                 )
                 if sched_new_req.diffusion_kv_metadata is not None:
                     new_state.extra["kv_computed_tokens"] = tuple(
@@ -1387,6 +1395,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                 if self.od_config.streaming_output
                                 else req.denoise_completed
                             )
+                            if finished and result is not None:
+                                self._maybe_send_stage_payload([req], [result])
                             runner_output_list.append(
                                 RunnerOutput(
                                     request_id=req.request_id,

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -261,6 +261,7 @@ class DiffusionWorker:
         self.init_snapshot: MemorySnapshot | None = None
         self.requested_memory: int | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
+        self._owns_sleep_pool = False
         self.lora_manager: DiffusionLoRAManager | None = None
         # Worker-side cache of (lora_request, lora_scale) per scheduled
         # request id. Used by step mode to recover LoRA identity for cached
@@ -1010,6 +1011,7 @@ class DiffusionWorker:
             allocator = CuMemAllocator.get_instance()
             if tag == "weights":
                 assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
+                self._owns_sleep_pool = True
             logger.info(f"[Worker {self.rank}] Activating Diffusion CuMem pool for tag: {tag}")
             return allocator.use_memory_pool(tag=tag)
         return nullcontext()
@@ -1019,25 +1021,36 @@ class DiffusionWorker:
         try:
             if self.model_runner is not None:
                 mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+                if mgr is None:
+                    mgr = getattr(self.model_runner, "_kv_transfer_manager", None)
                 try:
                     offload_backend = getattr(self.model_runner, "offload_backend", None)
                     if offload_backend is not None:
                         offload_backend.disable()
                 finally:
                     if mgr is not None:
-                        mgr.shutdown_prefetch()
+                        mgr.close()
         finally:
+            self.model_runner = None
+            self.lora_manager = None
+            self._sleep_saved_buffers = {}
             try:
-                shutdown_kv_connector()
+                if getattr(self, "_owns_sleep_pool", False):
+                    gc.collect()
+                    _get_cumem_allocator_class().get_instance().release_pools()
+                    self._owns_sleep_pool = False
             finally:
                 try:
-                    a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
-                    if a2a_permute is not None:
-                        a2a_permute.clear_a2a_permute_workspaces()
-                except Exception:
-                    logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
+                    shutdown_kv_connector()
                 finally:
-                    destroy_distributed_env()
+                    try:
+                        a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
+                        if a2a_permute is not None:
+                            a2a_permute.clear_a2a_permute_workspaces()
+                    except Exception:
+                        logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
+                    finally:
+                        destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:
@@ -1642,7 +1655,14 @@ class WorkerWrapperBase:
 
         # Re-initialize pipeline with custom pipeline if provided
         if self.uses_custom_pipeline:
-            self.worker.re_init_pipeline(self.custom_pipeline_args)
+            try:
+                self.worker.re_init_pipeline(self.custom_pipeline_args)
+            except Exception:
+                try:
+                    self.worker.shutdown()
+                except Exception:
+                    logger.exception("Failed to clean up worker after custom pipeline initialization failure")
+                raise
 
     def _prepare_worker_class(self) -> type:
         """

vllm_omni/diffusion/worker/stage_payload.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+
+logger = init_logger(__name__)
+
+
+def _to_device(value: Any, device: torch.device) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.to(device) if value.device != device else value
+    if isinstance(value, dict):
+        return {key: _to_device(item, device) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        moved = [_to_device(item, device) for item in value]
+        return moved if isinstance(value, list) else tuple(moved)
+    return value
+
+
+class DiffusionStagePayloadMixin(OmniConnectorModelRunnerMixin):
+    """Adapt diffusion prompts and outputs to shared connector transport."""
+
+    _STAGE_PAYLOAD_HANDLE_KEY = "_stage_payload_transfer"
+
+    def _stage_payload_broadcast_groups(self) -> tuple[Any, ...]:
+        from vllm_omni.diffusion.distributed.parallel_state import get_fs_group, get_hsdp_replicate_group
+
+        hsdp_groups = []
+        for get_group in (get_fs_group, get_hsdp_replicate_group):
+            try:
+                group = get_group()
+            except AssertionError:
+                continue
+            if group is not None and getattr(group, "world_size", 1) > 1:
+                hsdp_groups.append(group)
+        if hsdp_groups:
+            return tuple(hsdp_groups)
+
+        groups = list(super()._stage_payload_broadcast_groups())
+        try:
+            from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+
+            sp_group = get_sp_group()
+        except (AssertionError, ImportError):
+            sp_group = None
+        if sp_group is not None and getattr(sp_group, "world_size", 1) > 1:
+            groups.append(sp_group)
+        return tuple(groups)
+
+    def _maybe_recv_stage_payload(self, req: OmniDiffusionRequest) -> None:
+        if OmniDiffusionRequest.is_dummy_run_request_id(req.request_id):
+            return
+        prompt = getattr(req, "prompt", None)
+        if not isinstance(prompt, dict):
+            return
+        handle = prompt.pop(self._STAGE_PAYLOAD_HANDLE_KEY, None)
+        expected_keys = tuple(getattr(self.od_config, "stage_input_payload_keys", ()) or ())
+        if not isinstance(handle, dict) and not expected_keys:
+            return
+        from_stage, to_stage = self._kv_transfer_manager.recv_stages
+        if not isinstance(handle, dict) and (from_stage is None or to_stage is None):
+            raise RuntimeError(f"Stage {self.od_config.stage_id} expects a payload but has no incoming edge")
+        sender_info = getattr(req, "payload_sender_info", None) or getattr(req, "kv_sender_info", None)
+        if isinstance(sender_info, dict):
+            sender_stage = handle.get("from_stage", from_stage) if isinstance(handle, dict) else from_stage
+            sender_info = self._kv_transfer_manager._resolve_sender_info(sender_info, sender_stage)
+        payload = self.recv_stage_payload(
+            getattr(req, "external_req_id", None) or req.request_id,
+            str(from_stage),
+            str(to_stage),
+            sender_info=sender_info,
+            handle=handle if isinstance(handle, dict) else None,
+        )
+        additional = prompt.setdefault("additional_information", {})
+        if isinstance(payload, dict):
+            target_device = self._target_device or self.device
+            for name, value in payload.items():
+                if not expected_keys or name in expected_keys:
+                    additional[name] = _to_device(value, target_device)
+        required_keys = expected_keys or (handle.get("payload_keys", ()) if isinstance(handle, dict) else ())
+        missing = [name for name in required_keys if additional.get(name) is None]
+        if missing or (isinstance(handle, dict) and not required_keys and not isinstance(payload, dict)):
+            raise RuntimeError(f"Stage payload unavailable for {req.request_id}; missing keys: {missing}")
+
+    def _maybe_send_stage_payload(
+        self,
+        reqs: list[OmniDiffusionRequest],
+        outputs: list[DiffusionOutput],
+    ) -> None:
+        """Publish complete leader outputs; keep inline values if the put fails.
+
+        Pipelines must gather sharded outputs before calling this adapter.
+        Every rank publishes the same handle and drops only transferred keys.
+        """
+        if all(OmniDiffusionRequest.is_dummy_run_request_id(req.request_id) for req in reqs):
+            return
+        payload_keys = tuple(getattr(self.od_config, "stage_output_payload_keys", ()) or ())
+        if not payload_keys:
+            return
+        from_stage, to_stage = self._kv_transfer_manager.send_stages
+        if not from_stage or not to_stage:
+            logger.warning("Stage %s declares payload keys but has no outgoing edge", self.od_config.stage_id)
+            return
+        handles: dict[str, dict[str, Any]] = {}
+        if self.is_data_transfer_rank():
+            connector = self._stage_payload_connector()
+            if connector is not None:
+                for req, output in zip(reqs, outputs):
+                    if OmniDiffusionRequest.is_dummy_run_request_id(req.request_id):
+                        continue
+                    custom = getattr(output, "custom_output", None)
+                    if not isinstance(custom, dict):
+                        continue
+                    payload = {key: custom[key] for key in payload_keys if custom.get(key) is not None}
+                    if not payload:
+                        continue
+                    external_req_id = getattr(req, "external_req_id", None) or req.request_id
+                    _, _, put_key, _ = self._stage_payload_recv_spec(external_req_id, from_stage, to_stage)
+                    try:
+                        success, size, metadata = connector.put(from_stage, to_stage, put_key, payload)
+                    except Exception as exc:
+                        logger.warning("Stage payload put failed for %s: %s", put_key, exc)
+                        continue
+                    if not success:
+                        logger.warning("Stage payload %s was rejected; keeping the inline payload", put_key)
+                        continue
+                    handles[req.request_id] = {
+                        "key": put_key,
+                        "from_stage": from_stage,
+                        "to_stage": to_stage,
+                        "size_bytes": int(size),
+                        "metadata": metadata,
+                        "payload_keys": list(payload),
+                    }
+                    logger.debug("Stage payload put %s size=%s keys=%s", put_key, size, list(payload))
+        handles = self._broadcast_tp_payload_packet(handles) or {}
+        for req, output in zip(reqs, outputs):
+            handle = handles.get(req.request_id)
+            custom = getattr(output, "custom_output", None)
+            if handle is None or not isinstance(custom, dict):
+                continue
+            custom[self._STAGE_PAYLOAD_HANDLE_KEY] = handle
+            for name in handle["payload_keys"]:
+                custom.pop(name, None)

vllm_omni/diffusion/worker/utils.py

@@ -84,6 +84,7 @@ class StepRequestState:
     # Opaque model-owned preprocessing result. Scheduler-owned Diffusion KV
     # request state is removed before this Worker state is built.
     prepared_layout: Any | None = None
+    external_req_id: str | None = None
 
     # ── Encoded prompts (set once by prepare_encode) ──
     prompt_embeds: torch.Tensor | None = None

vllm_omni/distributed/omni_connectors/connectors/base.py

@@ -60,6 +60,21 @@ class OmniConnectorBase(ABC):
         """Clean up resources for a request."""
         pass
 
+    def get_with_deadline(
+        self,
+        from_stage: str,
+        to_stage: str,
+        get_key: str,
+        metadata: dict[str, Any] | None = None,
+        *,
+        deadline: float,
+    ) -> tuple[Any, int] | None:
+        """Receive with a monotonic deadline; blocking backends must override this."""
+        raise NotImplementedError(f"{type(self).__name__} does not support deadline-aware receive")
+
+    def abandon_get(self, get_key: str) -> None:
+        """Retire unresolved discovery attempts; never cancel an active DMA READ."""
+
     @abstractmethod
     def health(self) -> dict[str, Any]:
         """Return health status and metrics."""

vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py

@@ -91,6 +91,8 @@ class NixlConnector(OmniConnectorBase):
         self._pending: dict[str, _PendingPayload] = {}
         self._published: dict[str, dict[str, Any]] = {}
         self._state_lock = threading.RLock()
+        self._close_lock = threading.Lock()
+        self._close_thread: threading.Thread | None = None
         self._remote_agents: list[str] = []
         self._metrics: dict[str, int] = {
             "puts": 0,
@@ -159,6 +161,7 @@ class NixlConnector(OmniConnectorBase):
             self._lease_thread.start()
         self._transfer_wakeup = threading.Event()
         self._deferred_transfers: list[_DeferredTransfer] = []
+        self._abandoned_queries: list[tuple[str, str, str | None, str]] = []
         self._transfer_thread: threading.Thread | None = None
         if self._role != "sender":
             self._transfer_thread = threading.Thread(
@@ -325,6 +328,25 @@ class NixlConnector(OmniConnectorBase):
             logger.error("NixlConnector put failed for %s", put_key, exc_info=True)
             return False, 0, None
 
+    def get_with_deadline(self, from_stage, to_stage, get_key, metadata=None, *, deadline):
+        """Bound control-plane and DMA waits without freeing active READ buffers."""
+        self._req_local.deadline = deadline
+        try:
+            if time.monotonic() >= deadline:
+                return None
+            return self.get(from_stage, to_stage, get_key, metadata)
+        finally:
+            self._req_local.deadline = None
+
+    def abandon_get(self, get_key: str) -> None:
+        claims = getattr(self._req_local, "claims", {})
+        with self._state_lock:
+            for query_key, claim_id in list(claims.items()):
+                if query_key[1] == get_key:
+                    self._abandoned_queries.append((*query_key, claim_id))
+                    del claims[query_key]
+        self._transfer_wakeup.set()
+
     def get(
         self,
         from_stage: str,
@@ -347,6 +369,9 @@ class NixlConnector(OmniConnectorBase):
                 return None
 
             source_metadata = metadata
+            deadline = getattr(self._req_local, "deadline", None)
+            if deadline is not None and time.monotonic() >= deadline:
+                return None
 
             tensor_specs = metadata.get("tensor_specs")
             if not isinstance(tensor_specs, list):
@@ -448,7 +473,14 @@ class NixlConnector(OmniConnectorBase):
             for local_reg_descs in local_reg_descs_list:
                 self._safe_call(self._agent.deregister_memory, local_reg_descs)
             if source_metadata is not None:
-                self._notify_transfer_done(get_key, source_metadata)
+                if not self._notify_transfer_done(get_key, source_metadata):
+                    self._defer_transfer(
+                        _DeferredTransfer(
+                            tensors=[],
+                            source_key=get_key,
+                            source_metadata=source_metadata,
+                        )
+                    )
 
     def cleanup(self, request_id: str) -> None:
         pending = self._take_pending(request_id)
@@ -474,11 +506,27 @@ class NixlConnector(OmniConnectorBase):
             return self._pending.pop(request_id)
 
     def close(self) -> None:
+        with self._close_lock:
+            self._close_once()
+            if not self._closed and self._close_thread is None:
+                self._close_thread = threading.Thread(
+                    target=self._finish_close_in_background, name="nixl-close", daemon=True
+                )
+                self._close_thread.start()
+
+    def _finish_close_in_background(self) -> None:
+        retry = threading.Event()
+        while not retry.wait(max(self._poll_interval_s, 0.1)):
+            with self._close_lock:
+                self._close_once()
+                if self._closed:
+                    return
+
+    def _close_once(self) -> None:
         if self._closed:
             return
         self._closed = True
-        # Stop accepting puts/gets, but keep serving completion ACKs for outstanding
-        # remote READs. A later close() call can finish teardown once drained.
+        # Keep completion ACKs available until the automatic closer drains ownership.
         for request_id in list(self._pending):
             self.cleanup(request_id)
         with self._state_lock:
@@ -486,25 +534,33 @@ class NixlConnector(OmniConnectorBase):
                 _RETAINED_PRODUCERS.add(self)
                 self._closed = False
                 self._closing = True
-                logger.warning("NIXL close deferred: remote READ claims still own source allocations")
+                if self._close_thread is None:
+                    logger.warning("NIXL close deferred: remote READ claims still own source allocations")
                 return
         _RETAINED_PRODUCERS.discard(self)
         self._stop_event.set()
         self._lease_wakeup.set()
         self._transfer_wakeup.set()
-        if self._listener_thread is not None:
-            self._listener_thread.join(timeout=5.0)
-            self._listener_thread = None
-        if self._lease_thread is not None:
-            self._lease_thread.join(timeout=5.0)
-            self._lease_thread = None
-        if self._transfer_thread is not None:
-            self._transfer_thread.join(timeout=5.0)
-            self._transfer_thread = None
-        while self._deferred_transfers:
-            self._reap_deferred_transfers()
-            if self._deferred_transfers:
-                time.sleep(max(self._poll_interval_s, 0.01))
+        for thread_name in ("_listener_thread", "_lease_thread", "_transfer_thread"):
+            thread = getattr(self, thread_name)
+            if thread is None:
+                continue
+            if thread is not threading.current_thread():
+                thread.join(timeout=5.0)
+            if thread.is_alive():
+                _RETAINED_PRODUCERS.add(self)
+                self._closed = False
+                self._closing = True
+                return
+            setattr(self, thread_name, None)
+        self._reap_deferred_transfers()
+        if self._deferred_transfers or self._abandoned_queries:
+            _RETAINED_PRODUCERS.add(self)
+            self._closed = False
+            self._closing = True
+            if self._close_thread is None:
+                logger.warning("NIXL close deferred: local DMA resources retained for automatic cleanup")
+            return
         if self._zmq_ctx is not None:
             # destroy() rather than term(): REQ sockets live in thread-local
             # caches this thread cannot reach, and term() blocks until every
@@ -525,7 +581,7 @@ class NixlConnector(OmniConnectorBase):
 
     def health(self) -> dict[str, Any]:
         return {
-            "status": "unhealthy" if self._closed else "healthy",
+            "status": "unhealthy" if self._closed or getattr(self, "_closing", False) else "healthy",
             "pending_requests": len(self._pending),
             **self._metrics,
         }
@@ -593,19 +649,26 @@ class NixlConnector(OmniConnectorBase):
         return self._query_metadata_at(get_key, *endpoint)
 
     def _query_metadata_at(
-        self, get_key: str, host: str, port: int, *, generation: str | None = None
+        self, get_key: str, host: str, port: int, *, generation: str | None = None, claim_id: str | None = None
     ) -> dict[str, Any] | None:
         """Fetch transfer metadata for ``get_key`` from a producer's ROUTER socket.
 
         Each call performs one bounded query so a missing key cannot starve
         other requests in the shared receive loop. The caller retries later.
+        Retries after a lost reply reuse ownership until metadata is received.
         """
         zmq_addr = f"tcp://{host}:{port}"
+        claims = getattr(self._req_local, "claims", None)
+        if claims is None:
+            claims = {}
+            self._req_local.claims = claims
+        query_key = (zmq_addr, get_key, generation)
+        claim_id = claims.setdefault(query_key, claim_id or uuid.uuid4().hex)
         request = _GET_META_MSG + msgspec.msgpack.encode(
-            {"key": get_key, "generation": generation, "claim_id": uuid.uuid4().hex}
+            {"key": get_key, "generation": generation, "claim_id": claim_id}
         )
-        sock = self._get_req_socket(zmq_addr, self._metadata_query_timeout_ms)
         try:
+            sock = self._get_req_socket(zmq_addr, self._metadata_query_timeout_ms)
             sock.send(request)
             reply = sock.recv()
         except Exception:
@@ -613,10 +676,13 @@ class NixlConnector(OmniConnectorBase):
             logger.debug("NixlConnector handshake query to %s failed for %s", zmq_addr, get_key, exc_info=True)
             return None
         if reply == _META_NOT_FOUND:
+            claims.pop(query_key, None)
             return None
-        return msgspec.msgpack.decode(reply)
+        metadata = msgspec.msgpack.decode(reply)
+        claims.pop(query_key, None)
+        return metadata
 
-    def _notify_transfer_done(self, get_key: str, metadata: dict[str, Any]) -> None:
+    def _notify_transfer_done(self, get_key: str, metadata: dict[str, Any]) -> bool:
         """Tell the producer its buffer is drained so it can deregister now.
 
         Without this the producer would hold the registration until the lease
@@ -624,11 +690,11 @@ class NixlConnector(OmniConnectorBase):
         """
         endpoint = self._metadata_endpoint(metadata)
         if endpoint is None or self._zmq_ctx is None:
-            return
+            return True
         host, port = endpoint
         zmq_addr = f"tcp://{host}:{port}"
-        sock = self._get_req_socket(zmq_addr)
         try:
+            sock = self._get_req_socket(zmq_addr)
             sock.send(
                 _XFER_DONE_MSG
                 + msgspec.msgpack.encode(
@@ -636,9 +702,11 @@ class NixlConnector(OmniConnectorBase):
                 )
             )
             sock.recv()
+            return True
         except Exception:
             self._invalidate_req_socket(zmq_addr)
             logger.debug("NixlConnector failed to notify completion for %s", get_key, exc_info=True)
+            return False
 
     def _handshake_listener_loop(self) -> None:
         router = self._zmq_ctx.socket(zmq.ROUTER)
@@ -686,8 +754,7 @@ class NixlConnector(OmniConnectorBase):
                     or not claim
                     or request.get("generation") not in (None, pending.generation)
                     or (not pending.claims and time.monotonic() >= pending.deadline)
-                    or self._closed
-                    or getattr(self, "_closing", False)
+                    or ((self._closed or getattr(self, "_closing", False)) and claim not in pending.claims)
                 ):
                     return _META_NOT_FOUND
                 # Publish ownership before descriptors can leave this lock.
@@ -723,6 +790,12 @@ class NixlConnector(OmniConnectorBase):
             sock.connect(zmq_addr)
             cache[zmq_addr] = sock
         timeout_ms = self._handshake_timeout_ms if timeout_ms is None else timeout_ms
+        deadline = getattr(self._req_local, "deadline", None)
+        if deadline is not None:
+            remaining_ms = int((deadline - time.monotonic()) * 1000 / 2)
+            if remaining_ms <= 0:
+                raise TimeoutError("NIXL receive deadline expired")
+            timeout_ms = min(timeout_ms, remaining_ms)
         sock.setsockopt(zmq.SNDTIMEO, timeout_ms)
         sock.setsockopt(zmq.RCVTIMEO, timeout_ms)
         return sock
@@ -751,6 +824,9 @@ class NixlConnector(OmniConnectorBase):
 
     def _wait_for_transfer(self, handle: int, request_id: str) -> None:
         deadline = time.monotonic() + self._transfer_timeout_s
+        receive_deadline = getattr(self._req_local, "deadline", None)
+        if receive_deadline is not None:
+            deadline = min(deadline, receive_deadline)
         while True:
             state = self._agent.check_xfer_state(handle)
             if state == "DONE":
@@ -759,7 +835,7 @@ class NixlConnector(OmniConnectorBase):
                 raise RuntimeError(f"NIXL transfer for {request_id} failed with state={state}")
             if time.monotonic() >= deadline:
                 raise TimeoutError(f"NIXL transfer for {request_id} timed out")
-            time.sleep(self._poll_interval_s)
+            time.sleep(min(self._poll_interval_s, max(0.0, deadline - time.monotonic())))
 
     def _transfers_may_be_active(self, handles: list[Any]) -> bool:
         try:
@@ -815,6 +891,22 @@ class NixlConnector(OmniConnectorBase):
             self._reap_deferred_transfers()
 
     def _reap_deferred_transfers(self) -> None:
+        with self._state_lock:
+            queries = list(self._abandoned_queries)
+        for query in queries:
+            address, key, generation, claim_id = query
+            endpoint = address.removeprefix("tcp://").rsplit(":", 1)
+            metadata = self._query_metadata_at(
+                key, endpoint[0], int(endpoint[1]), generation=generation, claim_id=claim_id
+            )
+            query_key = (address, key, generation)
+            unresolved = query_key in getattr(self._req_local, "claims", {})
+            if metadata is None and unresolved:
+                continue
+            if metadata is not None and not self._notify_transfer_done(key, metadata):
+                self._defer_transfer(_DeferredTransfer(tensors=[], source_key=key, source_metadata=metadata))
+            with self._state_lock:
+                self._abandoned_queries.remove(query)
         with self._state_lock:
             transfers = list(self._deferred_transfers)
         for transfer in transfers:
@@ -834,7 +926,8 @@ class NixlConnector(OmniConnectorBase):
             ):
                 transfer.tensors.clear()
                 if transfer.source_metadata is not None:
-                    self._notify_transfer_done(transfer.source_key, transfer.source_metadata)
+                    if not self._notify_transfer_done(transfer.source_key, transfer.source_metadata):
+                        continue
                     transfer.source_metadata = None
                 with self._state_lock:
                     if transfer in self._deferred_transfers:

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -3,6 +3,7 @@
 
 import fcntl
 import os
+import time
 from multiprocessing import shared_memory as shm_pkg
 from typing import Any
 
@@ -24,6 +25,11 @@ class SharedMemoryConnector(OmniConnectorBase):
     the connector silently falls back to key-based lookup.
     """
 
+    def get_with_deadline(self, from_stage, to_stage, get_key, metadata=None, *, deadline):
+        if time.monotonic() >= deadline:
+            return None
+        return self.get(from_stage, to_stage, get_key, metadata)
+
     def __init__(self, config: dict[str, Any]):
         self.config = config
         self.stage_id = config.get("stage_id", -1)
@@ -68,13 +74,15 @@ class SharedMemoryConnector(OmniConnectorBase):
         deserialized = False
         try:
             with open(lock_file, "rb+") as lockf:
-                fcntl.flock(lockf, fcntl.LOCK_EX)
+                fcntl.flock(lockf, fcntl.LOCK_EX | fcntl.LOCK_NB)
                 data_bytes = shm_read_bytes(shm_handle)
                 fcntl.flock(lockf, fcntl.LOCK_UN)
             obj = self.deserialize_obj(data_bytes)
             result = (obj, int(shm_handle.get("size", 0)))
             deserialized = True
             return result
+        except BlockingIOError:
+            return None
         except Exception as e:
             logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
             return None

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -1376,6 +1376,16 @@ class OmniKVTransferManager:
                 logger.exception("Failed to shut down KV prefetch executor")
             self._prefetch_executor = None
 
+    def close(self) -> None:
+        """Stop prefetch work before closing the connector owned by this manager."""
+        try:
+            self.shutdown_prefetch()
+        finally:
+            connector = self._connector
+            self._connector = False
+            if connector:
+                connector.close()
+
     @staticmethod
     def _record_stream_for_prefetched(data: dict[str, Any]) -> None:
         """``record_stream(current_stream)`` on GPU tensors in *data*.

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import importlib
 import inspect
+import time
 from collections import deque
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, cast
@@ -192,8 +193,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         metadata: dict[str, Any] | None = None,
     ) -> Any:
         """Receive one ordinary non-KV stage payload on the local leader rank only."""
-        tp_group = self._get_local_tp_group()
-        if tp_group is None or getattr(tp_group, "world_size", 1) <= 1:
+        if not self._stage_payload_broadcast_groups():
             if metadata is None:
                 return connector.get(from_stage, to_stage, connector_get_key)
             return connector.get(from_stage, to_stage, connector_get_key, metadata)
@@ -243,13 +243,81 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             return dict(payload)
         return payload
 
-    def _broadcast_tp_payload_packet(self, packet: Any) -> Any:
-        """Broadcast one ordinary payload packet from TP rank 0 when TP is active."""
+    def _stage_payload_broadcast_groups(self) -> tuple[Any, ...]:
+        """Return ordered orthogonal payload groups; runners may supply TP x SP."""
         tp_group = self._get_local_tp_group()
         if tp_group is None or getattr(tp_group, "world_size", 1) <= 1:
-            return packet
-        leader_packet = packet if self.is_data_transfer_rank() else None
-        return tp_group.broadcast_object(leader_packet, src=0)
+            return ()
+        return (tp_group,)
+
+    def _broadcast_tp_payload_packet(self, packet: Any, *, tensor_payload: bool = False) -> Any:
+        """Fan out in group order, using each group's local leader as source."""
+        for group in self._stage_payload_broadcast_groups():
+            leader_packet = packet if group.rank_in_group == 0 else None
+            if tensor_payload:
+                delivered = group.broadcast_object(
+                    isinstance(packet, dict) if group.rank_in_group == 0 else None, src=0
+                )
+                if delivered:
+                    packet = group.broadcast_tensor_dict(leader_packet, src=0)
+            else:
+                packet = group.broadcast_object(leader_packet, src=0)
+        return packet
+
+    def _stage_payload_connector(self) -> Any | None:
+        try:
+            if getattr(self, "_synchronous_payload_transport", False):
+                return self._kv_transfer_manager.connector
+            return self._omni_connector
+        except Exception as exc:
+            logger.warning("Stage payload connector unavailable: %s", exc)
+            return None
+
+    def recv_stage_payload(
+        self,
+        external_req_id: str,
+        from_stage: str,
+        to_stage: str,
+        *,
+        sender_info: Any = None,
+        handle: dict[str, Any] | None = None,
+        retry_seconds: float = 2.0,
+    ) -> Any:
+        """Synchronously receive one payload, then fan out on the model thread.
+
+        The same monotonic deadline covers discovery and backend transfer waits.
+        This path does not register work with the asynchronous poller.
+        """
+        from_stage, to_stage, get_key, metadata = self._stage_payload_recv_spec(
+            external_req_id, from_stage, to_stage, sender_info=sender_info, handle=handle
+        )
+        if not get_key:
+            return None
+        result = None
+        if self.is_data_transfer_rank():
+            connector = self._stage_payload_connector()
+            if connector is not None:
+                deadline = time.monotonic() + retry_seconds
+                try:
+                    while True:
+                        result = connector.get_with_deadline(from_stage, to_stage, get_key, metadata, deadline=deadline)
+                        if result is not None:
+                            break
+                        remaining = deadline - time.monotonic()
+                        if remaining <= 0:
+                            break
+                        time.sleep(min(0.05, remaining))
+                except Exception as exc:
+                    logger.warning("Stage payload get failed for %s: %s", get_key, exc)
+                finally:
+                    try:
+                        connector.abandon_get(get_key)
+                    except Exception as exc:
+                        logger.warning("Stage payload claim retirement failed for %s: %s", get_key, exc)
+        payload = self._broadcast_tp_payload_packet(result[0] if result else None, tensor_payload=True)
+        if payload is None and self.is_data_transfer_rank():
+            logger.warning("Stage payload %s was not delivered; caller must validate inline fallback", get_key)
+        return payload
 
     def _apply_staged_payloads_locked(self, staged_payloads: dict[str, Any]) -> None:
         for req_id, payload in staged_payloads.items():
@@ -980,6 +1048,31 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
     #  Chunk-level poll / send  (ported from OmniChunkTransferAdapter)
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _stage_payload_recv_spec(
+        external_req_id: str,
+        from_stage: str,
+        to_stage: str,
+        chunk_id: int = 0,
+        sender_info: Any = None,
+        handle: dict[str, Any] | None = None,
+    ) -> tuple[str, str, str | None, dict[str, Any] | None]:
+        """Resolve a connector-independent request key and transfer metadata."""
+        metadata = None
+        if isinstance(sender_info, dict):
+            host = sender_info.get("host")
+            port = sender_info.get("zmq_port")
+            if host and port:
+                metadata = {"source_host": str(host), "source_port": int(port)}
+        if isinstance(handle, dict):
+            return (
+                str(handle.get("from_stage", from_stage)),
+                str(handle.get("to_stage", to_stage)),
+                handle.get("key"),
+                handle.get("metadata") or metadata,
+            )
+        return str(from_stage), str(to_stage), f"{external_req_id}_{from_stage}_{chunk_id}", metadata
+
     def _poll_single_request(self, req_id: str, *, publish_ready: bool = True) -> bool:
         """Poll connector for one chunk of a request (non-blocking)."""
         connector = self._omni_connector
@@ -1002,29 +1095,28 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         target_stage_id = self._stage_id - 1
         chunk_id = self._get_req_chunk[req_id]
         external_req_id = self._request_ids_mapping.get(req_id, req_id)
-        connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}"
         request = self._pending_load_reqs.get(req_id)
-        sender_info = getattr(request, "payload_sender_info", None)
-        metadata = None
-        if isinstance(sender_info, dict):
-            host = sender_info.get("host")
-            port = sender_info.get("zmq_port")
-            if host and port:
-                metadata = {"source_host": str(host), "source_port": int(port)}
+        from_stage, to_stage, connector_get_key, metadata = self._stage_payload_recv_spec(
+            external_req_id,
+            str(target_stage_id),
+            str(self._stage_id),
+            chunk_id,
+            getattr(request, "payload_sender_info", None),
+        )
 
         if self._async_chunk:
             result = self._recv_async_chunk_result(
                 connector,
-                str(target_stage_id),
-                str(self._stage_id),
+                from_stage,
+                to_stage,
                 connector_get_key,
                 metadata,
             )
         else:
             result = self._recv_full_payload_result(
                 connector,
-                str(target_stage_id),
-                str(self._stage_id),
+                from_stage,
+                to_stage,
                 connector_get_key,
                 metadata,
             )

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py

@@ -169,15 +169,24 @@ class _OmniConnectorRuntimeMixin:
         self,
         model_config: OmniModelConfig,
         kv_transfer_manager: OmniKVTransferManager | None = None,
+        *,
+        synchronous: bool = False,
     ) -> None:
         """Initialize connectors and background threads.
 
         Args:
             model_config: Stage-level model config with connector settings.
             kv_transfer_manager: Existing KV transfer manager to delegate to.
+            synchronous: Borrow the manager's lazy connector without background
+                threads or changes to its KV callbacks. The manager retains ownership.
         """
+        self._synchronous_payload_transport = synchronous
+        if synchronous and kv_transfer_manager is None:
+            raise ValueError("Synchronous payload transport requires a KV transfer manager")
         self._omni_connector: OmniConnectorBase | None = (
-            self._create_connector(model_config) if _should_create_payload_connector(model_config) else None
+            self._create_connector(model_config)
+            if not synchronous and _should_create_payload_connector(model_config)
+            else None
         )
         self._kv_transfer_manager = kv_transfer_manager
 
@@ -188,7 +197,9 @@ class _OmniConnectorRuntimeMixin:
             stage_id = int(stage_id)
         self._stage_id: int = stage_id if isinstance(stage_id, int) else 0
 
-        self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config)
+        self._custom_process_func_path, self._custom_process_func = (
+            (None, None) if synchronous else self._load_custom_func(model_config)
+        )
         self._custom_process_batch_func = self._load_custom_batch_func(self._custom_process_func)
         self._custom_process_payload_kwarg = self._connector_payload_kwarg(self._custom_process_func)
         self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg()
@@ -226,7 +237,7 @@ class _OmniConnectorRuntimeMixin:
         self._from_tp: int = rank_cfg["from_tp"]
         self._to_tp: int = rank_cfg["to_tp"]
         self._local_rank: int = rank_cfg["local_rank"]
-        if self._kv_transfer_manager is not None:
+        if self._kv_transfer_manager is not None and not synchronous:
             self._kv_transfer_manager.kv_send_key_builder = self.get_rank_aware_kv_send_keys
             self._kv_transfer_manager.kv_recv_key_builder = self.get_rank_aware_kv_keys
             self._kv_transfer_manager.kv_payload_merger = self._merge_rank_sharded_kv_payloads
@@ -329,6 +340,8 @@ class _OmniConnectorRuntimeMixin:
 
     def shutdown_omni_connectors(self) -> None:
         """Stop background threads and release connector resources."""
+        if getattr(self, "_synchronous_payload_transport", False):
+            return
         self._stop_event.set()
         if self._recv_thread is not None:
             self._recv_thread.join(timeout=5)
@@ -1071,15 +1084,14 @@ class _OmniConnectorRuntimeMixin:
     def is_data_transfer_rank(self) -> bool:
         """Whether this rank should participate in data (non-KV) transfer.
 
-        Ordinary stage payloads are TP-identical, so exactly one TP rank
-        should talk to the connector. When TP is initialized, use TP rank 0
-        so the connector leader matches TP-local broadcast source rank.
+        Ordinary stage payloads are identical across the configured payload
+        groups. Only the rank leading every group talks to the connector.
         Otherwise fall back to LOCAL_RANK==0 for the single-rank case.
         """
-        tp_group = self._get_local_tp_group()
-        if tp_group is not None and getattr(tp_group, "world_size", 1) > 1:
-            return getattr(tp_group, "rank_in_group", 0) == 0
-        return self._local_rank == 0
+        groups = self._stage_payload_broadcast_groups()
+        if groups:
+            return all(getattr(group, "rank_in_group", 0) == 0 for group in groups)
+        return getattr(self, "_local_rank", 0) == 0
 
     def get_kv_connector_key(
         self,

vllm_omni/engine/orchestrator.py

@@ -2270,20 +2270,20 @@ class OrchestratorBase:
             else:
                 diffusion_prompt = req_state.prompt
 
+            submit_kwargs = self._diffusion_submit_kwargs(req_id, src_stage_id, next_client, req_state, output)
+            payload_sender_info = self._build_payload_sender_info(src_stage_id, request_id=req_id)
+            if payload_sender_info is not None:
+                submit_kwargs["payload_sender_info"] = payload_sender_info
             if already_submitted:
-                replica_id = await next_pool.submit_update(req_id, req_state, diffusion_prompt)
+                replica_id = await next_pool.submit_update(
+                    req_id, req_state, diffusion_prompt, submit_kwargs=submit_kwargs
+                )
             else:
                 replica_id = await next_pool.submit_initial(
                     req_id,
                     req_state,
                     diffusion_prompt,
-                    submit_kwargs=self._diffusion_submit_kwargs(
-                        req_id,
-                        src_stage_id,
-                        next_client,
-                        req_state,
-                        output,
-                    ),
+                    submit_kwargs=submit_kwargs,
                     params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
                 )
             self._on_stage_submitted(

vllm_omni/engine/stage_client.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Shared stage-client typing for vLLM-Omni runtime surfaces."""
 
 from __future__ import annotations

vllm_omni/engine/stage_engine_core_client.py

@@ -152,11 +152,13 @@ class StageEngineCoreClientBase(StageClientBase):
 
         self.engine_outputs: Any = None
         self.client_addresses = dict(client_addresses or {})
+        self.vllm_config = vllm_config
         self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None)
         self._stage_hf_config = getattr(getattr(vllm_config, "model_config", None), "hf_config", None)
         self._kv_sender_host = self._resolve_contact_host()
         self._kv_sender_info: dict[str, Any] | None = None
         self._kv_sender_initialized = False
+        self._payload_sender_info: dict[str, Any] | None = None
 
         client_name = self.__class__.__name__
         logger.info(
@@ -213,8 +215,8 @@ class StageEngineCoreClientBase(StageClientBase):
                 )
             raise
 
-        self._payload_sender_info = self._build_payload_sender_info()
         self._initialize_kv_sender_endpoint()
+        self._payload_sender_info = self._build_payload_sender_info()
 
         logger.info(
             "[%s] stage-%s [rep-%s] EngineCore running",
@@ -301,16 +303,34 @@ class StageEngineCoreClientBase(StageClientBase):
     def _build_payload_sender_info(self) -> dict[str, Any] | None:
         model_config = getattr(self.vllm_config, "model_config", None)
         connector_config = getattr(model_config, "stage_connector_config", None)
-        if not isinstance(connector_config, dict):
-            return None
-        extra = connector_config.get("extra")
-        if not isinstance(extra, dict):
-            return None
-        if isinstance(extra.get("outgoing"), dict):
-            extra = extra["outgoing"]
-        elif extra.get("role") != "sender":
-            return None
-        base_port = extra.get("zmq_port", 50051)
+        extra = connector_config.get("extra") if isinstance(connector_config, dict) else None
+        outgoing = extra.get("outgoing") if isinstance(extra, dict) else None
+
+        # An intermediate stage's outgoing listener is independent of its
+        # incoming KV endpoint. Never advertise the upstream sender here.
+        sender_info = getattr(self, "_kv_sender_info", None)
+        if not isinstance(outgoing, dict) and isinstance(sender_info, dict):
+            sender_host = sender_info.get("host")
+            sender_port = sender_info.get("zmq_port")
+            if sender_host is not None and sender_port is not None:
+                return {
+                    "host": str(sender_host),
+                    "zmq_port": int(sender_port) - KV_TRANSFER_PORT_OFFSET,
+                }
+
+        kv_connector_config = self._get_kv_connector_config()
+        if isinstance(outgoing, dict):
+            extra = outgoing
+            base_port = extra.get("zmq_port", 50051)
+        else:
+            if not isinstance(extra, dict) or extra.get("role") != "sender":
+                extra = kv_connector_config
+            if not isinstance(extra, dict) or extra.get("role") != "sender":
+                return None
+            base_port = extra.get("zmq_port", 50051)
+            kv_port = kv_connector_config.get("zmq_port") if isinstance(kv_connector_config, dict) else None
+            if kv_port is not None:
+                base_port = int(os.path.expandvars(str(kv_port))) - KV_TRANSFER_PORT_OFFSET
         if base_port is None:
             return None
         sender_host = self._resolve_sender_host_from_config(extra)

vllm_omni/engine/stage_pool.py

@@ -998,6 +998,9 @@ class StagePool:
                     "Diffusion list-prompt batch requests are no longer supported. "
                     "Submit multiple independent requests to use scheduler batching."
                 )
+            payload_sender_info = getattr(request, "payload_sender_info", None)
+            if payload_sender_info is not None:
+                submit_kwargs.setdefault("payload_sender_info", payload_sender_info)
             replica_id = await self._pick_or_select(
                 request_id,
                 affinity_request_id=affinity_request_id,
@@ -1050,8 +1053,10 @@ class StagePool:
         request: Any,
         *,
         prompt_text: Any = None,
+        submit_kwargs: dict[str, Any] | None = None,
     ) -> int:
         """Submit a streaming update to an already admitted request."""
+        submit_kwargs = submit_kwargs or {}
         params = req_state.sampling_params_list[self.stage_id]
         if self.stage_type == "diffusion":
             params = OmniDiffusionSamplingParams.from_params(params)
@@ -1069,7 +1074,7 @@ class StagePool:
                     "Diffusion list-prompt batch requests are no longer supported. "
                     "Submit multiple independent requests to use scheduler batching."
                 )
-            await self._diffusion_client(replica_id).add_request_async(request_id, request, params)
+            await self._diffusion_client(replica_id).add_request_async(request_id, request, params, **submit_kwargs)
         else:
             # Refresh the shared output-processor state before yielding to the
             # stage client so streaming segments are merged against the latest
@@ -1082,7 +1087,7 @@ class StagePool:
                     request_index=0,
                     queue=None,
                 )
-                await self._llm_client(replica_id).add_request_async(request)
+                await self._llm_client(replica_id).add_request_async(request, **submit_kwargs)
             except Exception:
                 rollback = getattr(self.output_processor, "remove_request", None)
                 if callable(rollback):

vllm_omni/model_executor/models/minimax_h3/conditioning.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""MiniMax H3 text-conditioning contract."""
+"""MiniMax H3 validated text and unified encoder-conditioning contracts."""
 
 from __future__ import annotations
 

vllm_omni/model_executor/models/minimax_h3/encoder_processing.py

@@ -395,6 +395,8 @@ def _effective_audio_inputs(
     max_standalone_seconds: float,
 ) -> list[tuple[torch.Tensor, int]]:
     """Return the waveforms exactly as the audio WVAE will consume them."""
+    if not math.isfinite(max_standalone_seconds) or max_standalone_seconds <= 0:
+        raise ValueError("max_standalone_seconds must be positive and finite")
     audio_inputs = [item for item in video_audios if item is not None]
     audio_inputs.extend(
         (waveform[..., : int(round(max_standalone_seconds * sample_rate))], sample_rate)
@@ -683,6 +685,16 @@ def encode_media(
     Every visual participant must call this function with the same references.
     Component scopes allow the diffusion runner to stage one codec at a time.
     """
+    # Validate before the non-leader return so all participants reject bad input.
+    audio_inputs = _effective_audio_inputs(
+        media.video_audios,
+        media.audios,
+        max_standalone_seconds=float(media.num_frames) / MINIMAX_H3_FPS,
+    )
+    embedded_audio_count = sum(item is not None for item in media.video_audios)
+    validate_reference_audio_waveforms([item for item in media.video_audios if item is not None])
+    if media.audio_mode != "lock_source":
+        validate_reference_audio_waveforms(media.audios)
     visual_rows: list[torch.Tensor] = []
     visual_shapes: list[tuple[int, int, int]] = []
     video_edit_clean_rows: torch.Tensor | None = None
@@ -720,12 +732,6 @@ def encode_media(
     audio_lengths: list[int] = []
     audio_edit_clean_rows: torch.Tensor | None = None
     audio_edit_source_t = 0
-    embedded_audio_count = sum(item is not None for item in media.video_audios)
-    audio_inputs = _effective_audio_inputs(
-        media.video_audios,
-        media.audios,
-        max_standalone_seconds=float(media.num_frames) / MINIMAX_H3_FPS,
-    )
     if audio_inputs or media.audio_edit is not None:
         if audio_vae is None:
             raise RuntimeError("MiniMax H3 audio WVAE is not resident on the encoder leader")

vllm_omni/model_executor/models/minimax_h3/pipeline.py

@@ -28,6 +28,7 @@ MINIMAX_H3_PIPELINE = PipelineConfig(
             model_arch="MiniMaxH3Encoder",
             engine_output_type="latent",
             prompt_transform_func=f"{_PROCESSOR}.prepare_encoder_prompt",
+            custom_process_next_stage_input_func=f"{_PROCESSOR}.encoder2diffusion_full_payload",
             sampling_constraints={
                 "max_tokens": 1,
                 "temperature": 0.0,
@@ -45,6 +46,7 @@ MINIMAX_H3_PIPELINE = PipelineConfig(
             requires_multimodal_data=False,
             model_arch="MiniMaxH3Pipeline",
             custom_process_input_func=f"{_PROCESSOR}.encoder2diffusion",
+            stage_input_payload_keys=("encoder_output",),
             omni_kv_config={"need_recv_cache": False},
             model_path_resolver=f"{_DIFFUSION_PIPELINE}.resolve_minimax_h3_diffusion_model_path",
             inline_diffusion=True,

vllm_omni/model_executor/stage_input_processors/minimax_h3.py

@@ -8,7 +8,7 @@ import copy
 from collections.abc import Mapping, Sequence
 from typing import Any
 
-from vllm_omni.data_entry_keys import unflatten_payload
+from vllm_omni.data_entry_keys import OmniPayloadStruct, to_dict, unflatten_payload
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_executor.models.minimax_h3.conditioning import (
     MINIMAX_H3_CONDITION_LABELS_KEY,
@@ -92,6 +92,34 @@ def _global_request_id(prompt: Mapping[str, Any]) -> str | None:
     return str(value) if value is not None else None
 
 
+def _encoder_conditioning(payload: Any) -> MiniMaxH3EncoderConditioning:
+    if isinstance(payload, OmniPayloadStruct):
+        payload = to_dict(payload)
+    if not isinstance(payload, Mapping):
+        raise RuntimeError("MiniMax H3 encoder returned no conditioning payload")
+    try:
+        return MiniMaxH3EncoderConditioning.from_omni_payload(unflatten_payload(dict(payload)))
+    except ValueError as exc:
+        raise RuntimeError(str(exc)) from exc
+
+
+def encoder2diffusion_full_payload(
+    *,
+    pooling_output: Any = None,
+    **kwargs: Any,
+) -> dict[str, Any] | None:
+    """Pack all three encoder components for direct worker-to-worker transfer.
+
+    Returning the diffusion-ready structure here keeps the DiT worker free of
+    any H3-specific unpacking: the generic receive path merges this dict into
+    the request's ``additional_information``.
+    """
+    del kwargs
+    if pooling_output is None:
+        return None
+    return {"encoder_output": _encoder_conditioning(pooling_output).to_omni_payload()}
+
+
 def encoder2diffusion(
     source_outputs: list[Any],
     prompt: Any = None,
@@ -126,22 +154,25 @@ def encoder2diffusion(
         output_count = len(outputs) if isinstance(outputs, list) else 0
         raise RuntimeError(f"MiniMax H3 encoder must return exactly one completion, got {output_count}")
     payload = getattr(outputs[0], "multimodal_output", None)
-    if not isinstance(payload, Mapping):
-        raise RuntimeError("MiniMax H3 encoder returned no conditioning payload")
-    try:
-        conditioning = MiniMaxH3EncoderConditioning.from_omni_payload(unflatten_payload(dict(payload)))
-    except ValueError as exc:
-        raise RuntimeError(str(exc)) from exc
+    # Successful connector sends omit the inline payload. The diffusion runner
+    # merges encoder_output before forward; original media still needs cleanup.
+    conditioning = _encoder_conditioning(payload) if payload is not None else None
 
     additional_information = dict(diffusion_prompt.get("additional_information") or {})
-    additional_information.pop("hidden_states", None)
+    hidden_states = dict(additional_information.get("hidden_states") or {})
+    hidden_states.pop("layers", None)
+    if hidden_states:
+        additional_information["hidden_states"] = hidden_states
+    else:
+        additional_information.pop("hidden_states", None)
     meta = dict(additional_information.get("meta") or {})
     meta.pop(MINIMAX_H3_ENCODER_REQUEST_KEY, None)
     if meta:
         additional_information["meta"] = meta
     else:
         additional_information.pop("meta", None)
-    additional_information["encoder_output"] = conditioning.to_omni_payload()
+    if conditioning is not None:
+        additional_information["encoder_output"] = conditioning.to_omni_payload()
     diffusion_prompt["additional_information"] = additional_information
     diffusion_prompt["multi_modal_data"] = None
     diffusion_prompt.pop("model_intermediate_buffer", None)