diff --git a/docs/design/feature/omni_connectors/nixl_connector.md b/docs/design/feature/omni_connectors/nixl_connector.md index 7dfd73946..4cd58e7eb 100644 --- a/docs/design/feature/omni_connectors/nixl_connector.md +++ b/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 diff --git a/docs/design/feature/omni_connectors/stage_payload.md b/docs/design/feature/omni_connectors/stage_payload.md new file mode 100644 index 000000000..7cc429400 --- /dev/null +++ b/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. diff --git a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md index 54ea616af..f9f6d685b 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md +++ b/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. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index fdba7e90a..f9a5b4ea8 100644 --- a/tests/config/test_config_factory.py +++ b/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"] diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index d589fde2b..1c8d4e487 100644 --- a/tests/config/test_omni_config.py +++ b/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", diff --git a/tests/diffusion/distributed/test_parallel_state_sp_groups.py b/tests/diffusion/distributed/test_parallel_state_sp_groups.py index 0febb1925..b42e97ba2 100644 --- a/tests/diffusion/distributed/test_parallel_state_sp_groups.py +++ b/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)) diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py b/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py index b73804433..90758b449 100644 --- a/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py +++ b/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 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 7e4461247..13f2f5a68 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -27,6 +27,13 @@ def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ return result +def _turbo_spec(filename: str) -> TurboSpec: + """Build the spec a loaded artifact of this name would carry.""" + spec = parse_turbo_filename(filename) + assert spec is not None + return spec + + def _encoder_payload( *, task: str = "t2va", @@ -57,6 +64,282 @@ def _encoder_prompt(text: str = "test") -> dict[str, Any]: } +@pytest.mark.parametrize("field", ["hidden_states", "token_tags"]) +@pytest.mark.parametrize("invalid", ["dtype", "strides", "sparse"]) +def test_unified_encoder_preserves_main_text_validation(field, invalid): + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderConditioning + + payload = _encoder_payload() + target, key = ( + (payload["hidden_states"], "output") if field == "hidden_states" else (payload["meta"], "token_role_ids") + ) + shape = (2, 5120) if field == "hidden_states" else (2,) + dtype = torch.bfloat16 if field == "hidden_states" else torch.int64 + value = torch.ones(shape, dtype=dtype) + if invalid == "dtype": + value = value.float() + elif invalid == "strides": + value = torch.ones((4, *shape[1:]), dtype=dtype)[::2] + else: + value = value.to_sparse() + payload["hidden_states"]["output"] = torch.ones(2, 5120, dtype=torch.bfloat16) + payload["meta"]["token_role_ids"] = torch.ones(2, dtype=torch.int64) + target[key] = value + with pytest.raises(ValueError, match="minimax_h3.text_conditioning/v1"): + MiniMaxH3EncoderConditioning.from_omni_payload(payload) + + +@pytest.mark.parametrize("column_tags", [False, True]) +def test_unified_encoder_roundtrips_all_conditioning(column_tags): + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderConditioning + + conditioning = MiniMaxH3EncoderConditioning( + hidden_states=torch.ones(2, 5120, dtype=torch.bfloat16), + token_tags=torch.tensor([0, 1]), + task="ref2va", + height=768, + width=1344, + num_frames=124, + latent_t=37, + audio_t=207, + visual_condition=torch.arange(96, dtype=torch.float32).reshape(1, 96), + visual_condition_shapes=((1, 2, 2),), + audio_condition=torch.ones(160, 32), + audio_condition_lengths=(80,), + ref_blocks=({"kind": "video_audio", "ref_audio_t": 80, "latent_t": 1, "latent_h": 2, "latent_w": 2},), + keyframe_frame_indices=(0, -1), + ) + payload = conditioning.to_omni_payload() + if column_tags: + payload["meta"]["token_role_ids"] = payload["meta"]["token_role_ids"].unsqueeze(-1) + restored = MiniMaxH3EncoderConditioning.from_omni_payload(payload) + for name in ("hidden_states", "token_tags", "visual_condition", "audio_condition"): + torch.testing.assert_close(getattr(restored, name), getattr(conditioning, name)) + for name in ("visual_condition_shapes", "audio_condition_lengths", "ref_blocks", "keyframe_frame_indices"): + assert getattr(restored, name) == getattr(conditioning, name) + + +@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"]) +@pytest.mark.parametrize("text_key", [None, "encoder_output", "text_encoder_output"]) +def test_local_encoder_modes_preserve_validated_text_and_unified_handoff(task, text_key): + from PIL import Image + + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderConditioning + + pipeline = _distilled_pipeline([], {"fl2va": None, "ref2va": None}) + pipeline.load_text_encoder = text_key is None + pipeline.load_vae_encoder = True + text = {"hidden_states": torch.ones(2, 5120, dtype=torch.bfloat16), "token_tags": torch.tensor([0, 1])} + pipeline.encode_prompt = Mock(return_value=(text["hidden_states"], text["token_tags"])) + pipeline.video_vae = Mock() + pipeline.video_vae.is_distributed_enabled.return_value = False + pipeline.video_vae.encode_image.side_effect = lambda image: torch.ones( + (image.height // 32) * (image.width // 32), 96 + ) + pipeline.audio_vae = Mock() + prompt = {"prompt": "A person waves."} + if task != "t2va": + prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))} + if text_key is not None: + prompt["additional_information"] = {text_key: text} + sampling = OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=96, + num_inference_steps=2, + extra_args={"task": task, "aspect_ratio": "1:1"}, + ) + conditioning, window_embeddings = pipeline._prepare_local_conditioning( + prompt, sampling, require_external_text=text_key is not None + ) + assert window_embeddings is None + assert pipeline.encode_prompt.call_count == int(text_key is None) + assert pipeline.video_vae.encode_image.call_count == int(task != "t2va") + restored = MiniMaxH3EncoderConditioning.from_omni_payload(conditioning.to_omni_payload()) + pipeline.load_text_encoder = pipeline.load_vae_encoder = False + context = pipeline._prepare_request_inputs( + {"additional_information": {"encoder_output": restored.to_omni_payload()}}, sampling + ) + assert context["task"] == task + assert context["num_steps"] == 2 + torch.testing.assert_close(context["text_embeddings"], text["hidden_states"]) + assert pipeline.video_vae.encode_image.call_count == int(task != "t2va") + + +def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances: list["FakeEncoder"] = [] + + def __init__(self, **kwargs): + self.pushes = [] + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames, *, on_consumed=None): + self.pushes.append(np.array(frames, copy=True)) + if on_consumed is not None: + on_consumed() + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + chunk_value_range = (0.0, 1.0) + + def decode_with_chunks(self, latent, *, on_chunk): + for value in (0.0, 0.25, 0.5): + on_chunk(torch.full((1, 3, 3, 2, 2), value)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr( + "vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", + FakeEncoder, + ) + + output = pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + batch_frames=4, + ) + + assert output == b"mp4" + encoder = FakeEncoder.instances[-1] + assert [frames.shape[0] for frames in encoder.pushes] == [6, 3] + assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8)) + assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8)) + assert encoder.pushes[1].shape == (3, 2, 2, 3) + + +def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch): + """A client's encoder options must survive the worker-side pre-encode path.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances: list["FakeEncoder"] = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames, *, on_consumed=None): + del frames + if on_consumed is not None: + on_consumed() + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + chunk_value_range = (0.0, 1.0) + + def decode_with_chunks(self, latent, *, on_chunk): + on_chunk(torch.zeros(1, 3, 1, 2, 2)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr("vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", FakeEncoder) + + pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + video_codec_options={"preset": "ultrafast"}, + ) + + assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"} + + +@pytest.mark.parametrize( + ("codec_extra", "expected"), + [ + ({}, {"preset": "ultrafast", "threads": "0"}), + ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}), + ({"video_codec_options": {}}, {}), + ({"video_codec_options": None}, None), + ], +) +@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)]) +def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.supported_tasks = frozenset({"t2va"}) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline.device = torch.device("cpu") + pipeline.od_config = SimpleNamespace() + pipeline.load_text_encoder = False + pipeline.load_vae_encoder = False + pipeline.text_encoder = object() + pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) + pipeline._quality_policy = Mock() + pipeline._cache_dit_runtime = Mock() + pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + pipeline.decode_to_mp4 = Mock(return_value=b"mp4") + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_frames=124, + num_inference_steps=50, + extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra}, + ) + batch = DiffusionRequestBatch( + [OmniDiffusionRequest(prompt=_encoder_prompt("test"), sampling_params=sampling, request_id="codec-defaults")] + ) + + pipeline.forward(batch) + + assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected + assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames + + +def test_video_codec_options_are_normalized_for_the_encoder(): + from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options + + assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == { + "preset": "ultrafast", + "threads": "0", + } + assert normalize_video_codec_options(None) is None + with pytest.raises(ValueError, match="video_codec_options"): + normalize_video_codec_options("preset=ultrafast") + + def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -205,9 +488,9 @@ def _write_partition_index(path, *, partition, tasks): @pytest.mark.skip(reason="upstream utility was removed from this branch") def test_modular_diffusers_index_is_resolved_generically(tmp_path): + from vllm_omni.config.resolver import resolve_omni_config from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model - from vllm_omni.entrypoints.utils import resolve_model_config_path (tmp_path / "modular_model_index.json").write_text( json.dumps( @@ -221,7 +504,16 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path): assert is_diffusion_model(str(tmp_path)) assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline" - assert resolve_model_config_path(str(tmp_path)) is None + resolved = resolve_omni_config( + str(tmp_path), + trust_remote_code=False, + deploy_config_path=None, + cli_overrides=None, + stage_overrides=None, + strategy_config_path=None, + ) + assert resolved.config_path is None + assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline" config = OmniDiffusionConfig(model=str(tmp_path)) config.enrich_config() @@ -305,6 +597,9 @@ def test_encoder_task_validation_and_transformer_routing(): with pytest.raises(OmniClientError, match="serves"): pipeline._resolve_task("ref2va", {}, turbo_spec=fl2v_spec) assert pipeline._resolve_task("ref2va", {"image": object()}, turbo_spec=ref2v_spec) == "ref2va" + assert pipeline._resolve_task("t2va") == "t2va" + assert pipeline._resolve_task("fl2va") == "fl2va" + assert pipeline._resolve_task("ref2va") == "ref2va" pipeline.partition = "fl2va" pipeline.supported_tasks = frozenset({"t2va", "fl2va"}) assert pipeline._resolve_task(None, {"audio": object()}, audio_mode="lock_source") == "t2va" @@ -441,7 +736,10 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( revision=None, task_type=task_type, quantization_config=None, - parallel_config=DiffusionParallelConfig(cfg_parallel_size=1), + parallel_config=DiffusionParallelConfig( + cfg_parallel_size=1, + text_encoder_tp_size=1, + ), ) if load_text_encoder is not None: od_config.model_loaded["text_encoder"] = load_text_encoder @@ -474,6 +772,25 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"] assert pipeline._dit_modules == expected_dit_modules assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2) + if expect_encoder: + assert created["text_encoder"] == [str(component_path / "text_encoder")] + assert tokenizer_loader.call_count == 1 + assert processor_loader.call_count == 1 + assert pipeline.text_encoder is not None + assert pipeline.tokenizer is not None + assert pipeline.processor is not None + assert pipeline._encoder_modules == ["text_encoder"] + assert pipeline.text_encoder_tp_size == 1 + else: + assert created["text_encoder"] == [] + tokenizer_loader.assert_not_called() + processor_loader.assert_not_called() + assert pipeline.text_encoder is None + assert pipeline.tokenizer is None + assert pipeline.processor is None + assert pipeline._encoder_modules == [] + assert pipeline.text_encoder_tp_size == 0 + assert pipeline.text_encoder_group is None if expected_partition == "ref2va": assert pipeline._transformer_for_task("ref2va") is pipeline.transformer assert [source.model_or_path for source in pipeline.weights_sources] == [ @@ -580,6 +897,26 @@ def test_joint_postprocess_is_multiprocessing_picklable(): assert result["fps"] == 24 +def test_joint_video_output_is_quantized_before_transfer(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _minimax_h3_post_process, + _prepare_minimax_h3_video_output, + ) + + video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5) + expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1) + prepared = _prepare_minimax_h3_video_output(video) + + assert prepared.dtype == torch.uint8 + assert prepared.is_contiguous() + assert prepared.shape == (2, 2, 4, 5, 3) + torch.testing.assert_close(prepared, expected, rtol=0, atol=0) + + result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6))) + assert result["video"][0].dtype == np.uint8 + np.testing.assert_array_equal(result["video"][0], expected[0].numpy()) + + def test_cfg_parallel_is_rejected_for_distilled_checkpoint(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -839,6 +1176,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): def test_packed_attention_skips_mask_for_packed_mask_free_backend(): + from vllm_omni.diffusion.attention.backends.abstract import ( + VideoTokenLayout, + VideoTokenSpan, + ) from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( MiniMaxH3Attention, ) @@ -865,6 +1206,10 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): torch.nn.Module.__init__(attention) attention.attention = FakeAttention() q = torch.randn(8, 2, 4) + video_layout = VideoTokenLayout( + used_len=5, + video_spans=(VideoTokenSpan(start=3, latent_grid=(1, 1, 2), role="target"),), + ) attention._run_packed_attention( q, @@ -873,6 +1218,9 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), max_seqlen=5, packed_total=8, + video_layout=video_layout, + vsa_prefix_segments=(1, 2), + gate_compress=torch.ones_like(q), ) metadata = attention.attention.metadata @@ -880,6 +1228,11 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): assert metadata.attn_mask is None assert metadata.extra["valid_kv_length"] == 5 assert metadata.extra["npu_attn_varlen"] is True + assert metadata.video_layout is video_layout + assert metadata.extra["vsa_h3_prefix_segments"] == (1, 2) + assert "vsa_h3_video_shape" not in metadata.extra + assert "vsa_h3_target_start" not in metadata.extra + assert metadata.extra["gate_compress"].shape == (1, *q.shape) packed_padding = metadata.packed_padding assert packed_padding is not None assert packed_padding.q_length == 5 @@ -1325,6 +1678,12 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): from vllm_omni.diffusion.models.minimax_h3.encoder import ( MiniMaxH3Qwen3VLEncoder, ) + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + MiniMaxH3Pipeline, + ) + from vllm_omni.diffusion.offloader.component_utils import ( + get_encoder_block_groups, + ) encoder = MiniMaxH3Qwen3VLEncoder( "/nonexistent/text_encoder", @@ -1337,6 +1696,25 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): # The stub has no parameters, so it never contributes to the runner's # strict missing-parameter check on non-encoder ranks. assert list(encoder.named_parameters()) == [] + assert ( + get_encoder_block_groups( + encoder, + "text_encoder", + MiniMaxH3Pipeline._offload_plan, + strict=True, + ) + == [] + ) + + # An arbitrary empty module is not silently treated as a distributed + # stub; explicit plans still validate their declared paths. + with pytest.raises(ValueError, match=r"text_encoder\.vision\.blocks was not found"): + get_encoder_block_groups( + nn.Module(), + "text_encoder", + MiniMaxH3Pipeline._offload_plan, + strict=True, + ) def test_global_quant_config_is_shared_by_dit_and_encoder(): @@ -1403,22 +1781,149 @@ def test_text_encoder_linear_delegates_quantization_to_vllm_factory(): quant_config.get_quant_method.assert_called_once_with(linear, prefix) -def test_distributed_layerwise_offload_stages_vae_component(): +def test_no_offload_keeps_text_encoder_resident(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, enable_layerwise_offload=False, - enable_distributed_layerwise_offload=True, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + input_ids = torch.tensor([1, 2]) + + actual = pipeline._encode_text_hidden(input_ids, {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_not_called() + + +def test_model_offload_uses_hooked_text_encoder_call(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + ) + expected = torch.ones(2, 3) + pipeline.text_encoder = Mock(return_value=expected) + pipeline._model_cpu_offload_modules = [pipeline.text_encoder] + input_ids = torch.tensor([1, 2]) + vision_kwargs = {"pixel_values": torch.ones(1, 4)} + + actual = pipeline._encode_text_hidden(input_ids, vision_kwargs) + + assert actual is expected + pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs) + pipeline.text_encoder.load_to_device.assert_not_called() + + +@pytest.mark.parametrize( + ("enable_layerwise", "enable_distributed"), + [(True, False), (False, True)], +) +def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + diffusion_offload_config=None, + enable_cpu_offload=False, + enable_layerwise_offload=enable_layerwise, + enable_distributed_layerwise_offload=enable_distributed, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + vae = Mock() + with pipeline._component_on_device(vae): + pass + vae.load_to_device.assert_called_once_with() + vae.offload_to_cpu.assert_called_once_with() + + +def test_layerwise_encoder_selection_releases_text_encoder(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + }, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + +def test_layerwise_dit_only_keeps_text_encoder_resident(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={"mode": "layer", "components": ["dit"]}, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_not_called() + + +def test_dit_encoder_selection_keeps_vae_resident(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={ + "mode": "layer", + "components": ["dit", "text_encoder"], + }, ) component = Mock() with pipeline._component_on_device(component): - component.load_to_device.assert_called_once_with() - component.offload_to_cpu.assert_not_called() + pass - component.offload_to_cpu.assert_called_once_with() + component.load_to_device.assert_not_called() + component.offload_to_cpu.assert_not_called() def test_distributed_layerwise_resident_blocks_are_stage_scoped(): @@ -1550,6 +2055,60 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): assert next(video_vae.parameters()).dtype == torch.float32 +@pytest.mark.parametrize("fail_encode", [False, True]) +def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode): + from PIL import Image + + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + events = [] + input_image = Image.new("RGB", (32, 32)) + + @contextmanager + def track_context(device): + assert device == torch.device("cpu") + events.append("enter") + try: + yield + finally: + events.append("exit") + + class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.parallel_tiling = True + + def encode_images(self, image, *, use_fp16_latent): + assert events == ["enter"] + assert image is input_image + assert use_fp16_latent + events.append("encode") + if fail_encode: + raise RuntimeError("encode failed") + return [torch.ones(1, 1, 2, 2)] + + video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE) + torch.nn.Module.__init__(video_vae) + video_vae.model = FakeModel() + video_vae.config_dict = { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + } + monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context) + + if fail_encode: + with pytest.raises(RuntimeError, match="encode failed"): + video_vae.encode_image(input_image) + else: + rows = video_vae.encode_image(input_image) + torch.testing.assert_close(rows, torch.ones(1, 4)) + + assert events == ["enter", "encode", "exit"] + assert video_vae.model.parallel_tiling + + def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module @@ -2379,6 +2938,115 @@ def test_g4_standalone_audio_duration_and_total_duration_contract(): ) +def test_ref2va_video_soundtrack_does_not_consume_standalone_audio_budget(): + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MiniMaxH3EncoderConditioning, + MiniMaxH3EncoderMediaInput, + MiniMaxH3TextConditioning, + ) + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media + + pipeline = _distilled_pipeline([], {"ref2va": None}) + # A 2.35-second video soundtrack and a 15-second standalone reference + # each satisfy their modality's separate 15-second duration budget. + embedded = torch.full((188, 32), 1.0) + standalone = torch.full((1200, 32), 2.0) + video_vae = Mock() + video_vae.encode_image.return_value = torch.ones(1, 96) + video_vae.encode_video.return_value = (torch.ones(1, 96), (1, 2, 2)) + audio_vae = Mock() + audio_vae.encode_waveform.side_effect = [(embedded, 94), (standalone, 600)] + base = MiniMaxH3EncoderConditioning.from_omni_payload(_encoder_payload(task="ref2va", num_frames=362)) + media = MiniMaxH3EncoderMediaInput( + task=base.task, + height=base.height, + width=base.width, + num_frames=base.num_frames, + latent_t=base.latent_t, + audio_t=base.audio_t, + images=(torch.zeros(32, 32, 3, dtype=torch.uint8),), + videos=(torch.zeros(1, 32, 32, 3, dtype=torch.uint8),), + video_audios=((torch.zeros(1, 37600), 16000),), + audios=((torch.zeros(1, 15 * 16000), 16000),), + ) + encoded = encode_media(media, video_vae=video_vae, audio_vae=audio_vae, emit_conditioning=True) + conditioning = MiniMaxH3EncoderConditioning.from_components( + MiniMaxH3TextConditioning(base.hidden_states, base.token_tags), encoded + ) + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_inference_steps=50, + extra_args={"task": "ref2va", "duration": 15.0}, + ) + + context = pipeline._prepare_request_inputs( + {"additional_information": {"encoder_output": conditioning.to_omni_payload()}}, + sampling=sampling, + ) + + assert context["audio_condition_lengths"] == [94, 600] + assert [block["kind"] for block in context["ref_blocks"]] == ["image", "video_audio", "audio"] + torch.testing.assert_close(context["audio_condition"], torch.cat([embedded, standalone])) + + +def test_ref2va_audio_duration_validation_precedes_rank_branch(): + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media + + media = MiniMaxH3EncoderMediaInput( + task="ref2va", + height=32, + width=32, + num_frames=0, + latent_t=1, + audio_t=1, + ) + with pytest.raises(ValueError, match="max_standalone_seconds must be positive"): + encode_media( + media, + video_vae=None, + audio_vae=None, + emit_conditioning=False, + ) + + +def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration(): + from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput + from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media + + observed = [] + + class FakeAudioVAE: + def encode_waveform(self, waveform, sample_rate): + observed.append((waveform.clone(), sample_rate)) + return waveform, 100 + + waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40) + media = MiniMaxH3EncoderMediaInput( + task="ref2va", + height=32, + width=32, + num_frames=60, + latent_t=1, + audio_t=100, + audios=((waveform, 10),), + ) + encoded = encode_media( + media, + video_vae=None, + audio_vae=FakeAudioVAE(), + emit_conditioning=True, + ) + + assert observed[0][0].shape == (1, 25) + assert observed[0][1] == 10 + assert encoded.audio_condition.shape == (1, 25) + assert encoded.audio_condition_lengths == (100,) + + @pytest.mark.parametrize( ("field", "value", "message"), [ @@ -2669,178 +3337,6 @@ def test_pad_seq_len_request_validation_accepts_aligned_values(): assert _resolve_pad_seq_len(MINIMAX_H3_MAX_PAD_SEQ_LEN) == MINIMAX_H3_MAX_PAD_SEQ_LEN -def _turbo_spec(filename: str) -> TurboSpec: - """Build the spec a loaded artifact of this name would carry.""" - spec = parse_turbo_filename(filename) - assert spec is not None - return spec - - -def test_joint_video_output_is_quantized_before_transfer(): - from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( - _minimax_h3_post_process, - _prepare_minimax_h3_video_output, - ) - - video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5) - expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1) - prepared = _prepare_minimax_h3_video_output(video) - - assert prepared.dtype == torch.uint8 - assert prepared.is_contiguous() - assert prepared.shape == (2, 2, 4, 5, 3) - torch.testing.assert_close(prepared, expected, rtol=0, atol=0) - - result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6))) - assert result["video"][0].dtype == np.uint8 - np.testing.assert_array_equal(result["video"][0], expected[0].numpy()) - - -def test_no_offload_keeps_text_encoder_resident(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, - enable_layerwise_offload=False, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - input_ids = torch.tensor([1, 2]) - - actual = pipeline._encode_text_hidden(input_ids, {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_not_called() - - -def test_model_offload_uses_hooked_text_encoder_call(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=True, - enable_layerwise_offload=False, - ) - expected = torch.ones(2, 3) - pipeline.text_encoder = Mock(return_value=expected) - pipeline._model_cpu_offload_modules = [pipeline.text_encoder] - input_ids = torch.tensor([1, 2]) - vision_kwargs = {"pixel_values": torch.ones(1, 4)} - - actual = pipeline._encode_text_hidden(input_ids, vision_kwargs) - - assert actual is expected - pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs) - pipeline.text_encoder.load_to_device.assert_not_called() - - -@pytest.mark.parametrize( - ("enable_layerwise", "enable_distributed"), - [(True, False), (False, True)], -) -def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - diffusion_offload_config=None, - enable_cpu_offload=False, - enable_layerwise_offload=enable_layerwise, - enable_distributed_layerwise_offload=enable_distributed, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_called_once_with() - - vae = Mock() - with pipeline._component_on_device(vae): - pass - vae.load_to_device.assert_called_once_with() - vae.offload_to_cpu.assert_called_once_with() - - -def test_layerwise_encoder_selection_releases_text_encoder(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, - enable_layerwise_offload=False, - enable_distributed_layerwise_offload=False, - diffusion_offload_config={ - "mode": "layer", - "components": ["text_encoder"], - }, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_called_once_with() - - -def test_layerwise_dit_only_keeps_text_encoder_resident(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, - enable_layerwise_offload=False, - enable_distributed_layerwise_offload=False, - diffusion_offload_config={"mode": "layer", "components": ["dit"]}, - ) - pipeline.text_encoder = Mock() - expected = torch.ones(2, 3) - pipeline.text_encoder.encode_ids.return_value = expected - - actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) - - assert actual is expected - pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_not_called() - - -def test_dit_encoder_selection_keeps_vae_resident(): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.od_config = SimpleNamespace( - enable_cpu_offload=False, - enable_layerwise_offload=False, - enable_distributed_layerwise_offload=False, - diffusion_offload_config={ - "mode": "layer", - "components": ["dit", "text_encoder"], - }, - ) - component = Mock() - - with pipeline._component_on_device(component): - pass - - component.load_to_device.assert_not_called() - component.offload_to_cpu.assert_not_called() - - @pytest.mark.parametrize( ("capability", "allow_tf32"), [((8, 6), True), ((9, 0), False), ((10, 0), True), ((10, 3), True), (None, True)], @@ -2900,120 +3396,6 @@ def test_keyframe_encode_cpu_keeps_cudnn_settings(monkeypatch): assert not cudnn.allow_tf32 -@pytest.mark.parametrize("fail_encode", [False, True]) -def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode): - from PIL import Image - - from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module - - events = [] - input_image = Image.new("RGB", (32, 32)) - - @contextmanager - def track_context(device): - assert device == torch.device("cpu") - events.append("enter") - try: - yield - finally: - events.append("exit") - - class FakeModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(1)) - self.parallel_tiling = True - - def encode_images(self, image, *, use_fp16_latent): - assert events == ["enter"] - assert image is input_image - assert use_fp16_latent - events.append("encode") - if fail_encode: - raise RuntimeError("encode failed") - return [torch.ones(1, 1, 2, 2)] - - video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE) - torch.nn.Module.__init__(video_vae) - video_vae.model = FakeModel() - video_vae.config_dict = { - "latent_channels": 1, - "latents_mean": [0.0], - "latents_std": [1.0], - } - monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context) - - if fail_encode: - with pytest.raises(RuntimeError, match="encode failed"): - video_vae.encode_image(input_image) - else: - rows = video_vae.encode_image(input_image) - torch.testing.assert_close(rows, torch.ones(1, 4)) - - assert events == ["enter", "encode", "exit"] - assert video_vae.model.parallel_tiling - - -def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): - from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod - - class FakeEncoder: - instances: list["FakeEncoder"] = [] - - def __init__(self, **kwargs): - self.pushes = [] - self.kwargs = kwargs - self.__class__.instances.append(self) - - def push(self, frames, *, on_consumed=None): - self.pushes.append(np.array(frames, copy=True)) - if on_consumed is not None: - on_consumed() - - def finish(self): - return b"mp4" - - def abort(self): - raise AssertionError("unexpected abort") - - class FakeAudioVAE: - def decode_latent(self, latent): - return torch.zeros(1, 1, 2) - - class FakeVideoVAE: - chunk_value_range = (0.0, 1.0) - - def decode_with_chunks(self, latent, *, on_chunk): - for value in (0.0, 0.25, 0.5): - on_chunk(torch.full((1, 3, 3, 2, 2), value)) - - pipeline = object.__new__(mod.MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.audio_vae = FakeAudioVAE() - pipeline.video_vae = FakeVideoVAE() - pipeline.device = torch.device("cpu") - monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) - monkeypatch.setattr( - "vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", - FakeEncoder, - ) - - output = pipeline.decode_to_mp4( - torch.zeros(1), - torch.zeros(1), - height=2, - width=2, - batch_frames=4, - ) - - assert output == b"mp4" - encoder = FakeEncoder.instances[-1] - assert [frames.shape[0] for frames in encoder.pushes] == [6, 3] - assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8)) - assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8)) - assert encoder.pushes[1].shape == (3, 2, 2, 3) - - def _peer_rank_preencode_pipeline(monkeypatch): """An H3 pipeline whose video VAE is a distributed rank that owns no output. @@ -3083,119 +3465,6 @@ def test_peer_vae_rank_post_decode_reaches_post_processing(monkeypatch): assert mod._minimax_h3_post_process(output.output)["video"] == [b""] -def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch): - """A client's encoder options must survive the worker-side pre-encode path.""" - from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod - - class FakeEncoder: - instances: list["FakeEncoder"] = [] - - def __init__(self, **kwargs): - self.kwargs = kwargs - self.__class__.instances.append(self) - - def push(self, frames, *, on_consumed=None): - del frames - if on_consumed is not None: - on_consumed() - - def finish(self): - return b"mp4" - - def abort(self): - raise AssertionError("unexpected abort") - - class FakeAudioVAE: - def decode_latent(self, latent): - return torch.zeros(1, 1, 2) - - class FakeVideoVAE: - chunk_value_range = (0.0, 1.0) - - def decode_with_chunks(self, latent, *, on_chunk): - on_chunk(torch.zeros(1, 3, 1, 2, 2)) - - pipeline = object.__new__(mod.MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.audio_vae = FakeAudioVAE() - pipeline.video_vae = FakeVideoVAE() - pipeline.device = torch.device("cpu") - monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) - monkeypatch.setattr("vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", FakeEncoder) - - pipeline.decode_to_mp4( - torch.zeros(1), - torch.zeros(1), - height=2, - width=2, - video_codec_options={"preset": "ultrafast"}, - ) - - assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"} - - -@pytest.mark.parametrize( - ("codec_extra", "expected"), - [ - ({}, {"preset": "ultrafast", "threads": "0"}), - ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}), - ({"video_codec_options": {}}, {}), - ({"video_codec_options": None}, None), - ], -) -@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)]) -def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames): - from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline - from vllm_omni.diffusion.request import OmniDiffusionRequest - from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - pipeline = object.__new__(MiniMaxH3Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.partition = "fl2va" - pipeline.supported_tasks = frozenset({"t2va"}) - pipeline.default_video_shift = 12.0 - pipeline.default_audio_shift = 3.0 - pipeline.device = torch.device("cpu") - pipeline.od_config = SimpleNamespace() - pipeline.load_text_encoder = False - pipeline.load_vae_encoder = False - pipeline.text_encoder = object() - pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) - pipeline._quality_policy = Mock() - pipeline._cache_dit_runtime = Mock() - pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) - pipeline.decode_to_mp4 = Mock(return_value=b"mp4") - sampling = OmniDiffusionSamplingParams( - width=1344, - height=768, - fps=24, - num_frames=124, - num_inference_steps=50, - extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra}, - ) - batch = DiffusionRequestBatch( - [OmniDiffusionRequest(prompt=_encoder_prompt("test"), sampling_params=sampling, request_id="codec-defaults")] - ) - - pipeline.forward(batch) - - assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected - assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames - - -def test_video_codec_options_are_normalized_for_the_encoder(): - from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options - - assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == { - "preset": "ultrafast", - "threads": "0", - } - assert normalize_video_codec_options(None) is None - with pytest.raises(ValueError, match="video_codec_options"): - normalize_video_codec_options("preset=ultrafast") - - @pytest.mark.parametrize("duration", [60, 75]) def test_long_video_shape_requires_explicit_opt_in(duration): from vllm_omni.errors import OmniClientError diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py index 510a7b33e..18f1cd9fa 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py +++ b/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"] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py index 824ad2c63..083a43ef2 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/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 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index bebc5a06f..755899649 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/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, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py index fa7410b46..123483c23 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py +++ b/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(): diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index f5e4aaa2c..c76bee391 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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) diff --git a/tests/diffusion/test_diffusion_stage_payload.py b/tests/diffusion/test_diffusion_stage_payload.py new file mode 100644 index 000000000..a5666f7a7 --- /dev/null +++ b/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) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 0b22f6411..bdfe6728c 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/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"), [ diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index 3eadc53c1..3e602426a 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/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( diff --git a/tests/diffusion/test_stage_payload_kv_coexistence.py b/tests/diffusion/test_stage_payload_kv_coexistence.py new file mode 100644 index 000000000..613d20dd3 --- /dev/null +++ b/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()) diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index fb765e8fd..b31a8aa0a 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/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.""" diff --git a/tests/distributed/omni_connectors/test_kv_async_prefetch.py b/tests/distributed/omni_connectors/test_kv_async_prefetch.py index b596e97bb..296940ab6 100644 --- a/tests/distributed/omni_connectors/test_kv_async_prefetch.py +++ b/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") diff --git a/tests/distributed/omni_connectors/test_nixl_connector.py b/tests/distributed/omni_connectors/test_nixl_connector.py index d22ad17a4..49bf24b3b 100644 --- a/tests/distributed/omni_connectors/test_nixl_connector.py +++ b/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 diff --git a/tests/distributed/omni_connectors/test_nixl_connector_native.py b/tests/distributed/omni_connectors/test_nixl_connector_native.py index e0936a02d..956bb126e 100644 --- a/tests/distributed/omni_connectors/test_nixl_connector_native.py +++ b/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) diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index bba6002a0..bb502253a 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/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) diff --git a/tests/e2e/online_serving/test_wan22_t2v.py b/tests/e2e/online_serving/test_wan22_t2v.py index f3f21ae7d..d254a9de7 100644 --- a/tests/e2e/online_serving/test_wan22_t2v.py +++ b/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, ), diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 25e662d27..5f3d11471 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/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) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 5c9c8076f..c5ba1ada2 100644 --- a/tests/helpers/runtime.py +++ b/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: diff --git a/tests/helpers/tests/test_runtime_startup_timeout.py b/tests/helpers/tests/test_runtime_startup_timeout.py index 0eb6a91bf..9432a2bd1 100644 --- a/tests/helpers/tests/test_runtime_startup_timeout.py +++ b/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 diff --git a/tests/model_executor/models/test_minimax_h3_encoder_payload.py b/tests/model_executor/models/test_minimax_h3_encoder_payload.py index 110a0259a..c3f64dcd6 100644 --- a/tests/model_executor/models/test_minimax_h3_encoder_payload.py +++ b/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: diff --git a/tests/model_executor/stage_input_processors/test_minimax_h3.py b/tests/model_executor/stage_input_processors/test_minimax_h3.py index aac16d29f..f44f71aa6 100644 --- a/tests/model_executor/stage_input_processors/test_minimax_h3.py +++ b/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()) diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index bdf9936a4..9006d58c8 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/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.""" diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index e5ac4505d..9303cce5c 100644 --- a/vllm_omni/config/omni_config.py +++ b/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"), diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 584f6af26..95008a70b 100644 --- a/vllm_omni/config/stage_config.py +++ b/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: diff --git a/vllm_omni/deploy/minimax_h3_disaggregated.yaml b/vllm_omni/deploy/minimax_h3_disaggregated.yaml index 25a561c27..2856dff6a 100644 --- a/vllm_omni/deploy/minimax_h3_disaggregated.yaml +++ b/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 diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_nixl.yaml new file mode 100644 index 000000000..59ab77f2e --- /dev/null +++ b/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 diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml index ec11796b6..33a2116c7 100644 --- a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml +++ b/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 diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_turbo_nixl.yaml new file mode 100644 index 000000000..b14002461 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 8f26ebe36..eb025a617 100644 --- a/vllm_omni/diffusion/data.py +++ b/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": } 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"}: diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index ac2059f56..5543577fa 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/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: diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index f583aba14..f1b00155f 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/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: diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 8fc1514ef..f2c9854ee 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/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( diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index b5be82567..b606c39d7 100644 --- a/vllm_omni/diffusion/request.py +++ b/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 diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 1a6fc1aff..14c1c79a9 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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, } ) ) diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 220158f1f..fc795b6b6 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index fe5c1a21f..162718256 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 5f97344df..95db7453f 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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: """ diff --git a/vllm_omni/diffusion/worker/stage_payload.py b/vllm_omni/diffusion/worker/stage_payload.py new file mode 100644 index 000000000..bfdecd229 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index efc935a5d..c8244075e 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/connectors/base.py b/vllm_omni/distributed/omni_connectors/connectors/base.py index 0df428f2f..92257d9de 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/base.py +++ b/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.""" diff --git a/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py b/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py index 7dda744f6..2b84574ec 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/nixl_connector.py +++ b/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: diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index f2c01da9c..ae45a1a3c 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 77731b013..c2e841b8c 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/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*. diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 23c218c71..265a67dca 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/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, ) diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 2ed040703..3349c2166 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/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, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 5cd42c35e..deec36199 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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( diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 7ee949146..78f5add9d 100644 --- a/vllm_omni/engine/stage_client.py +++ b/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 diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index e87d19bb4..a87654cf2 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 46a34dfdb..cb79501cb 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/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): diff --git a/vllm_omni/model_executor/models/minimax_h3/conditioning.py b/vllm_omni/model_executor/models/minimax_h3/conditioning.py index 5d362f440..1444b3348 100644 --- a/vllm_omni/model_executor/models/minimax_h3/conditioning.py +++ b/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 diff --git a/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py b/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py index 7ea4db459..08c298559 100644 --- a/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py +++ b/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") diff --git a/vllm_omni/model_executor/models/minimax_h3/pipeline.py b/vllm_omni/model_executor/models/minimax_h3/pipeline.py index 085e2508a..e7ce1599f 100644 --- a/vllm_omni/model_executor/models/minimax_h3/pipeline.py +++ b/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, diff --git a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py index 78fd0225f..17559ce88 100644 --- a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py +++ b/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)