[CI/Build] Match perf warmups to benchmark concurrency

Use enough warmup requests to cover a full concurrent batch while preserving the existing two-request minimum.

Signed-off-by: Sy03 <24841857+Sy0307@users.noreply.github.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -43,291 +43,21 @@ DLO weight collective.
 
 Host storage is selected separately from the transfer protocol. The loader can
 produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
-otherwise it uses the ordinary loader. In no-AllGather mode, the Phase B host
-weight cache can explicitly normalize supported final ordinary-loader DiT
-tensors into an immutable node-local mmap entry. Ordinary pinned weights remain
-the default after checkpoint mmap fallback when the cache is not selected; cache
-publication and validation add startup and disk lifecycle costs. Replicas can
-share either proven-compatible checkpoint pages or, when selected, equivalent
-final runtime layouts. Under the existing
-pinned-memory policy, the backend attempts to make the complete final mapping
-directly H2D-copyable when the platform provides registration, avoiding the
-recurrent host packing otherwise required by mmap. An optional positive limit
-adds a per-worker safety ceiling. CUDA is the first implementation.
+otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
+the same node can share immutable checkpoint pages when direct mmap is
+selected, while the ordinary-loader fallback still keeps a private runtime
+copy per process.
 
 The Phase A shared-mmap support boundary is TP1. TP greater than one is an
 ordinary-loader compatibility path: DLO can consume the resulting TP-local
-tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
-can share each final TP coordinate among equivalent processes.
+tensors, but those configurations do not use checkpoint mmap and must not be
+used to claim shared-mmap host-memory savings.
 
-The compatibility matrix below defines the current implementation boundary.
-Hardware and performance qualification is tracked separately from this
-long-lived design contract.
+The compatibility matrix below describes the current implementation. The
+unit-level guards are covered, but not every parallelism combination has a
+full model-and-hardware end-to-end test.
 
-## Architecture
-
-DLO is divided into four planes: configuration, loader-owned host storage,
-backend setup, and the per-block inference hot path. Keeping these decisions
-separate is intentional. In particular, an mmap cache is a host-storage choice;
-it does not imply a DLO AllGather, and disabling DLO AllGather does not decide
-whether the host source is private memory, checkpoint mmap, or host weight cache
-mmap. The shared control plane is described once below; the AllGather and
-no-AllGather startup and hot paths are then shown separately.
-
-### Component ownership
-
-The map includes optional mode-specific components, but the ownership rules
-and loader-to-backend boundary are shared.
-
-```mermaid
-flowchart TB
-    CLI["CLI / EngineArgs"] --> Config["OmniDiffusionConfig<br/>validation and operator policy"]
-    Config --> Runner["DiffusionModelRunner<br/>startup coordinator"]
-
-    subgraph LoaderPlane["Loader-owned host-storage plane"]
-        Loader["DiffusersPipelineLoader"]
-        Checkpoint["checkpoint mmap preflight"]
-        Ordinary["ordinary load + callbacks + calibration"]
-        HostWeightCache["host_weight_cache<br/>final-layout publication / join"]
-        Contract["HostWeightPlan<br/>backing + bindings + transforms"]
-        Loader --> Checkpoint
-        Loader -->|"registered final layout requested"| Ordinary
-        Checkpoint -->|"compatible"| Contract
-        Checkpoint -->|"fallback"| Ordinary
-        Ordinary -->|"host weight cache enabled"| HostWeightCache
-        HostWeightCache -->|"enabled and compatible"| Contract
-    end
-
-    Runner -->|"load_model"| Loader
-    Loader -->|"pipeline + optional HostWeightPlan"| Runner
-
-    Config --> Factory["get_offload_backend<br/>OffloadConfig normalization"]
-    Runner -->|"device + optional plan"| Factory
-    Factory --> Backend["DistributedLayerwiseOffloadBackend<br/>resource and transfer owner"]
-    Topology["existing DP / SP process groups"] --> Backend
-    Discovery["ModuleDiscovery + optional OffloadPlan"] --> Backend
-    Contract -.->|"passed once by the runner"| Backend
-
-    Backend -.->|"host weight cache + pinned-memory policy"| Registration["host_registration<br/>platform-neutral lifecycle"]
-    Registration --> CUDA["CUDA registration backend"]
-    Backend --> Hooks["DistributedLayerwiseOffloadHook<br/>one circular hook per streamed block"]
-    Forward["pipeline.forward"] --> Hooks
-    Hooks --> DataPath["two device slots<br/>H2D only or H2D + AllGather"]
-
-    Shutdown["DiffusionWorker.shutdown"] -->|"disable before process-group teardown"| Backend
-```
-
-| Component | Owns | Must not own |
-| --- | --- | --- |
-| CLI, engine arguments, and `OmniDiffusionConfig` | User policy, defaults, validation, and propagation to every worker | Checkpoint inspection, cache publication, or buffer allocation |
-| `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, pinned-memory policy, and optional registration ceiling | Host-weight compatibility decisions |
-| `DiffusersPipelineLoader` | Ordinary loading lifecycle and the decision to return an optional `HostWeightPlan` | DLO process groups or block-stream scheduling |
-| Checkpoint mmap planner | Proving that checkpoint tensors can represent the required runtime layout before ordinary DiT materialization is skipped | Host weight cache identity or transfer policy |
-| `host_weight_cache` | Normalizing final ordinary-loader tensors into an immutable node-local entry and returning a validated plan | Model mutation, host registration, H2D, or collectives |
-| `HostWeightPlan` | The one-way loader-to-offloader contract: backing kind, tensor bindings, source coverage, and any bounded transforms | Cache construction or runtime scheduling |
-| `ModuleDiscovery` and model `OffloadPlan` | Discovering DiT/components and declaring block, resident, and on-demand paths | Loader or cache behavior embedded in model pipelines |
-| `DistributedLayerwiseOffloadBackend` | Consuming the exact plan, selecting the existing group, mapping/repointing storage, allocating shared buffers and streams, optional host registration, installing hooks, and cleanup | Re-running loader compatibility checks or orchestrating replicas |
-| `DistributedLayerwiseOffloadHook` | The block hot path: prefetch, optional DLO AllGather, readiness events, parameter repointing, and replacement with offload placeholders | Cache keys, manifest validation, or request routing |
-| `host_registration` and platform backend | All-or-nothing registration lifecycle for existing immutable mappings | Choosing which weights are cached or changing the two-slot device-buffer contract |
-| `DiffusionWorker` | Calling backend teardown before destroying the distributed environment | Managing backend-owned mappings or registrations directly |
-
-The configuration fields are routed to the component that can enforce their
-contract. The cache root and lock timeout are consumed by the loader.
-`dlo_use_host_weight_cache` explicitly selects final-layout cache publication
-instead of direct-checkpoint mmap. It defaults to false so the ordinary
-pinned-loader path remains available without requiring an opt-out.
-With `pin_cpu_memory=True`, selecting the host weight cache deliberately chooses
-the transform-complete layout even when direct checkpoint mmap is available,
-because only that final layout can be registered without deferred transforms.
-The backend consumes `dlo_host_weight_cache_pin_limit_gib` as an optional
-registration ceiling; zero applies no additional ceiling and does not select
-the cache by itself.
-`dlo_use_allgather` controls loader plan routing because the host weight cache
-is no-AllGather-only, and it selects the backend transfer protocol. This is
-coordination through typed configuration and `HostWeightPlan`, not through
-model-specific global state.
-
-### Runtime host-layout identity (v1)
-
-Host weight cache v1 identifies the loader-produced host representation, not
-the topology used later to move weights to devices. These are separate
-contracts:
-
-- `RuntimeHostLayout` describes tensor ownership and representation after the
-  ordinary loader has applied its parallel-aware transforms. This is the
-  layout stored in the host weight cache.
-- `DLOTransferLayout` describes whether each worker copies its local weights
-  directly or participates in the existing DLO AllGather protocol. It is not
-  part of the host weight cache in v1; AllGather continues to use its existing
-  process-owned host shards.
-
-The v1 physical-sharing invariant is one cache entry per TP coordinate,
-shared by equivalent DP and SP consumers on the same node. The identity
-therefore includes TP world size and TP rank, because different TP coordinates
-own different tensors. It excludes DP world size/rank and SP rank, because
-those dimensions do not change the final tensor ownership supported by v1.
-SP world size and implementation details remain conservative compatibility
-guards until equivalent layouts across those configurations are proven.
-
-PP, HSDP/DTensor, expert parallelism, and CFG parallelism are rejected rather
-than assigned speculative cache identities. Supporting one of these layouts
-requires first defining its loader-produced ownership, representation, and
-sharing domain; it does not follow automatically from supporting its runtime
-collectives. Consequently, this phase implements a TP-coordinate runtime host
-layout, not a general parallel-layout cache.
-
-### Shared startup and lifecycle
-
-The shared control path selects host backing and transfers its ownership to
-the backend before either transfer mode installs hooks:
-
-```mermaid
-sequenceDiagram
-    participant C as OmniDiffusionConfig
-    participant R as DiffusionModelRunner
-    participant L as DiffusersPipelineLoader
-    participant B as DLO backend
-    participant W as DiffusionWorker
-
-    C->>R: validated offload policy
-    R->>L: load_model(config)
-    L->>L: select and validate host backing
-    L-->>R: pipeline + optional HostWeightPlan
-    R->>B: enable(config, device, optional plan)
-    B->>B: discover modules and consume host backing
-    B-->>R: mode-specific backend is ready
-    W->>B: shutdown calls disable()
-    B->>B: synchronize, unregister, remove hooks, close mappings
-```
-
-The loader returns a plan only after its contract is valid: checkpoint plans
-have preflight-proven source coverage and metadata, while host weight cache
-plans also have full manifest, file, and mapped-content validation. The runner
-then transfers ownership exactly once with
-`model_loader.take_host_weight_plan()`. If no plan is returned, the backend
-requires materialized ordinary-loader tensors. If a plan caused ordinary DiT
-materialization to be skipped, failure to create or enable its consumer is a
-hard startup error rather than a fallback to uninitialized parameters.
-
-At inference time, the model pipeline remains the caller: normal block forward
-calls activate the installed hooks. Both modes use one copy stream and two
-shared device slots sized for the largest streamed block. The hooks do not
-call the loader or cache.
-
-### AllGather startup and hot path
-
-```mermaid
-sequenceDiagram
-    participant B as DLO backend
-    participant H as block hooks
-    participant G as existing DP or SP group
-
-    B->>B: choose existing group and prepare persistent local shards
-    B->>B: allocate two device slots and a communication stream
-    B->>H: install hooks and prime the first block
-    loop each streamed block
-        H->>H: copy the next local shard to the alternate slot
-        H->>G: all_gather_into_tensor(next block)
-        G-->>H: complete next block is ready
-        H->>H: execute current block and rotate slots
-    end
-```
-
-AllGather uses the existing DP group when DP is greater than one, otherwise
-the existing SP group when SP is greater than one. TP is never selected as
-DLO's weight-sharding group. Every participating rank must request blocks in
-the same order; fast P2P reduces communication cost but does not remove that
-lockstep contract. An effective group size of one issues no collective and
-uses rank-local transfer.
-
-### No-AllGather startup and hot path
-
-```mermaid
-sequenceDiagram
-    participant L as DiffusersPipelineLoader
-    participant C as host_weight_cache
-    participant B as DLO backend
-    participant P as platform registration
-    participant H as block hooks
-
-    alt host weight cache explicitly selected
-        L->>L: load, mutate, validate, and calibrate final tensors
-        L->>C: publish or join final-layout entry
-        C-->>L: validated host weight cache plan or fail-closed fallback
-        L-->>B: cache plan or materialized fallback via runner
-    else automatic checkpoint selection
-        L->>L: checkpoint mmap preflight
-        alt loader selects a compatible checkpoint plan
-            L-->>B: checkpoint plan via runner
-        else ordinary loading
-            L->>L: load, mutate, validate, and calibrate final tensors
-            L-->>B: materialized tensors via runner
-        end
-    end
-    B->>B: consume complete rank-local block sources
-    opt host weight cache plan and pinned-memory policy
-        B->>P: register the complete immutable mapping
-        P-->>B: direct-H2D source or bounded-staging fallback
-    end
-    B->>H: install hooks and prime the first block
-    loop each streamed block
-        alt registered host weight cache mapping
-            H->>H: copy complete next block directly to device
-        else pageable mmap mapping
-            H->>H: pack into a bounded staging slot, then copy to device
-        else ordinary runtime tensors
-            H->>H: copy complete next block to device
-        end
-        H->>H: execute current block and rotate slots
-    end
-```
-
-No-AllGather creates no DLO weight group and issues no DLO collective.
-Workers may therefore schedule requests independently. Ordinary TP and SP
-collectives remain inside model execution; host registration changes only
-the source of the complete-block H2D copy.
-
-### Host backing and transfer composition
-
-| Loader-selected host backing | Backend transfer mode | Setup result | Per-block hot path |
-| --- | --- | --- | --- |
-| Ordinary runtime tensors | AllGather group larger than one | Each rank retains a pinned local weight shard | shard H2D → DLO AllGather → device slot |
-| Checkpoint mmap | AllGather group larger than one | Each rank prepares its persistent shard, then releases the source mapping | shard H2D → DLO AllGather → device slot |
-| Ordinary runtime tensors | Rank-local / no-AllGather | Each rank retains the complete loader-produced local layout | complete-block H2D → device slot |
-| Checkpoint mmap, or host weight cache with pinned memory disabled/registration unavailable | Rank-local / no-AllGather | Immutable mapping plus two bounded host-staging slots | mmap views → staging slot → device slot |
-| Host weight cache mmap with pinned memory enabled and registration available | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot |
-
-The three axes therefore remain independent:
-
-1. **Offload policy:** whether DLO is enabled and which layers are resident.
-2. **Host backing:** ordinary tensors, checkpoint mmap, or final host weight cache
-   mmap, selected by the loader.
-3. **Transfer protocol:** rank-local H2D or H2D plus DLO AllGather, selected by
-   the backend from validated configuration and existing topology.
-
-### Failure and lifecycle boundaries
-
-- Invalid option combinations fail during configuration validation.
-- Checkpoint mmap preflight failure falls back before ordinary loading is
-  skipped.
-- Host weight cache incompatibility, timeout, or publication failure retains
-  the already-valid ordinary tensors.
-- Host-registration failure rolls back and retains bounded staging. A rollback
-  failure aborts startup because unmapping a range still owned by the platform
-  would be unsafe.
-- During shutdown, DLO synchronizes outstanding device work, unregisters host
-  mappings, removes hooks, and closes mmap handles only after safe unregister;
-  the worker destroys its process groups afterward.
-
-DLO deliberately does not own DP replica placement, independent-engine request
-routing, cross-engine admission control, or cross-node cache distribution.
-Those belong to an external router/orchestrator. It also does not replace TP or
-SP model collectives; it only adds an optional weight collective over an
-already-existing DP or SP group.
-
-## Shared design
-
-The following topology and loader contracts apply to both transfer modes.
+## Design
 
 ### DLO consumes the existing parallel topology
 
@@ -335,10 +65,15 @@ DLO does not create a new DP, TP, or SP topology. It reads the configured
 `DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
 standard distributed groups have been initialized.
 
-AllGather mode selects one existing group for DLO weight sharding; no-AllGather
-mode creates no DLO group. Neither mode changes ordinary TP or SP model
-collectives. The mode-specific group and scheduling contracts are described
-below.
+The DLO weight-sharding group is selected as follows:
+
+1. Use the existing DP group when `data_parallel_size > 1`.
+2. When DP is one and SP is greater than one, use the SP group.
+3. Otherwise, run rank-locally without a DLO process group.
+
+TP is deliberately not used as DLO's AllGather group. HSDP has its own
+parameter-sharding lifecycle and is not allowed to be sharded a second time by
+DLO's AllGather path.
 
 ### The loader owns host-weight planning
 
@@ -370,24 +105,7 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
 transforms. Model-specific direct-layout knowledge, when required, lives in a
 checkpoint adapter beside the ordinary loader.
 
-## AllGather design
-
-### Group selection
-
-The DLO weight-sharding group is selected in this order:
-
-1. Use the existing DP group when `data_parallel_size > 1`.
-2. When DP is one and SP is greater than one, use the existing SP group.
-3. Otherwise, use an effective group size of one and issue no collective.
-
-TP is deliberately not used as DLO's AllGather group. HSDP has its own
-parameter-sharding lifecycle and cannot be sharded a second time by DLO.
-
-### Sharded transfer and double buffering
-
-The loader may provide ordinary runtime tensors or a compatible checkpoint
-plan. During backend setup, either source is normalized into persistent local
-shards. The final-layout host weight cache is not selected in AllGather mode.
+### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
 `1 / group_size` of each streamable block in pinned host memory. The next
@@ -395,15 +113,6 @@ block's shard is copied to a device buffer and reconstructed with
 `all_gather_into_tensor` on a communication stream while the current block is
 executing.
 
-![AllGather DLO sharded host weights and lockstep reconstruction](../../figures/dlo/dlo_allgather_sharded_host.svg)
-
-The figure uses DP2 with TP2 to expose both parallel dimensions. DLO creates
-one AllGather group across the DP replicas for each TP coordinate; TP0 workers
-reconstruct TP0 blocks together and TP1 workers reconstruct TP1 blocks
-together. Ordinary TP collectives remain inside each DP replica. Unlike the
-no-AllGather host weight cache path, these persistent pinned host shards are
-process-owned, and every participating rank must follow the same block order.
-
 ```text
 Compute:    [Block N]             [Block N+1]          [Block N+2]
 H2D:                      [shard N+1]           [shard N+2]
@@ -429,29 +138,11 @@ When DP is greater than one, the engine can process one request per DP rank in
 the same denoising wave. Because AllGather is a collective, all participating
 requests must take the same execution path at every denoising step.
 
-Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
-independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
-plus shared mmap storage avoids coupling their request schedules and failure
-domains.
-
-## No-AllGather design
-
-### Rank-local transfer
+### Rank-local path without DLO AllGather
 
 With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
 one and streams complete blocks using H2D copies only. The host backing may be
-a loader-approved checkpoint mapping, a normalized host weight cache mapping, or
-ordinary runtime tensors.
-
-![No-AllGather DLO shared host weights and independent runtime orchestration](../../figures/dlo/dlo_no_allgather_shared_host.svg)
-
-The figure uses TP2 to make per-coordinate physical sharing explicit. Each
-engine contains two rank-local workers and devices. The cache keeps one
-immutable entry per TP coordinate, so equivalent TP0 workers across engines
-share the TP0 entry and equivalent TP1 workers share the separate TP1 entry.
-An external router or orchestrator chooses each engine's requests; DLO
-coordinates only the worker-local mapping, hooks, transfer streams, events,
-and device slots.
+either a loader-approved checkpoint mapping or ordinary runtime tensors.
 
 For direct mmap, each process retains immutable safetensors views and uses two
 bounded pinned host staging slots. Processes on the same node that map the same
@@ -460,17 +151,10 @@ the persistent private full-model copy per pure-DP process, but each process
 still packs and transfers every complete block. Sharing is node-local; each
 node has its own page cache.
 
-For a host weight cache mapping, the default pinned-memory policy attempts to
-replace those host slots with direct copies from the shared final tensor views.
-Disabling pinned memory or a failed registration retains staging behavior.
-
-When direct mmap preflight fails or the host weight cache is selected, the
-regular model loader remains responsible for preparing each rank's weights,
-including TP-local tensors or HSDP-managed parameters. Ordinary pinned tensors
-remain the default without cache selection. When the host weight cache is
-explicitly selected,
-supported CPU layouts enter the cache; unsupported layouts keep one private
-runtime copy per process.
+When direct mmap preflight fails, the regular model loader remains responsible
+for preparing each rank's weights, including TP-local tensors or HSDP-managed
+parameters. In that fallback, each pure-DP process keeps a private full runtime
+copy.
 
 This mode means:
 
@@ -481,109 +165,18 @@ This mode means:
 - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
   DLO's additional weight AllGather is disabled.
 - Pure DP deployments share one checkpoint-backed copy per node when direct
-  mmap is selected. The host weight cache can share equivalent ordinary-loader
-  outputs; its fallback keeps one private runtime copy per rank.
+  mmap is selected; the ordinary-loader fallback keeps one private runtime
+  copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
-### Post-load host weight cache
-
-The Phase B cache preserves the same ownership boundary. DLO neither derives
-cache keys nor interprets model loaders. No-AllGather DLO first attempts a
-compatible direct-checkpoint mmap plan. When that plan is unavailable, every
-rank completes the ordinary loader, custom callbacks, post-load processing,
-strict validation, and calibration. When cache use or registration is
-explicitly selected, the loader:
-
-1. discovers final CPU DiT parameters and persistent buffers;
-2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
-   ambiguous layouts before changing the model;
-3. hashes tensor names, metadata, and final bytes;
-4. builds or joins one content-addressed entry under a strict POSIX lock; and
-5. returns a `HostWeightPlan(backing_kind="host_weight_cache")` only after file,
-   manifest, and remapped-tensor content validation succeeds.
-
-The identity excludes DP size/rank and SP rank. It includes TP world size/rank
-and conservative SP implementation/world-size guards. DP replicas and
-equivalent SP ranks can therefore share; different TP coordinates cannot.
-Final tensor content is authoritative, so an unmodeled loader/config change
-that changes bytes creates a different entry instead of risking silent
-sharing.
-
-Publication writes sharded safetensors and a JSON manifest into a same-
-filesystem temporary directory, flushes them, and atomically renames the
-directory. The manifest includes per-file SHA-256 digests and the final tensor
-content digest. Kernel `flock` ownership distinguishes a live writer from a
-stale lock file: process death releases the lock automatically. Stale temporary
-entries are removed only by the next lock owner. A bounded wait or any cache
-error retains the already-valid ordinary tensors.
-
-DLO maps this plan without moving the DiT to `meta` and without rerunning any
-post-load hook. It changes tensor storage while preserving `Parameter` objects,
-then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
-streaming hooks. The source mappings remain open through DLO disable. With
-pinned memory disabled, or when registration is unavailable, they feed the
-existing bounded two-slot host-staging path. Every streamed block is first
-packed from pageable mmap storage into a staging slot and is then copied to the
-device. This removes model-sized private pinned storage, but adds a recurrent
-host-to-host copy and can amplify TP synchronization waits when several ranks
-compete for host bandwidth.
-
-This version is deliberately steady-state-oriented. It requires roughly one
-local-disk model copy per runtime identity, has no eviction, and still performs
-ordinary loading plus full content-hash passes in every rank. Cache publication
-or validation failure falls back to the already-valid ordinary host tensors.
-Skip-load-on-hit is a separate follow-up after loader side effects can be
-modeled safely.
-
-### Registered host weight cache transfer
-
-Selecting the host weight cache with `pin_cpu_memory=True` opts a worker into
-registration of the complete mapping when its platform supports it. It also
-changes loader selection: no-AllGather uses the ordinary loader and final-layout
-host weight cache even when a TP1 direct-checkpoint plan is available. A
-checkpoint binding can carry a deferred adapter such as
-grouped-QKV reordering. Registering that raw source would not remove the
-adapter's recurrent CPU allocation/copy, whereas the host weight cache contains
-the transform-complete final tensors.
-
-After mapping and validating every final tensor, DLO groups source storages by
-backing file, page-aligns their address ranges, and coalesces overlapping or
-adjacent ranges within that mapping. When a positive per-worker ceiling is
-configured, the complete byte count is checked against it before any platform
-mutation. The CUDA backend
-first requires `cudaDevAttrHostRegisterReadOnlySupported`, then registers every
-range with the read-only host-registration flag and verifies that PyTorch
-recognizes every mapped tensor as pinned. A device without that capability
-retains bounded staging; vLLM-Omni does not replace immutable mappings with a
-writable or copy-on-write CUDA-specific layout. Other registration failures
-roll back any ranges already registered and preserve the existing two-slot
-staging path. Partial direct transfer is deliberately unsupported. A rollback
-failure aborts worker initialization so a still-registered range is retained
-until process/context teardown rather than being unsafely unmapped.
-
-On success, the hooks copy each immutable tensor view directly into its offset
-in the existing flattened rotating device buffer on the copy stream. Parameter
-repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
-direct path trades recurrent host packing for finer-grained H2D operations. A
-packed cache schema should be considered only if that additional format and
-publication complexity becomes necessary.
-
-Registration is per process and platform context; physical file-cache pages
-remain shared and are not copied into one private host allocation per worker.
-The operator must nevertheless budget page-locked memory and satisfy the
-platform and OS registration limits for every worker. During shutdown, DLO
-first synchronizes outstanding device work, unregisters all ranges, and then
-closes the safetensors mappings. Worker teardown invokes this cleanup before
-destroying the distributed environment and platform context.
-
 ## Parallelism compatibility
 
 | Parallelism | DLO + AllGather | DLO without AllGather |
 |---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages; the optional host weight cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
-| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active without a DLO weight collective. Host weight cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
-| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Host weight cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings avoids the recurrent staging copy but does not remove ordinary TP collectives. |
-| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted only with ordinary-loader rank-local tensors. HSDP owns parameter sharding and its own gathers; host weight cache v1 rejects HSDP/DTensor layouts. |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
+| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
+| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
 
 ### Combined dimensions
 
@@ -599,9 +192,7 @@ destroying the distributed environment and platform context.
 - **HSDP + DP or TP:** rejected independently by the diffusion parallel
   configuration.
 
-## Mode constraints
-
-### AllGather constraints
+## Request and loading constraints
 
 AllGather DP multi-concurrency requires:
 
@@ -609,54 +200,109 @@ AllGather DP multi-concurrency requires:
 - the same `num_inference_steps` for all requests in a wave; and
 - identical request arguments that affect the collective execution path.
 
-Per-tensor online FP8 linears can use DLO AllGather after the ordinary loader
-finalizes their runtime weights and scales. Other online methods must use
-`--dlo-no-use-allgather` or disable online quantization until their runtime
-layouts are validated.
-
-### No-AllGather constraints
-
 The no-AllGather path does not impose these DLO-specific synchronized-wave
-requirements. Host weight cache v1 is no-AllGather only and rejects all
-quantization configurations, HSDP/DTensor, expert parallelism, CFG
-parallelism, PP, non-contiguous tensors, and tied/shared storage.
-
-### Shared loading constraint
+requirements.
 
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
-ordinary loader.
+ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
+ordinary loader finalizes their runtime weights and scales; DLO then shards and
+reconstructs those tensors with their recorded layouts. Other online methods
+must use `--dlo-no-use-allgather` or disable online quantization until their
+runtime layouts are validated.
 
-## Validation status
+A normalized runtime mmap cache, built through the ordinary loader, is the
+proposed general mechanism for sharing transformed TP or quantized layouts.
+That cache and its publication/lifecycle protocol are intentionally outside
+this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
 
-Historical test counts and hardware measurements are intentionally not kept
-in this design document because they can drift as the implementation changes.
-Revalidation must run against the final commit and publish commit-pinned
-artifacts before making hardware or performance claims. Track that work in
-[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
+## Validation coverage
+
+Current source-level validation includes:
+
+- HSDP + DLO + AllGather rejection;
+- HSDP + DLO without AllGather acceptance at configuration level;
+- loader preflight fallback for TP, HSDP, online quantization, unknown custom
+  loaders, missing keys, and shape/dtype mismatches;
+- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
+  sharding of finalized weights and scales;
+- exact loader-to-backend plan transfer and ordinary-loader fallback;
+- rank-local mmap source retention, bounded two-slot staging, and adapter
+  transforms without parameter-side flags;
+- resident-layer requests requiring no-AllGather;
+- DP request-wave validation for denoising-step compatibility;
+- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
+  tests.
+
+### B300 parallel-topology smoke matrix
+
+A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
+CUDNN attention backend, 256x256 output, two denoising steps, and
+`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
+VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
+or shared-mmap host-memory savings.
+
+| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
+|---|---:|---:|---:|---:|
+| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
+| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
+| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
+| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
+
+Within each topology, the AllGather and no-AllGather video and audio outputs
+were byte-identical. All four runs completed without an `ERROR` or traceback
+and released their device allocations. For DP4xTP1, no-AllGather direct mmap
+reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
+125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
+ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
+AllGather, because DP replicas did not share checkpoint-backed runtime
+weights. This is a functional and memory smoke test, not a production-quality
+performance or output-quality benchmark.
+
+### Host-memory measurement
+
+A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
+ordinary-loader fallback with direct mmap. Both runs used
+DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
+256x256 four-second request. The ordinary-loader workers were sampled after
+initialization. The mmap workers were sampled after one completed request, so
+the checkpoint working set had been faulted into the page cache; this is the
+more conservative point for mmap.
+
+The values below come from `/proc/<worker>/smaps_rollup` and include the whole
+worker, not only the DiT. The stable rank-to-rank difference comes from other
+pipeline components, so each worker should be compared with the same worker in
+the other storage mode.
+
+| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
+|---|---:|---:|---:|---:|---:|
+| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
+| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
+| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
+
+The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
+31.20 GiB `Pss_File`, which is the proportional charge expected when the same
+resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
+from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
+1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
+because it counts a shared physical page in every process that maps it; summed
+PSS is the appropriate node-memory comparison.
+
+The highest-value missing coverage is broader end-to-end numerical and
+lifecycle comparison against ordinary layerwise offload for DP+SP,
+HSDP+SP+no-AllGather, and TP greater than one across additional models and
+target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
+change the Phase A direct-mmap TP1 support boundary.
 
 ## Recommendations
 
-### AllGather
-
-- Use **DP + DLO AllGather** when ranks already execute the same block path in
-  lockstep and belong to the same fast P2P domain. It shards persistent host
-  weights across that group.
+- Use **DP + DLO AllGather** for the supported throughput and host-memory
+  scaling path.
 - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
   not the goal.
-
-### No-AllGather
-
 - Use **no-AllGather** when independent replica execution is required. TP1
-  direct-mmap deployments can share checkpoint pages per node. Enable the
-  host weight cache only when independent replicas have repeated final layouts
-  and benchmark the target CPU/memory topology, especially for matching
-  coordinates across multiple TP engines. With the default pinned-memory
-  policy, supported CUDA workers prefer registered direct H2D. Configure a
-  positive ceiling when an explicit page-locking bound is required, or disable
-  pinned memory when pageable staging is operationally preferable.
-
-### HSDP
-
-- Do not combine HSDP with DLO AllGather or host weight cache v1. HSDP with
-  no-AllGather remains an ordinary-loader rank-local path.
+  direct-mmap deployments can share checkpoint pages per node; other layouts
+  retain the ordinary loader's private host memory behavior and are outside
+  the Phase A shared-mmap support scope.
+- Prefer **HSDP alone** for production HSDP deployments until the combined
+  HSDP + DLO no-AllGather path has broader end-to-end coverage.

docs/design/figures/dlo/dlo_allgather_sharded_host.svg

@@ -1,165 +0,0 @@
-<svg xmlns="http://www.w3.org/2000/svg" width="1280" height="850" viewBox="0 0 1280 850" role="img" aria-labelledby="title desc">
-  <title id="title">AllGather DLO sharded host weights and lockstep device reconstruction</title>
-  <desc id="desc">An external orchestrator submits one synchronized wave to a DP2 by TP2 diffusion engine. Each DP replica contains TP0 and TP1 workers with private pinned host shards and two rotating GPU weight slots. DLO AllGather groups connect matching TP coordinates across DP replicas, reconstructing full rank-local blocks on each device in lockstep. Ordinary TP groups remain within each DP replica.</desc>
-  <defs>
-    <marker id="request-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
-      <path d="M 0 0 L 10 5 L 0 10 z" fill="#7c3aed"/>
-    </marker>
-    <marker id="h2d-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
-      <path d="M 0 0 L 10 5 L 0 10 z" fill="#d97706"/>
-    </marker>
-    <marker id="collective-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
-      <path d="M 0 0 L 10 5 L 0 10 z" fill="#7c3aed"/>
-    </marker>
-    <style>
-      .canvas { fill: #ffffff; stroke: #e2e8f0; }
-      .node { fill: #f8fafc; stroke: #cbd5e1; }
-      .engine { fill: #ffffff; stroke: #cbd5e1; }
-      .replica { fill: #f8fafc; stroke: #cbd5e1; }
-      .router { fill: #f5f3ff; stroke: #c4b5fd; }
-      .scope { fill: #f8fafc; stroke: #cbd5e1; }
-      .host { fill: #fff7ed; stroke: #fdba74; }
-      .shard { fill: #d97706; }
-      .gpu { fill: #f0fdf4; stroke: #86efac; }
-      .slot-current { fill: #16a34a; }
-      .slot-next { fill: #bbf7d0; stroke: #22c55e; }
-      .lockstep { fill: #f5f3ff; stroke: #c4b5fd; }
-      .title { font: 700 25px system-ui, -apple-system, Segoe UI, sans-serif; fill: #0f172a; }
-      .subtitle { font: 500 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
-      .section { font: 700 15px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
-      .section-sub { font: 500 11.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
-      .eyebrow { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; letter-spacing: .8px; fill: #64748b; }
-      .box-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
-      .box-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
-      .shard-title { font: 700 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
-      .shard-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
-      .slot-title { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
-      .slot-text { font: 500 9px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
-      .slot-next-title { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
-      .slot-next-text { font: 500 9px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
-      .request-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
-      .h2d-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b45309; }
-      .collective-label { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
-      .lock-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
-      .lock-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #5b21b6; }
-      .request-path { fill: none; stroke: #7c3aed; stroke-width: 2; marker-end: url(#request-arrow); }
-      .h2d-path { fill: none; stroke: #d97706; stroke-width: 2; marker-end: url(#h2d-arrow); }
-      .collective-path { fill: none; stroke: #7c3aed; stroke-width: 2.5; marker-start: url(#collective-arrow); marker-end: url(#collective-arrow); }
-    </style>
-  </defs>
-
-  <rect x="12" y="12" width="1256" height="826" rx="14" class="canvas"/>
-
-  <text id="title-text" x="40" y="50" class="title">AllGather DLO: sharded host weights, lockstep reconstruction</text>
-  <text x="40" y="74" class="subtitle">DP2 &#215; TP2 shown: DLO groups span DP replicas per TP coordinate; ordinary TP stays within each replica</text>
-
-  <!-- The external control plane submits one synchronized DP wave. -->
-  <rect x="390" y="94" width="500" height="68" rx="10" class="router"/>
-  <text x="640" y="121" text-anchor="middle" class="section">External router / orchestrator</text>
-  <text x="640" y="144" text-anchor="middle" class="section-sub">submits one synchronized DP wave &#183; outside vLLM-Omni DLO</text>
-
-  <rect x="30" y="184" width="1220" height="634" rx="14" class="node"/>
-  <text x="52" y="211" class="eyebrow">ONE NODE &#183; ONE SYNCHRONIZED ENGINE, PROCESS-OWNED HOST SHARDS</text>
-
-  <path d="M 640 162 L 640 232" class="request-path"/>
-  <text x="654" y="202" class="request-label">request wave</text>
-
-  <!-- One engine contains two DP replicas that execute in lockstep. -->
-  <rect x="45" y="232" width="1190" height="262" rx="12" class="engine"/>
-  <text x="68" y="261" class="section">One synchronized diffusion engine &#183; DP2 &#215; TP2</text>
-  <text x="68" y="281" class="section-sub">two DP requests follow the same block order; DLO reuses existing DP groups at each TP coordinate</text>
-
-  <!-- DP replica 0. -->
-  <rect x="65" y="296" width="515" height="186" rx="10" class="replica"/>
-  <text x="82" y="318" class="box-title">DP replica 0 &#183; request 0</text>
-  <text x="563" y="318" text-anchor="end" class="section-sub">ordinary TP group</text>
-
-  <rect x="82" y="332" width="473" height="64" rx="8" class="gpu"/>
-  <text x="98" y="351" class="eyebrow">TP0 &#183; GPU 0</text>
-  <text x="98" y="376" class="box-text">DLO backend &#183; local shard H2D</text>
-  <rect x="250" y="340" width="132" height="44" rx="6" class="slot-current"/>
-  <text x="316" y="357" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
-  <text x="316" y="374" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
-  <rect x="397" y="340" width="142" height="44" rx="6" class="slot-next"/>
-  <text x="468" y="357" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
-  <text x="468" y="374" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
-
-  <rect x="82" y="410" width="473" height="64" rx="8" class="gpu"/>
-  <text x="98" y="429" class="eyebrow">TP1 &#183; GPU 1</text>
-  <text x="98" y="454" class="box-text">DLO backend &#183; local shard H2D</text>
-  <rect x="250" y="418" width="132" height="44" rx="6" class="slot-current"/>
-  <text x="316" y="435" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
-  <text x="316" y="452" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
-  <rect x="397" y="418" width="142" height="44" rx="6" class="slot-next"/>
-  <text x="468" y="435" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
-  <text x="468" y="452" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
-
-  <!-- DP replica 1. -->
-  <rect x="700" y="296" width="515" height="186" rx="10" class="replica"/>
-  <text x="717" y="318" class="box-title">DP replica 1 &#183; request 1</text>
-  <text x="1198" y="318" text-anchor="end" class="section-sub">ordinary TP group</text>
-
-  <rect x="725" y="332" width="473" height="64" rx="8" class="gpu"/>
-  <text x="741" y="351" class="eyebrow">TP0 &#183; GPU 2</text>
-  <text x="741" y="376" class="box-text">DLO backend &#183; local shard H2D</text>
-  <rect x="893" y="340" width="132" height="44" rx="6" class="slot-current"/>
-  <text x="959" y="357" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
-  <text x="959" y="374" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
-  <rect x="1040" y="340" width="142" height="44" rx="6" class="slot-next"/>
-  <text x="1111" y="357" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
-  <text x="1111" y="374" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
-
-  <rect x="725" y="410" width="473" height="64" rx="8" class="gpu"/>
-  <text x="741" y="429" class="eyebrow">TP1 &#183; GPU 3</text>
-  <text x="741" y="454" class="box-text">DLO backend &#183; local shard H2D</text>
-  <rect x="893" y="418" width="132" height="44" rx="6" class="slot-current"/>
-  <text x="959" y="435" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
-  <text x="959" y="452" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
-  <rect x="1040" y="418" width="142" height="44" rx="6" class="slot-next"/>
-  <text x="1111" y="435" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
-  <text x="1111" y="452" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
-
-  <!-- DLO collectives connect matching TP coordinates across DP replicas. -->
-  <text x="640" y="350" text-anchor="middle" class="collective-label">DLO AG &#183; TP0 GROUP</text>
-  <path d="M 555 364 L 725 364" class="collective-path"/>
-  <text x="640" y="428" text-anchor="middle" class="collective-label">DLO AG &#183; TP1 GROUP</text>
-  <path d="M 555 442 L 725 442" class="collective-path"/>
-
-  <!-- All participants must reach the same collective in the same order. -->
-  <rect x="365" y="516" width="550" height="48" rx="8" class="lockstep"/>
-  <circle cx="395" cy="540" r="12" fill="#7c3aed"/>
-  <text x="395" y="545" text-anchor="middle" font-family="system-ui, sans-serif" font-size="15" font-weight="700" fill="#ffffff">&#8644;</text>
-  <text x="420" y="536" class="lock-title">DLO AllGather requires lockstep block order</text>
-  <text x="420" y="553" class="lock-text">every participant enters each block collective in the same sequence</text>
-
-  <!-- Each replica copies two process-owned host shards before reconstruction. -->
-  <path d="M 455 650 C 412 625, 340 610, 340 575 L 340 540 C 340 516, 318 500, 300 482" class="h2d-path"/>
-  <text x="276" y="618" text-anchor="middle" class="h2d-label">local shard H2D &#215;2</text>
-  <path d="M 825 650 C 868 625, 940 610, 940 575 L 940 540 C 940 516, 962 500, 980 482" class="h2d-path"/>
-  <text x="1004" y="618" text-anchor="middle" class="h2d-label">local shard H2D &#215;2</text>
-
-  <!-- Persistent host shards are private to workers, not mmap-shared. -->
-  <rect x="58" y="650" width="242" height="140" rx="10" class="scope"/>
-  <text x="78" y="676" class="eyebrow">GROUP TOPOLOGY</text>
-  <text x="78" y="702" class="box-title">DP2 &#215; TP2</text>
-  <text x="78" y="725" class="box-text">DLO groups: same TP coordinate</text>
-  <text x="78" y="746" class="box-text">across DP replicas</text>
-  <text x="78" y="767" class="box-text">TP stays the model-parallel group.</text>
-
-  <rect x="330" y="650" width="620" height="140" rx="10" class="host"/>
-  <text x="350" y="676" class="eyebrow">PROCESS-OWNED PINNED HOST SHARDS &#183; NOT HOST WEIGHT CACHE MMAP</text>
-  <rect x="370" y="692" width="260" height="60" rx="8" class="shard"/>
-  <text x="500" y="716" text-anchor="middle" class="shard-title">DP replica 0 &#183; private</text>
-  <text x="500" y="738" text-anchor="middle" class="shard-text">TP0 shard 0/2 &#183; TP1 shard 0/2</text>
-  <rect x="650" y="692" width="260" height="60" rx="8" class="shard"/>
-  <text x="780" y="716" text-anchor="middle" class="shard-title">DP replica 1 &#183; private</text>
-  <text x="780" y="738" text-anchor="middle" class="shard-text">TP0 shard 1/2 &#183; TP1 shard 1/2</text>
-  <text x="640" y="773" text-anchor="middle" class="box-text">Across each DLO group, shards total approximately one rank-local runtime block copy</text>
-
-  <rect x="980" y="650" width="242" height="140" rx="10" class="scope"/>
-  <text x="1000" y="676" class="eyebrow">HOT PATH &#183; PER WORKER</text>
-  <text x="1000" y="702" class="box-title">Prefetch next block</text>
-  <text x="1000" y="725" class="box-text">local pinned shard &#8594; GPU</text>
-  <text x="1000" y="746" class="box-text">DLO AllGather &#8594; full block</text>
-  <text x="1000" y="767" class="box-text">rotate slots; overlap compute</text>
-</svg>

docs/design/figures/dlo/dlo_no_allgather_shared_host.svg

@@ -1,202 +0,0 @@
-<svg xmlns="http://www.w3.org/2000/svg" width="1280" height="850" viewBox="0 0 1280 850" role="img" aria-labelledby="title desc">
-  <title id="title">No-AllGather DLO shared host weights and independent runtime orchestration</title>
-  <desc id="desc">An external orchestrator routes requests to two independently scheduled TP2 diffusion engines on one node. Each engine contains two rank-local workers and devices, with private loaders, mmap views, DLO runtime hooks, streams, events, and rotating GPU slots. Equivalent TP0 workers across engines map one immutable TP0 entry in the host weight cache, while equivalent TP1 workers map a separate TP1 entry, without a cross-engine DLO AllGather.</desc>
-  <defs>
-    <marker id="request-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
-      <path d="M 0 0 L 10 5 L 0 10 z" fill="#7c3aed"/>
-    </marker>
-    <marker id="mapping-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
-      <path d="M 0 0 L 10 5 L 0 10 z" fill="#d97706"/>
-    </marker>
-    <marker id="runtime-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
-      <path d="M 0 0 L 10 5 L 0 10 z" fill="#2563eb"/>
-    </marker>
-    <style>
-      .canvas { fill: #ffffff; stroke: #e2e8f0; }
-      .node { fill: #f8fafc; stroke: #cbd5e1; }
-      .engine { fill: #ffffff; stroke: #cbd5e1; }
-      .router { fill: #f5f3ff; stroke: #c4b5fd; }
-      .scope { fill: #f8fafc; stroke: #cbd5e1; }
-      .cache { fill: #fff7ed; stroke: #fdba74; }
-      .payload { fill: #d97706; }
-      .loader { fill: #f1f5f9; stroke: #cbd5e1; }
-      .mapping { fill: #fffbeb; stroke: #f59e0b; }
-      .hook { fill: #eff6ff; stroke: #60a5fa; }
-      .gpu { fill: #f0fdf4; stroke: #86efac; }
-      .device { fill: #ffffff; stroke: #86efac; }
-      .slot-current { fill: #16a34a; }
-      .slot-next { fill: #bbf7d0; stroke: #22c55e; }
-      .no-collective { fill: #fef2f2; stroke: #fca5a5; }
-      .title { font: 700 25px system-ui, -apple-system, Segoe UI, sans-serif; fill: #0f172a; }
-      .subtitle { font: 500 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
-      .section { font: 700 15px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
-      .section-sub { font: 500 11.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
-      .eyebrow { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; letter-spacing: .8px; fill: #64748b; }
-      .box-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
-      .box-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
-      .payload-title { font: 700 14px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
-      .payload-text { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
-      .device-title { font: 700 9px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
-      .slot-text { font: 700 8px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
-      .slot-next-text { font: 700 8px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
-      .arrow-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
-      .request-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
-      .mapping-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b45309; }
-      .no-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b91c1c; }
-      .no-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #7f1d1d; }
-      .request-path { fill: none; stroke: #7c3aed; stroke-width: 2; marker-end: url(#request-arrow); }
-      .mapping-path { fill: none; stroke: #d97706; stroke-width: 2; marker-end: url(#mapping-arrow); }
-      .runtime-path { fill: none; stroke: #2563eb; stroke-width: 2; marker-end: url(#runtime-arrow); }
-      .startup-path { fill: none; stroke: #94a3b8; stroke-width: 1.5; stroke-dasharray: 4 3; marker-end: url(#runtime-arrow); }
-    </style>
-  </defs>
-
-  <rect x="12" y="12" width="1256" height="826" rx="14" class="canvas"/>
-
-  <text id="title-text" x="40" y="50" class="title">No-AllGather DLO: TP-coordinate host sharing, independent engines</text>
-  <text x="40" y="74" class="subtitle">TP2 shown: each engine owns two devices; equivalent workers share one host payload per TP coordinate</text>
-
-  <!-- The control plane remains outside DLO. -->
-  <rect x="390" y="94" width="500" height="68" rx="10" class="router"/>
-  <text x="640" y="121" text-anchor="middle" class="section">External router / orchestrator</text>
-  <text x="640" y="144" text-anchor="middle" class="section-sub">independent request placement &#183; outside vLLM-Omni DLO</text>
-
-  <!-- One node contains independent engine runtimes over one shared physical backing. -->
-  <rect x="30" y="184" width="1220" height="634" rx="14" class="node"/>
-  <text x="52" y="211" class="eyebrow">ONE NODE &#183; PROCESS-LOCAL ENGINES OVER SHARED PHYSICAL HOST MEMORY</text>
-
-  <!-- Independent request routing from the external orchestrator. -->
-  <path d="M 500 162 C 456 181, 400 204, 332 232" class="request-path"/>
-  <text x="416" y="194" class="request-label">request A</text>
-  <path d="M 780 162 C 824 181, 880 204, 948 232" class="request-path"/>
-  <text x="864" y="194" class="request-label">request B</text>
-
-  <!-- Independent engine A. -->
-  <rect x="45" y="232" width="575" height="250" rx="12" class="engine"/>
-  <text x="68" y="261" class="section">Independent engine A</text>
-  <text x="68" y="281" class="section-sub">TP2 &#183; two rank-local workers &#183; own request schedule and DLO state</text>
-
-  <rect x="68" y="310" width="98" height="130" rx="8" class="loader"/>
-  <text x="117" y="335" text-anchor="middle" class="eyebrow">STARTUP &#215;2</text>
-  <text x="117" y="359" text-anchor="middle" class="box-title">rank-local</text>
-  <text x="117" y="377" text-anchor="middle" class="box-title">loaders</text>
-  <text x="117" y="402" text-anchor="middle" class="box-text">final transforms</text>
-  <text x="117" y="421" text-anchor="middle" class="box-text">publish / join</text>
-
-  <rect x="180" y="310" width="112" height="130" rx="8" class="mapping"/>
-  <text x="236" y="335" text-anchor="middle" class="eyebrow">HOST VIEWS &#215;2</text>
-  <text x="236" y="359" text-anchor="middle" class="box-title">TP0 mmap</text>
-  <text x="236" y="377" text-anchor="middle" class="box-title">TP1 mmap</text>
-  <text x="236" y="402" text-anchor="middle" class="box-text">process-local</text>
-  <text x="236" y="421" text-anchor="middle" class="box-text">optional registration</text>
-
-  <rect x="306" y="310" width="120" height="130" rx="8" class="hook"/>
-  <text x="366" y="335" text-anchor="middle" class="eyebrow">RUNTIME &#215;2</text>
-  <text x="366" y="359" text-anchor="middle" class="box-title">DLO backends</text>
-  <text x="366" y="377" text-anchor="middle" class="box-title">+ block hooks</text>
-  <text x="366" y="402" text-anchor="middle" class="box-text">private streams</text>
-  <text x="366" y="421" text-anchor="middle" class="box-text">events + state</text>
-
-  <rect x="440" y="310" width="157" height="130" rx="8" class="gpu"/>
-  <text x="518" y="330" text-anchor="middle" class="eyebrow">TWO DEVICES</text>
-  <rect x="450" y="340" width="137" height="42" rx="6" class="device"/>
-  <text x="518" y="353" text-anchor="middle" class="device-title">TP0 &#183; GPU A0</text>
-  <rect x="457" y="359" width="58" height="16" rx="4" class="slot-current"/>
-  <text x="486" y="370" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
-  <rect x="521" y="359" width="58" height="16" rx="4" class="slot-next"/>
-  <text x="550" y="370" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
-  <rect x="450" y="389" width="137" height="42" rx="6" class="device"/>
-  <text x="518" y="402" text-anchor="middle" class="device-title">TP1 &#183; GPU A1</text>
-  <rect x="457" y="408" width="58" height="16" rx="4" class="slot-current"/>
-  <text x="486" y="419" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
-  <rect x="521" y="408" width="58" height="16" rx="4" class="slot-next"/>
-  <text x="550" y="419" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
-
-  <path d="M 166 375 L 180 375" class="startup-path"/>
-  <path d="M 292 375 L 306 375" class="runtime-path"/>
-  <path d="M 426 375 L 440 375" class="runtime-path"/>
-
-  <!-- Independent engine B. -->
-  <rect x="660" y="232" width="575" height="250" rx="12" class="engine"/>
-  <text x="683" y="261" class="section">Independent engine B</text>
-  <text x="683" y="281" class="section-sub">TP2 &#183; two rank-local workers &#183; own request schedule and DLO state</text>
-
-  <rect x="683" y="310" width="98" height="130" rx="8" class="loader"/>
-  <text x="732" y="335" text-anchor="middle" class="eyebrow">STARTUP &#215;2</text>
-  <text x="732" y="359" text-anchor="middle" class="box-title">rank-local</text>
-  <text x="732" y="377" text-anchor="middle" class="box-title">loaders</text>
-  <text x="732" y="402" text-anchor="middle" class="box-text">final transforms</text>
-  <text x="732" y="421" text-anchor="middle" class="box-text">publish / join</text>
-
-  <rect x="795" y="310" width="112" height="130" rx="8" class="mapping"/>
-  <text x="851" y="335" text-anchor="middle" class="eyebrow">HOST VIEWS &#215;2</text>
-  <text x="851" y="359" text-anchor="middle" class="box-title">TP0 mmap</text>
-  <text x="851" y="377" text-anchor="middle" class="box-title">TP1 mmap</text>
-  <text x="851" y="402" text-anchor="middle" class="box-text">process-local</text>
-  <text x="851" y="421" text-anchor="middle" class="box-text">optional registration</text>
-
-  <rect x="921" y="310" width="120" height="130" rx="8" class="hook"/>
-  <text x="981" y="335" text-anchor="middle" class="eyebrow">RUNTIME &#215;2</text>
-  <text x="981" y="359" text-anchor="middle" class="box-title">DLO backends</text>
-  <text x="981" y="377" text-anchor="middle" class="box-title">+ block hooks</text>
-  <text x="981" y="402" text-anchor="middle" class="box-text">private streams</text>
-  <text x="981" y="421" text-anchor="middle" class="box-text">events + state</text>
-
-  <rect x="1055" y="310" width="157" height="130" rx="8" class="gpu"/>
-  <text x="1133" y="330" text-anchor="middle" class="eyebrow">TWO DEVICES</text>
-  <rect x="1065" y="340" width="137" height="42" rx="6" class="device"/>
-  <text x="1133" y="353" text-anchor="middle" class="device-title">TP0 &#183; GPU B0</text>
-  <rect x="1072" y="359" width="58" height="16" rx="4" class="slot-current"/>
-  <text x="1101" y="370" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
-  <rect x="1136" y="359" width="58" height="16" rx="4" class="slot-next"/>
-  <text x="1165" y="370" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
-  <rect x="1065" y="389" width="137" height="42" rx="6" class="device"/>
-  <text x="1133" y="402" text-anchor="middle" class="device-title">TP1 &#183; GPU B1</text>
-  <rect x="1072" y="408" width="58" height="16" rx="4" class="slot-current"/>
-  <text x="1101" y="419" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
-  <rect x="1136" y="408" width="58" height="16" rx="4" class="slot-next"/>
-  <text x="1165" y="419" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
-
-  <path d="M 781 375 L 795 375" class="startup-path"/>
-  <path d="M 907 375 L 921 375" class="runtime-path"/>
-  <path d="M 1041 375 L 1055 375" class="runtime-path"/>
-
-  <!-- The runtimes are deliberately not coupled by a DLO collective. -->
-  <rect x="365" y="516" width="550" height="44" rx="8" class="no-collective"/>
-  <circle cx="395" cy="538" r="12" fill="#dc2626"/>
-  <path d="M 389 532 L 401 544 M 401 532 L 389 544" stroke="#ffffff" stroke-width="2.2"/>
-  <text x="420" y="534" class="no-title">No cross-engine DLO AllGather</text>
-  <text x="420" y="551" class="no-text">shared host pages do not couple request order, progress, or failure domains</text>
-  <text x="640" y="592" text-anchor="middle" class="section-sub">Ordinary TP / SP model collectives, when configured, remain inside each engine</text>
-
-  <!-- Each engine maps both TP-coordinate payloads from the shared host cache. -->
-  <path d="M 455 650 C 412 625, 340 610, 340 570 L 340 525 C 340 478, 278 454, 236 440" class="mapping-path"/>
-  <text x="286" y="617" text-anchor="middle" class="mapping-label">TP0 + TP1 mappings</text>
-  <path d="M 825 650 C 868 625, 940 610, 940 570 L 940 525 C 940 478, 893 454, 851 440" class="mapping-path"/>
-  <text x="994" y="617" text-anchor="middle" class="mapping-label">TP0 + TP1 mappings</text>
-
-  <!-- One node-local physical backing, visually beneath its consumers. -->
-  <rect x="58" y="650" width="242" height="140" rx="10" class="scope"/>
-  <text x="78" y="676" class="eyebrow">SHARING SCOPE</text>
-  <text x="78" y="702" class="box-title">TP2: two coordinate entries</text>
-  <text x="78" y="725" class="box-text">one physical copy per coordinate</text>
-  <text x="78" y="746" class="box-text">A/TP0 + B/TP0 &#8594; TP0 entry</text>
-  <text x="78" y="767" class="box-text">A/TP1 + B/TP1 &#8594; TP1 entry</text>
-
-  <rect x="330" y="650" width="620" height="140" rx="10" class="cache"/>
-  <text x="350" y="676" class="eyebrow">NODE-LOCAL IMMUTABLE HOST WEIGHT CACHE &#183; ONE COPY PER TP COORDINATE</text>
-  <rect x="370" y="692" width="260" height="60" rx="8" class="payload"/>
-  <text x="500" y="717" text-anchor="middle" class="payload-title">TP0 physical payload</text>
-  <text x="500" y="738" text-anchor="middle" class="payload-text">shared by A/TP0 + B/TP0</text>
-  <rect x="650" y="692" width="260" height="60" rx="8" class="payload"/>
-  <text x="780" y="717" text-anchor="middle" class="payload-title">TP1 physical payload</text>
-  <text x="780" y="738" text-anchor="middle" class="payload-text">shared by A/TP1 + B/TP1</text>
-  <text x="640" y="773" text-anchor="middle" class="box-text">Identity includes final layout + TP coordinate; engine and DP replica are excluded</text>
-
-  <rect x="980" y="650" width="242" height="140" rx="10" class="scope"/>
-  <text x="1000" y="676" class="eyebrow">WHAT STAYS PRIVATE</text>
-  <text x="1000" y="702" class="box-title">Per worker / process</text>
-  <text x="1000" y="725" class="box-text">mmap object and validation state</text>
-  <text x="1000" y="746" class="box-text">optional host registration</text>
-  <text x="1000" y="767" class="box-text">hooks, streams, events, GPU slots</text>
-</svg>

docs/models/supported_models.md

@@ -39,6 +39,7 @@ th {
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
 | `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   | — |

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -1,127 +1,156 @@
 # Distributed Layerwise Offloading
 
-Distributed layerwise offloading (DLO) streams DiT blocks through two device
-buffers instead of keeping the complete DiT in HBM. It supports sharded host
-weights with AllGather or independently streamed rank-local weights. The
-no-AllGather path can optionally share node-local host weights through mmap.
+Distributed layerwise offloading (DLO) extends block streaming to multi-device
+deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
+the host weights and reconstructs each layer at runtime. Without AllGather,
+each rank streams a complete block independently. Compatible TP1 deployments
+can share checkpoint-backed host pages among processes on the same node;
+otherwise DLO streams the ordinary loader's rank-local tensors.
 
-See the [shared DLO component architecture](../../../design/feature/offloader/distributed_layerwise_offload.md#architecture)
-for ownership and lifecycle coordination. The mode-specific internals are
-documented separately in the
-[AllGather design](../../../design/feature/offloader/distributed_layerwise_offload.md#allgather-design)
-and
-[no-AllGather design](../../../design/feature/offloader/distributed_layerwise_offload.md#no-allgather-design).
+See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
+for the implementation contract and compatibility matrix.
 
-## Choose a mode
+## Execution model
 
-| Mode | Use when | Host weights | Runtime synchronization |
-| --- | --- | --- | --- |
-| DLO AllGather (default) | DP ranks execute the same block path in lockstep | About `1 / dp_size` per rank | DLO weight AllGather |
-| no-AllGather | Ranks or engines must schedule independently | Complete rank-local layout | No DLO weight collective |
-| no-AllGather + host weight cache | Equivalent independent workers share one node | Shared final mmap layout per TP coordinate | No DLO weight collective |
-| host weight cache + host registration | The platform supports registration and recurrent staging is too expensive | Shared registered mmap layout | No DLO weight collective |
+DLO overlaps three operations with a fixed two-block device buffer:
 
-AllGather is normally the best choice for synchronized DP. Host weight cache mode
-targets independently scheduled replicas, especially repeated TP engines on
-one node.
+```text
+Compute stream:  [Layer N]          [Layer N+1]        [Layer N+2]
+H2D stream:      [H2D shard N+1]    [H2D shard N+2]
+AllGather:       [AG N+1]           [AG N+2]
+Slots:           slot 0: Layer N    slot 1: Layer N+1
+```
+
+AllGather communicates only request-independent weight shards, so data-
+parallel ranks may process different requests concurrently.
 
 ## Usage
 
 ```bash
-# Sharded host weights with DLO AllGather
+# Four ranks with sharded host weights and AllGather
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4
 
-# Independently streamed rank-local weights
+# Standard-loader rank-local weights, without DLO AllGather
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
+  --data-parallel-size 4 \
   --dlo-no-use-allgather
 
-# Share final host weights; prefer registered direct H2D under the pinned-memory policy
+# Sequence parallel deployment
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
-  --dlo-no-use-allgather \
-  --dlo-use-host-weight-cache
-
-# Apply an 80 GiB per-worker safety ceiling to the same registration attempt
-vllm serve /path/to/model --omni \
-  --enable-distributed-layerwise-offload \
-  --tensor-parallel-size 2 \
-  --dlo-no-use-allgather \
-  --dlo-use-host-weight-cache \
-  --dlo-host-weight-cache-pin-limit-gib 80
+  --usp 4
 ```
 
-## Relevant options
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="/path/to/model",
+    enable_distributed_layerwise_offload=True,
+    dlo_use_allgather=True,
+)
+```
+
+## Flags
 
 | Flag | Meaning | Default |
 | --- | --- | --- |
 | `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
-| `--dlo-no-use-allgather` | Stream complete rank-local blocks independently | `false` |
-| `--dlo-use-host-weight-cache` | Use the final-layout cache and prefer registration when pinned memory is enabled | `false` |
-| `--dlo-host-weight-cache-pin-limit-gib GIB` | Optional per-worker registration ceiling; zero means no additional ceiling | `0` |
-| `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` |
+| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
+| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
+| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
+| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
 
-No-AllGather DLO automatically uses a compatible checkpoint mmap plan. When
-checkpoint mmap is unavailable, it preserves ordinary pinned loader weights by
-default. `--dlo-use-host-weight-cache` explicitly selects the transform-complete
-final-layout cache. Under the default `pin_cpu_memory=True` policy, DLO attempts
-to register the complete immutable mapping for direct H2D. Setting
-`pin_cpu_memory=False` through structured or programmatic configuration retains
-the staged mmap path instead. A positive pin limit is only an optional safety
-ceiling and does not select the cache. The cache uses
-`~/.cache/vllm-omni/dlo-host-weights`. Programmatic configuration may override
-`dlo_host_weight_cache_dir` and the writer lock timeout; these advanced storage
-controls are intentionally not separate CLI flags.
+## Host-weight loading
 
-## Operational notes
+The diffusion loader chooses host storage before DLO is enabled. It first
+attempts to build a complete, validated direct-checkpoint mmap plan. If names,
+coverage, shape, dtype, topology, or loader-callback compatibility cannot be
+proven, it runs the ordinary model loader instead. DLO consumes that result and
+does not make a second checkpoint-compatibility decision.
 
-- All workers that should share must see the same local, disk-backed cache
-  directory. Do not use tmpfs or a cross-node filesystem for this version.
-- Cache entries are immutable and validated before use. A cache or registration
-  failure keeps the ordinary loader weights or falls back to bounded staging.
-- The bounded-staging fallback is a memory-first mode: recurrent pageable-mmap to
-  pinned-slot copies can reduce throughput, especially across concurrent TP
-  engines. Use it intentionally by disabling pinned host memory, or when the
-  active platform cannot register the immutable mapping.
-- When a positive registration ceiling is configured, it must cover the
-  complete page-aligned mapping reported in the worker log. Registration is
-  all-or-nothing. Zero applies no additional ceiling.
-- CUDA is the first registration backend and requires read-only host
-  registration support for immutable cache mappings. Platforms or devices
-  without that capability continue to use bounded staging; pinned-memory policy
-  requests direct H2D but does not guarantee it.
-- Registration is process-local but does not duplicate the underlying file
-  pages. Each worker must still satisfy its platform and OS page-locking limits.
-- Shutdown unregisters host ranges before closing their mmap handles.
-- Host weight cache v1 has no automatic eviction. Stop all users of an entry
-  before deleting it.
+The shared-mmap optimization in this phase is supported only with TP1. TP
+greater than one falls back before model mutation to ordinary TP-aware loading.
+DLO may still consume those TP-local tensors, but this is a compatibility path:
+it does not share checkpoint-backed runtime weights across DP replicas and
+provides no shared-mmap host-memory guarantee.
 
-## Scheduling constraints
+The mmap plan skips only dedicated DiT weight sources. Other component sources,
+such as a text encoder loaded through the shared diffusion loader, continue to
+use their ordinary component loader. A checkpoint source that mixes DiT and
+non-DiT weights falls back completely rather than leaving an unplanned
+component uninitialized.
 
-AllGather ranks must request blocks in the same collective order. Concurrent DP
-requests may use different prompts, but they must follow the same denoising and
-block-execution path and set the same explicit `num_inference_steps`.
+With direct checkpoint mmap, the loader:
 
-No-AllGather workers do not have this DLO lockstep requirement. Ordinary TP or
-SP model collectives still synchronize ranks within each engine.
+1. saves non-persistent buffers such as RoPE frequencies;
+2. moves the normally created transformer to the meta device;
+3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
+4. applies any loader-owned bounded layout adapters while packing blocks;
+5. restores saved non-persistent buffers; and
+6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
+   hooks.
+
+For AllGather with a group larger than one, each process copies only its
+persistent shard and then releases the source mapping. For no-AllGather, each
+process keeps the mapping open and packs complete blocks through two bounded
+pinned staging slots. Processes mapping the same files on one node share the
+immutable pages; no-AllGather still performs a complete-block H2D copy in each
+process.
+
+When the effective DLO group size is one, `dlo_use_allgather=True` does not
+perform a collective and uses the same rank-local transfer behavior.
+
+## Declarative topology
+
+Models may declare an `OffloadPlan` instead of embedding offload logic:
+
+```python
+from vllm_omni.diffusion.offloader import OffloadPlan
+
+
+class MyPipeline(nn.Module):
+    _dit_modules = ["transformer"]
+    _offload_plan = OffloadPlan(
+        block_attrs={"transformer": ("blocks",)},
+        offload_submodules={"context_encoder": "layers"},
+    )
+```
+
+When no plan exists, discovery falls back to
+`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
+
+## Data-parallel concurrency
+
+With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
+up to `dp_size` requests per denoising step. Every concurrent request must set
+the same explicit `num_inference_steps`; `None` is rejected because every rank
+must enter each collective.
 
 ## Limitations
 
-- Direct checkpoint mmap currently requires TP1. The host weight cache can share
-  final ordinary-loader layouts between matching TP coordinates at TP1 or TP>1.
+- Direct checkpoint mmap currently requires TP1. TP greater than one is
+  outside the Phase A shared-mmap support scope and falls back before model
+  mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
+  but it provides no shared-mmap host-memory benefit or guarantee.
+- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
+  AllGather has limited end-to-end validation.
 - Per-tensor online FP8 linears use the ordinary loader and can run with either
-  DLO transfer path. Other online quantization methods require no-AllGather
-  until their runtime layouts are validated.
-- Host weight cache v1 rejects quantized, non-contiguous, aliased/tied,
-  device-only, HSDP, expert-parallel, CFG-parallel, and PP layouts.
-- HSDP plus DLO AllGather is unsupported, and host weight cache v1 rejects HSDP.
-- The host weight cache changes steady-state host backing, not startup loading:
-  each worker still performs ordinary loading and full content validation.
-- Skip-load-on-hit, cache eviction, and cross-node sharing remain follow-up
-  work in
-  [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+  DLO transfer path. With AllGather, every rank temporarily materializes the
+  complete FP8 model in host memory before DLO retains only its shard. Other
+  online quantization methods require no-AllGather until their runtime layouts
+  are validated.
+- Resident leading layers require `--dlo-no-use-allgather` and a model
+  `OffloadPlan` that declares eligible `resident_dit_paths`.
+- DP concurrency requires an explicit, identical inference-step count.
+
+Sharing transformed TP or quantized runtime layouts through a normalized mmap
+cache is a follow-up design in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
+the direct-checkpoint path.
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

docs/user_guide/diffusion/parallelism/hsdp.md

@@ -152,7 +152,8 @@ In `DiffusionParallelConfig`:
 **Constraints:**
 
 - `hsdp_replicate_size × hsdp_shard_size == world_size`
-- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism,
+  Pipeline Parallelism, or Expert Parallelism.
 
 ---
 

docs/user_guide/diffusion_features.md

@@ -168,6 +168,14 @@ The following tables show which models support each feature:
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **LongCat-Video-Avatar-1.5** |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |
+| **SANA-WM**                  |     ❌     |     ❌      |          ❌<sup>5</sup> |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+
+> Notes:
+> 5. SANA-WM cannot support sequence parallelism: its bidirectional gated delta
+>    recurrence carries state across frames, so a rank cannot denoise a slice of
+>    the token sequence in isolation. Doing so would need a distributed scan or
+>    an all-gather before every GDN block. The remaining ❌ columns are simply
+>    unvalidated on this model, not known-broken.
 
 > **Step execution note:** Helios supports single-request step execution only;
 > use `max_num_seqs=1`.
@@ -175,7 +183,7 @@ The following tables show which models support each feature:
 **Frame Interpolation Support**
 
 - **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
-- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM
 
 ### AudioGen
 

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -53,7 +53,6 @@ from typing import Any
 import torch
 from PIL import Image
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -139,12 +138,6 @@ async def main():
 
     # ---- Torch setup ----
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-    )
 
     # ---- Cache Config ----
     if args.cache_backend == "cache_dit":
@@ -171,7 +164,10 @@ async def main():
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
         diffusion_load_format="dummy",

examples/offline_inference/diffusion/lingbot_world_v2.py

@@ -154,8 +154,6 @@ def _positive_finite(value: float, flag: str) -> float:
 def build_omni_kwargs(
     args: argparse.Namespace,
     paths: LingBotPaths,
-    *,
-    parallel_config: Any,
 ) -> dict[str, Any]:
     """Build Omni configuration without overriding checkpoint class discovery."""
 
@@ -167,7 +165,7 @@ def build_omni_kwargs(
     return {
         "model": model,
         "flow_shift": flow_shift,
-        "parallel_config": parallel_config,
+        "tensor_parallel_size": args.tensor_parallel_size,
         "enforce_eager": args.enforce_eager,
         "model_config": {"lingbot_action_root": str(paths.action_root)},
     }
@@ -262,12 +260,10 @@ def main(argv: Sequence[str] | None = None) -> Path:
     # helper tests work on machines without the native vLLM runtime.
     from diffusers.utils import export_to_video
 
-    from vllm_omni.diffusion.data import DiffusionParallelConfig
     from vllm_omni.entrypoints.omni import Omni
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-    parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
-    omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
+    omni_kwargs = build_omni_kwargs(args, paths)
     omni = Omni(**omni_kwargs)
     try:
         outputs = omni.generate(

examples/offline_inference/diffusion/lingbot_world_v2_realtime.py

@@ -125,7 +125,6 @@ async def run(argv: Sequence[str] | None = None) -> Path:
     # require a CUDA-enabled vLLM installation.
     import torch
 
-    from vllm_omni.diffusion.data import DiffusionParallelConfig
     from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
     from vllm_omni.entrypoints.async_omni import AsyncOmni
     from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
@@ -141,7 +140,7 @@ async def run(argv: Sequence[str] | None = None) -> Path:
         model=args.model,
         engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
         enforce_eager=args.enforce_eager,
-        parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
+        tensor_parallel_size=args.tensor_parallel_size,
         max_num_seqs=1,
         model_config={
             "ar_diffusion_height": args.height,

examples/offline_inference/image_to_image/image_edit.py

@@ -96,7 +96,6 @@ from typing import Any
 import torch
 from PIL import Image
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
@@ -452,14 +451,6 @@ def main():
 
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        enable_expert_parallel=args.enable_expert_parallel,
-    )
-
     # Configure cache based on backend type
     cache_config = None
     if args.cache_backend == "cache_dit":
@@ -490,7 +481,11 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        enable_expert_parallel=args.enable_expert_parallel,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,

examples/offline_inference/image_to_video/image_to_video.py

@@ -56,7 +56,6 @@ import numpy as np
 import PIL.Image
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -503,17 +502,6 @@ def main():
         }
 
     profiler_enabled = args.profiler_config is not None
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        vae_patch_parallel_size=args.vae_patch_parallel_size,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-        pipeline_parallel_size=args.pipeline_parallel_size,
-    )
     omni_kwargs = dict(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
@@ -524,7 +512,15 @@ def main():
         diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
         diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
         enable_cpu_offload=args.enable_cpu_offload,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
         enforce_eager=args.enforce_eager,
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,

examples/offline_inference/longcat_video/end2end.py

@@ -15,7 +15,6 @@ import imageio
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import (
     prepare_longcat_video_avatar_model_for_omni,
 )
@@ -236,14 +235,6 @@ def main():
         dtype="bfloat16",
         enforce_eager=args.enforce_eager,
         additional_config=additional_config,
-        parallel_config=DiffusionParallelConfig(
-            ulysses_degree=1,
-            ring_degree=1,
-            cfg_parallel_size=1,
-            tensor_parallel_size=1,
-            vae_patch_parallel_size=1,
-            pipeline_parallel_size=1,
-        ),
     )
 
     audio_input: str | dict[str, str] | None = args.audio

examples/offline_inference/speech_to_video/speech_to_video.py

@@ -23,7 +23,6 @@ from typing import Any
 
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -250,17 +249,6 @@ def main():
             "scm_steps_policy": "dynamic",
         }
 
-    parallel_config = DiffusionParallelConfig(
-        tensor_parallel_size=args.tensor_parallel_size,
-        cfg_parallel_size=args.cfg_parallel_size,
-        vae_patch_parallel_size=args.vae_patch_parallel_size,
-        ring_degree=args.ring_degree,
-        ulysses_degree=args.ulysses_degree,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-    )
-
     profiler_enabled = args.profiler_config is not None
 
     omni_kwargs = dict(
@@ -271,7 +259,14 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        parallel_config=parallel_config,
+        tensor_parallel_size=args.tensor_parallel_size,
+        cfg_parallel_size=args.cfg_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
+        ring_degree=args.ring_degree,
+        ulysses_degree=args.ulysses_degree,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
         enforce_eager=args.enforce_eager,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         cache_backend=args.cache_backend,

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -38,7 +38,6 @@ from typing import Any
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -306,17 +305,13 @@ def main():
     print(f"  Seed: {args.seed}")
     print(f"{'=' * 60}\n")
 
-    parallel_config = DiffusionParallelConfig(
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-    )
-
     # Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio
     # is resolved from its config.
     omni_kwargs: dict[str, Any] = dict(
         model=args.model,
-        parallel_config=parallel_config,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,

examples/offline_inference/text_to_speech/ming_tts/end2end.py

@@ -13,6 +13,7 @@ import yaml
 from transformers import AutoTokenizer
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
 from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
     KEY_CFG,
     KEY_MAX_DECODE_STEPS,
@@ -21,7 +22,6 @@ from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
     SAMPLE_RATE,
 )
 from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
-from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
 
 try:
     from .runner import (
@@ -142,7 +142,8 @@ def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
         return _load_speaker_embedding(args.speaker_embedding)
     if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
         return None
-    embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
+    # Runs before the engine is built, so nothing has pre-fetched campplus.onnx.
+    embs = SpeakerEmbeddingExtractor(args.model, allow_download=True).extract_many(paths)
     if not embs:
         raise RuntimeError("Speaker extraction produced no embeddings")
     return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)

examples/offline_inference/text_to_video/text_to_video.py

@@ -10,7 +10,6 @@ from typing import Any
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -462,17 +461,6 @@ def main():
             "scm_steps_policy": "dynamic",
         }
 
-    # Configure parallel settings
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        vae_patch_parallel_size=args.vae_patch_parallel_size,
-        pipeline_parallel_size=args.pipeline_parallel_size,
-        enable_expert_parallel=args.enable_expert_parallel,
-    )
-
     profiler_enabled = args.profiler_config is not None
 
     omni_kwargs = dict(
@@ -481,7 +469,13 @@ def main():
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,
-        parallel_config=parallel_config,
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
+        enable_expert_parallel=args.enable_expert_parallel,
         enforce_eager=args.enforce_eager,
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -11,7 +11,6 @@ import numpy as np
 from PIL import Image
 from vllm.multimodal.media.audio import load_audio
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -197,14 +196,6 @@ def main() -> None:
             },
         )
 
-    parallel_config = DiffusionParallelConfig(
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-    )
-
     cache_config = None
     if args.cache_backend == "cache_dit":
         cache_config = {
@@ -222,7 +213,11 @@ def main() -> None:
 
     omni_kwargs = dict(
         model=args.model,
-        parallel_config=parallel_config,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
         model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,

recipes/NVIDIA/SANA-WM.md

@@ -0,0 +1,193 @@
+# SANA-WM
+
+> Camera-controllable first-frame image-to-video world model.
+
+## Summary
+
+- Vendor: Efficient-Large-Model / NVlabs SANA
+- Model: `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` (standard diffusers layout, converted offline from the NVlabs release; Stage-1 transformer + VAE only)
+- Task: First-frame image-to-video generation with camera control
+- Mode: Online serving with the OpenAI-compatible video API
+- Model weights: about 13 GB for the Stage-1 transformer (10 GB) and VAE (2.3 GB)
+- Local disk: reserve about 40 GB for the Hugging Face cache and runtime artifacts
+- Recommended GPU: 24 GB or larger CUDA GPU
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to serve SANA-WM through `/v1/videos` or
+`/v1/videos/sync`. The model takes a text prompt, a first-frame image, and
+either an action DSL string or explicit camera poses. vLLM-Omni serves the
+SANA-WM Stage-1 DiT, decoded through the SANA VAE. The optional LTX-2 refiner
+stage is not supported by this integration; it is a planned follow-up.
+
+## References
+
+- Upstream model card: <https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional>
+- Video API: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA RTX PRO 6000 Blackwell 96GB
+
+#### Capacity
+
+- Model storage: the Stage-1 transformer is about 10 GB and the VAE about
+  2.3 GB. The Gemma text encoder is a separate 4.9 GB repo. The engine prefetches
+  the whole model repo at startup (`allow_patterns=["*"]`), which is why the
+  Stage-1 weights live in their own repo — the two-stage one carries an
+  additional 84 GB `refiner/` that this path never loads.
+- Text encoder: the pipeline tries `google/gemma-2-2b-it` first, then falls back
+  to the ungated mirror `Efficient-Large-Model/gemma-2-2b-it`. The first repo is
+  gated, so without an accepted licence and `hf auth login` you will see one
+  failed load before the fallback succeeds — that is expected, not an error to
+  chase. Set `VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER` to pin a specific repo or a
+  local path and skip both.
+- Disk sizing: provision about 40 GB of local disk or Hugging Face cache volume
+  so the model, temporary downloads, and generated artifacts fit without cache
+  eviction.
+- GPU sizing: the default 1280x704, 161-frame, 60-step serving profile peaks at
+  22.6 GB of device memory and takes about 133 s to generate on one RTX PRO
+  6000 Blackwell. The peak lands in the VAE decode, which is why the pipeline
+  forces VAE tiling on regardless of `vae_use_tiling` — without it the same
+  request costs about 9 GB more, and a 321-frame one OOMs outright. On smaller
+  GPUs, lower `width`, `height`, or `num_frames` before serving production
+  requests.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA driver with CUDA runtime supported by your PyTorch
+  build
+- Recommended operator library: Triton, installed through the vLLM/vLLM-Omni
+  Python environment
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+The repo ships the standard Diffusers layout (`model_index.json` +
+`transformer/`, `vae/`), and its `model_index.json` names `SanaWmPipeline`, so
+the pipeline class resolves on its own.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+vllm serve BBBBruce/SANA-WM_bidirectional-stage1-diffusers \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091
+```
+
+No deploy config: single-stage diffusion models are deliberately absent from
+`OMNI_PIPELINES` (`vllm_omni/config/pipeline_registry.py`), so stage resolution
+falls back to the default stage config and a YAML's stage settings — including
+`default_sampling_params` — would not be applied. The production generation
+settings therefore live in the model (`num_inference_steps=60`,
+`guidance_scale=5.0`); a request that omits a field gets them. The examples
+below still pass every field explicitly so the numbers are visible.
+
+If you point this at the older two-stage repo
+(`BBBBruce/SANA-WM_bidirectional-diffusers`), startup fails with `Model class
+SanaWmTwoStagesPipeline not found in diffusion model registry`, because that
+repo's `model_index.json` names a class this build does not register. Add
+`--model-class-name SanaWmPipeline` to override it.
+
+#### Verification
+
+Use a short smoke request first:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "prompt=A slow forward camera move through a quiet city street." \
+  -F "negative_prompt=blurry, low quality, distorted, watermark" \
+  -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+  -F "width=1280" \
+  -F "height=704" \
+  -F "num_frames=9" \
+  -F "fps=16" \
+  -F "num_inference_steps=2" \
+  -F "guidance_scale=5.0" \
+  -F "seed=42" \
+  --form-string 'extra_params={"sana_wm":{"action":"w-8","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+  -o sana_wm_smoke.mp4
+```
+
+For a production-length request, note that the action durations must sum to
+`num_frames - 1` — the rollout includes the identity start pose — and a
+mismatch is rejected rather than padded or truncated:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "prompt=A slow forward camera move through a quiet city street." \
+  -F "negative_prompt=blurry, low quality, distorted, watermark" \
+  -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+  -F "width=1280" \
+  -F "height=704" \
+  -F "num_frames=161" \
+  -F "fps=16" \
+  -F "num_inference_steps=60" \
+  -F "guidance_scale=5.0" \
+  -F "seed=42" \
+  --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+  -o sana_wm_output.mp4
+```
+
+Use `POST /v1/videos` instead when you want job storage and polling rather than
+inline MP4 bytes. It accepts the same form fields as `/v1/videos/sync`.
+
+```bash
+create_response=$(curl -sS -X POST http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A slow forward camera move through a quiet city street." \
+  -F "negative_prompt=blurry, low quality, distorted, watermark" \
+  -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+  -F "width=1280" \
+  -F "height=704" \
+  -F "num_frames=161" \
+  -F "fps=16" \
+  -F "num_inference_steps=60" \
+  -F "guidance_scale=5.0" \
+  -F "seed=42" \
+  --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}')
+
+video_id=$(echo "$create_response" | jq -r '.id')
+curl -sS "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o sana_wm_output.mp4
+```
+
+#### Notes
+
+- Sequence parallelism is not supported. The bidirectional gated delta
+  recurrence carries state across frames, so a rank cannot denoise a slice
+  of the token sequence in isolation; supporting it needs a distributed scan
+  or an all-gather before the GDN blocks.
+
+- `input_reference` is required for the first frame. Use `image_reference` only
+  when you need a JSON-safe image URL or data URL instead of a multipart upload.
+- `sana_wm` must provide exactly one of `action` or `camera`.
+- Action strings use comma-separated `<keys>-<duration>` segments. Supported
+  keys are `w`, `a`, `s`, `d` for translation and `i`, `j`, `k`, `l` for
+  pitch/yaw rotation. The durations must sum to `num_frames - 1`.
+- Explicit camera control (alternative to `action`): pass
+  `"camera": {"poses": [...]}` where `poses` is a list of `num_frames`
+  camera-to-world 4x4 matrices (row-major, OpenCV `+X right, +Y down, +Z forward`
+  convention), e.g.
+  `extra_params={"sana_wm":{"camera":{"poses":[[[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], ...]},"intrinsics":{...}}}`.
+  Most callers should prefer `action`; explicit poses exist for callers that
+  already have a per-frame trajectory.
+- Explicit `intrinsics` are recommended and take the mapping form
+  `{"fx":640,"fy":640,"cx":640,"cy":352}` (for 1280x704). This `{fx,fy,cx,cy}`
+  mapping is the only accepted intrinsics form; omit `intrinsics` to derive them
+  from the output resolution. All four values must be finite, and `fx`/`fy`
+  must be positive — the ray map divides by them.
+- The video API returns decoded MP4 bytes and has no `output_type` field, so
+  raw Stage-1 latents are reachable only from the offline API
+  (`OmniDiffusionSamplingParams(output_type="latent")`); see
+  [`tests/e2e/offline_inference/test_sana_wm.py`](../../tests/e2e/offline_inference/test_sana_wm.py).
+  Putting `output_type` in `extra_params` does not work: it lands in
+  `sampling_params.extra_args`, while the pipeline reads the top-level field.

recipes/README.md

@@ -70,6 +70,7 @@ recipes/
 | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
 | [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
 | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
+| [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |

recipes/inclusionAI/Ming-omni-tts.md

@@ -5,20 +5,21 @@
 ## Summary
 
 - Vendor: inclusionAI
-- Model: `inclusionAI/Ming-omni-tts-0.5B`
+- Model: `inclusionAI/Ming-omni-tts-16.8B-A3B`, `inclusionAI/Ming-omni-tts-0.5B`
 - Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
 - Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API
 - Maintainer: Community
 
 ## References
 
-- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
+- [Huggingface: inclusionAI/Ming-omni-tts-16.8B-A3B](https://huggingface.co/inclusionAI/Ming-omni-tts-16.8B-A3B)
+- [Huggingface: inclusionAI/Ming-omni-tts-0.5B](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
 - Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts)
 
 
 ## Hardware Support
 
-This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment.
+This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B and the MoE 16.8B two-stage TTS pipeline deployment.
 Other hardware is welcome as community validation lands.
 
 ## CUDA
@@ -39,11 +40,7 @@ Other hardware is welcome as community validation lands.
 Launch the two-stage talker:
 
 ```bash
-vllm serve inclusionAI/Ming-omni-tts-0.5B \
-    --deploy-config vllm_omni/deploy/ming_tts.yaml \
-    --omni \
-    --port 8091 \
-    --enforce-eager
+vllm serve inclusionAI/Ming-omni-tts-16.8B-A3B --omni --port 8091
 ```
 
 #### Verification
@@ -54,7 +51,7 @@ Basic synthesis (save the WAV bytes):
 curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
-      "model": "inclusionAI/Ming-omni-tts-0.5B",
+      "model": "inclusionAI/Ming-omni-tts-16.8B-A3B",
       "input": "你好，这是 Ming 在线语音合成测试。",
       "response_format": "wav"
     }' --output ming_tts_basic.wav
@@ -74,7 +71,7 @@ Reference-audio zero-shot cloning:
 REF_AUDIO="$BASE/10002287-00000094.wav"
 jq -n --arg ref_audio "$REF_AUDIO" \
     '{
-       model: "inclusionAI/Ming-omni-tts-0.5B",
+       model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
        input: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。",
        ref_audio: $ref_audio,
        ref_text: "在此奉劝大家别乱打美白针。",
@@ -107,7 +104,7 @@ jq -n \
     --arg a "$REF_A" \
     --arg b "$REF_B" \
     '{
-       model: "inclusionAI/Ming-omni-tts-0.5B",
+       model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
        input: $input,
        ref_audio: [$a, $b],
        ref_text: $ref_text,

tests/config/test_omni_config.py

@@ -733,8 +733,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
 
 def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
     cfg = OmniStageDiffusionParallelConfig(
-        pipeline_parallel_size=2,
-        ulysses_degree=2,
+        ulysses_degree=4,
         use_hsdp=True,
         hsdp_shard_size=-1,
         hsdp_replicate_size=2,
@@ -745,10 +744,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
 
 
 def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with TP"):
         OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with DP"):
         OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
 
@@ -908,52 +907,6 @@ def test_structured_diffusion_config_rejects_invalid_compile_granularity():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
 
 
-@pytest.mark.parametrize(
-    ("enable_dlo", "use_allgather"),
-    [(False, False), (True, True)],
-)
-@pytest.mark.parametrize(
-    ("field_name", "field_value"),
-    [
-        ("dlo_use_host_weight_cache", True),
-        ("dlo_host_weight_cache_pin_limit_gib", 1.0),
-    ],
-)
-def test_structured_diffusion_config_restricts_host_weight_cache_to_no_allgather_dlo(
-    enable_dlo,
-    use_allgather,
-    field_name,
-    field_value,
-):
-    with pytest.raises(ValueError, match="requires no-AllGather distributed layerwise offload"):
-        omni_config_module._DiffusionConfigProjection(
-            enable_distributed_layerwise_offload=enable_dlo,
-            dlo_use_allgather=use_allgather,
-            **{field_name: field_value},
-        )
-
-
-def test_structured_diffusion_config_requires_cache_for_registration_ceiling():
-    with pytest.raises(ValueError, match="requires dlo_use_host_weight_cache=True"):
-        omni_config_module._DiffusionConfigProjection(
-            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=False,
-            dlo_host_weight_cache_pin_limit_gib=1.0,
-        )
-
-
-@pytest.mark.parametrize(
-    "field_name",
-    [
-        "dlo_host_weight_cache_lock_timeout",
-        "dlo_host_weight_cache_pin_limit_gib",
-    ],
-)
-def test_structured_diffusion_config_rejects_non_finite_host_weight_cache_values(field_name):
-    with pytest.raises(ValueError, match=rf"{field_name} must be finite"):
-        omni_config_module._DiffusionConfigProjection(**{field_name: float("inf")})
-
-
 def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage():
     pipeline = _resolve_pipeline_or_skip("qwen3_tts")
     legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]

tests/dfx/conftest.py

@@ -653,6 +653,7 @@ def run_benchmark(
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
     resource_label: str | None = None,
+    num_warmups: int = 2,
 ) -> dict[str, Any]:
     """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
 
@@ -677,7 +678,7 @@ def run_benchmark(
         + args
         + [
             "--num-warmups",
-            "2",
+            str(num_warmups),
             "--save-result",
             "--result-dir",
             os.environ.get("BENCHMARK_DIR", "tests"),

tests/dfx/perf/scripts/run_benchmark.py

@@ -279,5 +279,6 @@ def test_performance_benchmark(omni_server, benchmark_params):
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
             resource_label=resource_label,
+            num_warmups=max(2, int(concurrency)),
         )
         assert_result(result, params, num_prompt)

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -74,14 +74,6 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
 
 
-@pytest.fixture(autouse=True)
-def mock_dit_group(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
-        lambda: None,
-    )
-
-
 @pytest.fixture(autouse=True)
 def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -92,9 +92,8 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
     monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -229,9 +228,8 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
     monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -259,6 +257,33 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
     ]
 
 
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch):
+    """A destroyed diffusion PP group must not block the next initialization."""
+
+    class _DestroyableGroup:
+        def __init__(self) -> None:
+            self.destroy_calls = 0
+
+        def destroy(self) -> None:
+            self.destroy_calls += 1
+
+    pipeline_group = _DestroyableGroup()
+    for name in ("_DP", "_CFG", "_SP"):
+        monkeypatch.setattr(parallel_state, name, None)
+    monkeypatch.setattr(parallel_state, "_PP", pipeline_group)
+    monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None)
+    for name in ("_DP", "_PCP", "_TP", "_EP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+    monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False)
+
+    parallel_state.destroy_model_parallel()
+
+    assert pipeline_group.destroy_calls == 1
+    assert parallel_state.vllm_parallel_state._PP is None
+
+
 @pytest.mark.cpu
 @pytest.mark.core_model
 def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
@@ -295,9 +320,8 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
-    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
         monkeypatch.setattr(parallel_state, name, None)
     for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
         monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)

tests/diffusion/distributed/test_hsdp.py

@@ -234,7 +234,7 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_tp(self):
         """Test that HSDP and Tensor Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        with pytest.raises(ValueError, match="not compatible with TP"):
             DiffusionParallelConfig(
                 tensor_parallel_size=2,
                 use_hsdp=True,
@@ -243,7 +243,7 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_dp(self):
         """Test that HSDP and Data Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+        with pytest.raises(ValueError, match="not compatible with DP"):
             DiffusionParallelConfig(
                 data_parallel_size=2,
                 use_hsdp=True,
@@ -264,158 +264,6 @@ class TestDiffusionParallelConfigHSDP:
         assert config.hsdp_shard_size == 2  # auto: 4 // 2
 
 
-class TestStandaloneHSDPDetection:
-    """Tests for standalone HSDP detection and dit_parallel_size calculation.
-
-    These tests verify the logic used in initialize_model_parallel() to detect
-    standalone HSDP mode and calculate effective parallel sizes.
-
-    Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
-    """
-
-    @staticmethod
-    def compute_standalone_hsdp_params(
-        data_parallel_size: int = 1,
-        cfg_parallel_size: int = 1,
-        sequence_parallel_size: int = 1,
-        pipeline_parallel_size: int = 1,
-        tensor_parallel_size: int = 1,
-        fully_shard_degree: int = 1,
-        hsdp_replicate_size: int = 1,
-    ) -> dict:
-        """Compute standalone HSDP detection parameters.
-
-        This mirrors the logic in initialize_model_parallel().
-        """
-        dit_parallel_size = (
-            data_parallel_size
-            * cfg_parallel_size
-            * sequence_parallel_size
-            * pipeline_parallel_size
-            * tensor_parallel_size
-        )
-
-        # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
-        is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
-
-        # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
-        if is_standalone_hsdp:
-            effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
-        else:
-            effective_dit_parallel_size = dit_parallel_size
-
-        effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
-
-        return {
-            "original_dit_parallel_size": dit_parallel_size,
-            "is_standalone_hsdp": is_standalone_hsdp,
-            "effective_dit_parallel_size": effective_dit_parallel_size,
-            "effective_dp_size": effective_dp_size,
-        }
-
-    def test_standalone_hsdp_basic(self):
-        """Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 4
-        assert result["effective_dp_size"] == 4
-
-    def test_standalone_hsdp_with_replicate(self):
-        """Test standalone HSDP with replication (shard_size=4, replicate=2)."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=2,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 8  # 4 * 2
-        assert result["effective_dp_size"] == 8
-
-    def test_combined_hsdp_sp_not_standalone(self):
-        """Test HSDP combined with SP is NOT detected as standalone.
-
-        This is a regression test for the bug where the condition
-        `dit_parallel_size == fully_shard_degree` incorrectly matched
-        combined modes like SP=4 + HSDP=4.
-        """
-        result = self.compute_standalone_hsdp_params(
-            sequence_parallel_size=4,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 4
-        assert result["is_standalone_hsdp"] is False
-        # Should NOT override dp_size for combined mode
-        assert result["effective_dp_size"] == 1  # original data_parallel_size
-
-    def test_combined_hsdp_cfg_not_standalone(self):
-        """Test HSDP combined with CFG is NOT detected as standalone."""
-        result = self.compute_standalone_hsdp_params(
-            cfg_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_combined_hsdp_pp_not_standalone(self):
-        """Test HSDP combined with PP is NOT detected as standalone."""
-        result = self.compute_standalone_hsdp_params(
-            pipeline_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_no_hsdp_not_standalone(self):
-        """Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=1,
-        )
-        assert result["original_dit_parallel_size"] == 1
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_combined_multiple_parallelism_not_standalone(self):
-        """Test HSDP combined with multiple parallelism is NOT standalone."""
-        result = self.compute_standalone_hsdp_params(
-            sequence_parallel_size=2,
-            cfg_parallel_size=2,
-            fully_shard_degree=4,
-            hsdp_replicate_size=1,
-        )
-        assert result["original_dit_parallel_size"] == 4  # 2 * 2
-        assert result["is_standalone_hsdp"] is False
-        assert result["effective_dp_size"] == 1
-
-    def test_standalone_hsdp_large_shard(self):
-        """Test standalone HSDP with large shard size."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=8,
-            hsdp_replicate_size=1,
-        )
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 8
-        assert result["effective_dp_size"] == 8
-
-    def test_standalone_hsdp_large_replicate(self):
-        """Test standalone HSDP with large replicate size."""
-        result = self.compute_standalone_hsdp_params(
-            fully_shard_degree=4,
-            hsdp_replicate_size=4,
-        )
-        assert result["is_standalone_hsdp"] is True
-        assert result["effective_dit_parallel_size"] == 16  # 4 * 4
-        assert result["effective_dp_size"] == 16
-
-
 class TestHSDPShardConditions:
     """Tests for _hsdp_shard_conditions matching logic."""
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -12,7 +12,6 @@ import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
-from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
@@ -310,7 +309,6 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch
 
     assert loader.load_model(load_device="cpu") is model
     assert not loaded_ordinary_weights
-    assert plan.requires_backend
     assert loader.take_host_weight_plan() is plan
     assert loader.take_host_weight_plan() is None
 
@@ -348,7 +346,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
         backing_kind="checkpoint_mmap",
         bindings={
             "transformer.weight": TensorBinding(
-                storage_key="weight",
+                checkpoint_key="weight",
                 file_path="unused",
             )
         },
@@ -376,9 +374,8 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
     assert loader.take_host_weight_plan() is plan
 
 
-def test_dlo_plan_fallback_retains_ordinary_loader_weights_by_default(monkeypatch):
+def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-    import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
 
     od_config = SimpleNamespace(
         dtype=torch.float32,
@@ -403,202 +400,11 @@ def test_dlo_plan_fallback_retains_ordinary_loader_weights_by_default(monkeypatc
         lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
     )
 
-    def build_cache(*_args, **_kwargs):
-        pytest.fail("host weight cache must require explicit selection")
-
-    monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
-
     assert loader.load_model(load_device="cpu") is model
     assert calls == ["load", "process"]
     assert loader.take_host_weight_plan() is None
 
 
-def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
-    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-    import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
-
-    od_config = SimpleNamespace(
-        dtype=torch.float32,
-        parallel_config=SimpleNamespace(
-            use_hsdp=False,
-            tensor_parallel_size=1,
-            sequence_parallel_size=1,
-            data_parallel_size=2,
-            cfg_parallel_size=1,
-            pipeline_parallel_size=1,
-            ulysses_degree=1,
-            ring_degree=1,
-            allgather_degree=1,
-            ulysses_mode="strict",
-        ),
-        quantization_config=None,
-        enable_distributed_layerwise_offload=True,
-        dlo_use_allgather=False,
-        dlo_use_host_weight_cache=True,
-        dlo_host_weight_cache_dir=str(tmp_path),
-        dlo_host_weight_cache_lock_timeout=1.0,
-        model="unused",
-        revision=None,
-        model_config={},
-    )
-    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
-    model = nn.Module()
-    model.transformer = nn.Linear(2, 2, bias=False)
-    model.transformer.quant_method = UnquantizedLinearMethod()
-    events: list[str] = []
-    runtime_plan = HostWeightPlan(
-        backing_kind="host_weight_cache",
-        bindings={},
-        runtime_layout_key="layout-key",
-    )
-
-    def load_weights(_model):
-        events.append("load")
-        with torch.no_grad():
-            model.transformer.weight.fill_(1)
-
-    def process_weights(_model, _device):
-        events.append("process")
-        with torch.no_grad():
-            model.transformer.weight.add_(1)
-
-    def calibrate(_model):
-        events.append("calibrate")
-        with torch.no_grad():
-            model.transformer.weight.add_(1)
-
-    original_eval = model.eval
-
-    def eval_model():
-        events.append("eval")
-        return original_eval()
-
-    def build_cache(_pipeline, **kwargs):
-        events.append("cache")
-        assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
-        assert kwargs["tensor_parallel_rank"] == 0
-        assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
-        assert kwargs["quantization_config"] is None
-        return HostWeightPlanResult(runtime_plan)
-
-    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
-    loader.load_weights = load_weights  # type: ignore[method-assign]
-    loader._process_weights_after_loading = process_weights  # type: ignore[method-assign]
-    loader._apply_skip_softmax_calibration = calibrate  # type: ignore[method-assign]
-    model.eval = eval_model  # type: ignore[method-assign]
-    monkeypatch.setattr(
-        loader_mod,
-        "build_checkpoint_mmap_plan",
-        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
-    )
-    monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
-
-    assert loader.load_model(load_device="cpu") is model
-    assert events == ["load", "process", "calibrate", "eval", "cache"]
-    assert loader.take_host_weight_plan() is runtime_plan
-
-
-@pytest.mark.parametrize("pin_cpu_memory", [True, False])
-def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeypatch, tmp_path, pin_cpu_memory):
-    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-    import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
-
-    od_config = SimpleNamespace(
-        dtype=torch.float32,
-        parallel_config=SimpleNamespace(
-            use_hsdp=False,
-            tensor_parallel_size=1,
-            sequence_parallel_size=1,
-            data_parallel_size=1,
-            cfg_parallel_size=1,
-            pipeline_parallel_size=1,
-        ),
-        quantization_config=None,
-        enable_distributed_layerwise_offload=True,
-        dlo_use_allgather=False,
-        dlo_use_host_weight_cache=True,
-        pin_cpu_memory=pin_cpu_memory,
-        dlo_host_weight_cache_dir=str(tmp_path),
-        dlo_host_weight_cache_pin_limit_gib=0.0,
-        model="unused",
-    )
-    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
-    model = nn.Module()
-    model.transformer = nn.Linear(2, 2, bias=False)
-    events: list[str] = []
-    runtime_plan = HostWeightPlan(
-        backing_kind="host_weight_cache",
-        bindings={},
-        runtime_layout_key="layout-key",
-    )
-
-    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
-    loader.load_weights = lambda _model: events.append("load")  # type: ignore[method-assign]
-    loader._process_weights_after_loading = lambda *_args: events.append("process")  # type: ignore[method-assign]
-    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
-
-    def unexpected_checkpoint_plan(*_args, **_kwargs):
-        pytest.fail("selected host weight cache mode must not select the raw checkpoint mmap plan")
-
-    def build_cache(_pipeline, **_kwargs):
-        events.append("cache")
-        return HostWeightPlanResult(runtime_plan)
-
-    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
-    monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
-
-    assert loader.load_model(load_device="cpu") is model
-    assert events == ["load", "process", "cache"]
-    assert loader.take_host_weight_plan() is runtime_plan
-
-
-def test_dlo_host_weight_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
-    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-    import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
-
-    od_config = SimpleNamespace(
-        dtype=torch.float32,
-        parallel_config=SimpleNamespace(
-            use_hsdp=False,
-            tensor_parallel_size=1,
-            sequence_parallel_size=1,
-            data_parallel_size=2,
-            cfg_parallel_size=1,
-            pipeline_parallel_size=1,
-        ),
-        quantization_config=None,
-        enable_distributed_layerwise_offload=True,
-        dlo_use_allgather=False,
-        dlo_use_host_weight_cache=True,
-        dlo_host_weight_cache_dir=str(tmp_path),
-        model="unused",
-    )
-    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
-    model = nn.Module()
-    model.transformer = nn.Linear(2, 2, bias=False)
-    model.transformer.quant_method = object()
-    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
-    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
-    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
-    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
-    monkeypatch.setattr(
-        loader_mod,
-        "build_checkpoint_mmap_plan",
-        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
-    )
-
-    def build_cache(_pipeline, **kwargs):
-        assert kwargs["quantization_config"] == "module_quant_method"
-        return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
-
-    monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
-
-    original_weight = model.transformer.weight.detach().clone()
-    assert loader.load_model(load_device="cpu") is model
-    assert torch.equal(model.transformer.weight, original_weight)
-    assert loader.take_host_weight_plan() is None
-
-
 def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
     from vllm.model_executor.layers.quantization.online.fp8 import (
         Fp8PerTensorOnlineLinearMethod,

tests/diffusion/model_loader/test_host_weight_cache.py

@@ -1,260 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import fcntl
-import multiprocessing as mp
-import shutil
-from pathlib import Path
-
-import pytest
-import torch
-from torch import nn
-
-from vllm_omni.diffusion.model_loader.host_weight_cache import (
-    build_host_weight_cache_plan,
-    default_host_weight_cache_root,
-)
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
-
-_SPAWN_TIMEOUT_SECONDS = 120
-
-
-class _LoaderMarker:
-    pass
-
-
-class _Transformer(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.proj = nn.Linear(3, 2)
-        self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
-        self.register_buffer("temporary", torch.ones(1), persistent=False)
-        with torch.no_grad():
-            self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
-            self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
-
-
-class _Pipeline(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.transformer = _Transformer()
-
-
-def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
-    monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
-
-    assert default_host_weight_cache_root().endswith("/.cache/vllm-omni/dlo-host-weights")
-
-
-def _build(pipeline: nn.Module, cache_root: Path, **overrides):
-    kwargs = {
-        "dit_modules": (("transformer", pipeline.transformer),),
-        "loader_type": _LoaderMarker,
-        "cache_root": cache_root,
-        "lock_timeout_seconds": 2.0,
-        "max_shard_bytes": 32,
-        "model_identity": "test/model",
-        "revision": "revision",
-        "runtime_dtype": torch.float32,
-        "load_format": "default",
-        "loader_inputs": {"model_config": {"hidden_size": 3}},
-        "tensor_parallel_size": 1,
-        "tensor_parallel_rank": 0,
-        "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
-        "use_hsdp": False,
-        "enable_expert_parallel": False,
-        "quantization_config": None,
-        "cfg_parallel_size": 1,
-        "pipeline_parallel_size": 1,
-    }
-    kwargs.update(overrides)
-    return build_host_weight_cache_plan(pipeline, **kwargs)
-
-
-def _multiprocess_builder(cache_root: str, queue) -> None:
-    result = _build(_Pipeline(), Path(cache_root))
-    queue.put(
-        (
-            result.fallback_code,
-            result.plan.runtime_layout_key if result.plan is not None else None,
-            sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
-        )
-    )
-
-
-def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
-    first = _build(_Pipeline(), tmp_path)
-    second = _build(_Pipeline(), tmp_path)
-
-    assert first.plan is not None
-    assert second.plan is not None
-    assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
-    assert not first.plan.requires_backend
-    assert {binding.file_path for binding in first.plan.bindings.values()} == {
-        binding.file_path for binding in second.plan.bindings.values()
-    }
-    assert set(first.plan.bindings) == {
-        "transformer.persistent",
-        "transformer.proj.bias",
-        "transformer.proj.weight",
-    }
-
-
-def test_equivalent_local_model_paths_reuse_one_entry(tmp_path, monkeypatch):
-    model_dir = tmp_path / "model"
-    model_dir.mkdir()
-    model_link = tmp_path / "model-link"
-    model_link.symlink_to(model_dir, target_is_directory=True)
-    cache_root = tmp_path / "cache"
-    monkeypatch.setenv("HOME", str(tmp_path))
-    monkeypatch.chdir(tmp_path)
-
-    identities = (str(model_dir), "~/model", "model", str(model_link))
-    results = [_build(_Pipeline(), cache_root, model_identity=identity) for identity in identities]
-
-    assert all(result.plan is not None for result in results)
-    assert len({result.plan.runtime_layout_key for result in results if result.plan is not None}) == 1
-    assert len(list((cache_root / "v1").iterdir())) == 1
-
-
-def test_unstable_loader_identity_fails_closed(tmp_path):
-    result = _build(_Pipeline(), tmp_path, loader_inputs={"opaque": object()})
-
-    assert result.plan is None
-    assert result.fallback_code == "unstable_identity"
-    assert result.fallback_reason == "host weight cache identity does not support values of type builtins.object"
-    assert list(tmp_path.rglob("*.safetensors")) == []
-
-
-def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
-    pipeline = _Pipeline()
-    tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
-    tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
-    sp2 = _build(
-        pipeline,
-        tmp_path,
-        sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
-    )
-
-    assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
-    assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
-
-
-def test_final_runtime_content_is_authoritative(tmp_path):
-    first_pipeline = _Pipeline()
-    second_pipeline = _Pipeline()
-    with torch.no_grad():
-        second_pipeline.transformer.proj.weight[0, 0] += 1
-
-    first = _build(first_pipeline, tmp_path)
-    second = _build(second_pipeline, tmp_path)
-
-    assert first.plan is not None and second.plan is not None
-    assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
-
-
-def test_expert_parallel_ownership_fails_closed(tmp_path):
-    result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
-
-    assert result.plan is None
-    assert result.fallback_code == "unsupported_expert_parallel"
-    assert list(tmp_path.rglob("*.safetensors")) == []
-
-
-def test_pipeline_parallel_ownership_fails_closed(tmp_path):
-    result = _build(_Pipeline(), tmp_path, pipeline_parallel_size=2)
-
-    assert result.plan is None
-    assert result.fallback_code == "unsupported_pp"
-    assert list(tmp_path.rglob("*.safetensors")) == []
-
-
-def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
-    first = _build(_Pipeline(), tmp_path)
-    assert first.plan is not None
-    shard = Path(next(iter(first.plan.bindings.values())).file_path)
-    old_inode = shard.stat().st_ino
-    with shard.open("r+b") as handle:
-        handle.seek(-1, 2)
-        byte = handle.read(1)
-        handle.seek(-1, 2)
-        handle.write(bytes([byte[0] ^ 0xFF]))
-
-    second = _build(_Pipeline(), tmp_path)
-
-    assert second.plan is not None
-    assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
-    assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
-
-
-def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
-    initial = _build(_Pipeline(), tmp_path)
-    assert initial.plan is not None
-    cache_key = initial.plan.runtime_layout_key
-    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
-    shutil.rmtree(entry_dir)
-    lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
-
-    with lock_path.open("a+b") as lock_handle:
-        fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
-        result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
-
-    assert result.plan is None
-    assert result.fallback_code == "lock_timeout"
-
-
-def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
-    initial = _build(_Pipeline(), tmp_path)
-    assert initial.plan is not None
-    cache_key = initial.plan.runtime_layout_key
-    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
-    shutil.rmtree(entry_dir)
-    stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
-    stale_temp.mkdir(parents=True)
-    (stale_temp / "partial").write_bytes(b"partial")
-
-    result = _build(_Pipeline(), tmp_path)
-
-    assert result.plan is not None
-    assert not stale_temp.exists()
-
-
-@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
-def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
-    pipeline = _Pipeline()
-    if case == "alias":
-        pipeline.transformer.tied = pipeline.transformer.proj.weight
-        expected_code = "unsupported_alias"
-    elif case == "external_alias":
-        pipeline.encoder = nn.Module()
-        pipeline.encoder.shared = pipeline.transformer.proj.weight
-        expected_code = "unsupported_alias"
-    else:
-        pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
-        expected_code = "unsupported_tensor"
-
-    result = _build(pipeline, tmp_path)
-
-    assert result.plan is None
-    assert result.fallback_code == expected_code
-    assert list(tmp_path.rglob("*.safetensors")) == []
-
-
-def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
-    context = mp.get_context("spawn")
-    queue = context.Queue()
-    processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
-    for process in processes:
-        process.start()
-    # Spawn workers import the full vLLM stack before entering the target;
-    # loaded CI hosts can spend more than 30 seconds in that startup path.
-    results = [queue.get(timeout=_SPAWN_TIMEOUT_SECONDS) for _ in processes]
-    for process in processes:
-        process.join(timeout=_SPAWN_TIMEOUT_SECONDS)
-
-    assert all(process.exitcode == 0 for process in processes)
-    assert all(code is None for code, _, _ in results)
-    assert len({key for _, key, _ in results}) == 1
-    assert len({tuple(paths) for _, _, paths in results}) == 1
-    assert len(list((tmp_path / "v1").iterdir())) == 1

tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's camera-branch reduce-scatter.
+
+``out_proj_cam`` is row-parallel, so every rank produces a partial sum over the
+full hidden width, while the GDN stream it is added to is TP-local. Instead of
+all-reducing and then discarding all but one slice, the branch reduce-scatters
+straight to the local shard. These tests pin the shard each rank ends up with,
+including the bias slice, on both code paths:
+
+* the process-group path, which drives ``torch.distributed.reduce_scatter_tensor``
+  directly because vLLM-Omni's diffusion ``GroupCoordinator`` does not implement
+  ``reduce_scatter``;
+* the all-reduce-and-slice fallback taken when the coordinator exposes no
+  ``device_group``.
+
+CPU-only: the collectives and the rank/world-size lookups are mocked, so a
+single process can simulate an arbitrary TP degree.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmSelfAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+HIDDEN = 24
+
+
+def _host(bias: torch.Tensor | None) -> SimpleNamespace:
+    """A stand-in ``self`` exposing only what the method reads."""
+    return SimpleNamespace(out_proj_cam=SimpleNamespace(bias=bias))
+
+
+def _fake_reduce_scatter_tensor(rank: int, tp_size: int):
+    """Model ``reduce_scatter_tensor``: split dim 0 and keep chunk ``rank``.
+
+    The test drives a single process, so the other ranks' partials are already
+    folded into the input; scattering is all this has to reproduce.
+    """
+
+    def _impl(output: torch.Tensor, input_: torch.Tensor, group=None) -> None:
+        output.copy_(input_.chunk(tp_size, dim=0)[rank])
+
+    return _impl
+
+
+def _tp_group(with_device_group: bool):
+    return SimpleNamespace(device_group=object() if with_device_group else None)
+
+
+@pytest.mark.parametrize("tp_size", [2, 4])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_process_group_path_matches_all_reduce_then_slice(tp_size: int, rank: int) -> None:
+    """The scattered shard equals the slice the all-reduce path would take."""
+    local_width = HIDDEN // tp_size
+    torch.manual_seed(0)
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+    reference = torch.zeros(2, 3, local_width)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+        patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+        patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(True)),
+        patch("torch.distributed.reduce_scatter_tensor", side_effect=_fake_reduce_scatter_tensor(rank, tp_size)),
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+    start = rank * local_width
+    expected = (contrib + bias)[..., start : start + local_width]
+    assert out.shape == (2, 3, local_width)
+    torch.testing.assert_close(out, expected)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_fallback_path_matches_the_process_group_path(rank: int) -> None:
+    """Without a device_group the method all-reduces and slices, same result."""
+    tp_size = 2
+    local_width = HIDDEN // tp_size
+    torch.manual_seed(1)
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+    reference = torch.zeros(2, 3, local_width)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+        patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+        patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(False)),
+        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t),
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+    start = rank * local_width
+    torch.testing.assert_close(out, (contrib + bias)[..., start : start + local_width])
+
+
+def test_tp1_adds_bias_without_a_collective() -> None:
+    contrib = torch.randn(2, 3, HIDDEN)
+    bias = torch.randn(HIDDEN)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1),
+        patch("torch.distributed.reduce_scatter_tensor") as collective,
+    ):
+        out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, contrib)
+
+    assert collective.call_count == 0
+    torch.testing.assert_close(out, contrib + bias)
+
+
+def test_rejects_a_width_that_does_not_shard() -> None:
+    """A local width that is not 1/tp of the contribution is a wiring bug."""
+    contrib = torch.randn(1, 2, HIDDEN)
+    reference = torch.zeros(1, 2, HIDDEN // 3)  # tp=2 would need HIDDEN // 2
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch("torch.distributed.reduce_scatter_tensor"),
+        pytest.raises(ValueError, match="camera contribution width mismatch"),
+    ):
+        SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(None), contrib, reference)

tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's TensorParallelRMSNorm q/k fusion.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the collective is mocked, so a single process simulates one rank
+(all-reduce is identity). This exercises the cat/split logic a real multi-rank
+run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+    TensorParallelRMSNorm,
+    fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+
+def _make_norms(dim: int, tp_size: int) -> tuple[TensorParallelRMSNorm, TensorParallelRMSNorm]:
+    """Two norms with distinct random weights so a swapped q/k slice is caught."""
+    torch.manual_seed(0)
+    norm_q = TensorParallelRMSNorm(dim, tp_size=tp_size)
+    norm_k = TensorParallelRMSNorm(dim, tp_size=tp_size)
+    norm_q.weight.data.normal_()
+    norm_k.weight.data.normal_()
+    return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+    """Single-rank all-reduce: the reduction over one rank is the identity."""
+    return tensor
+
+
+def test_fused_matches_separate() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim, tp_size=2)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_one_all_reduce() -> None:
+    dim = 32
+    norm_q, norm_k = _make_norms(dim, tp_size=4)
+    q = torch.randn(1, 3, dim)
+    k = torch.randn(1, 3, dim)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 1
+        # Both sums travel in one tensor: last dim is 2, not 1.
+        assert collective.call_args.args[0].shape[-1] == 2
+
+        collective.reset_mock()
+        norm_q(q), norm_k(k)
+        assert collective.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "norms",
+    [
+        (RMSNorm(16), RMSNorm(16)),
+        (nn.Identity(), nn.Identity()),
+    ],
+    ids=["plain_rms_norm", "identity"],
+)
+def test_falls_back_for_non_tp_norms(norms: tuple[nn.Module, nn.Module]) -> None:
+    """qk_norm=False and TP=1 both hand in norms the fused path cannot use."""
+    norm_q, norm_k = norms
+    q = torch.randn(2, 4, 16)
+    k = torch.randn(2, 4, 16)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 0
+
+    torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
+
+
+def test_tp1_skips_the_collective() -> None:
+    norm_q, norm_k = _make_norms(16, tp_size=1)
+    q = torch.randn(2, 4, 16)
+    k = torch.randn(2, 4, 16)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+        assert collective.call_count == 0
+
+    torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)

tests/diffusion/models/sana_wm/test_request.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SANA-WM request normalization."""
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_wm.request import (
+    SANA_WM_DEFAULT_HEIGHT,
+    SANA_WM_DEFAULT_NUM_FRAMES,
+    SANA_WM_DEFAULT_WIDTH,
+    normalize_sana_wm_payload,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _prompt(**sana_wm):
+    """A valid request, with the action rollout kept in step with num_frames."""
+    block = {"num_frames": SANA_WM_DEFAULT_NUM_FRAMES, **sana_wm}
+    block.setdefault("action", f"w-{block['num_frames'] - 1}")
+    return {
+        "prompt": "a quiet city street",
+        "multi_modal_data": {"image": object()},
+        "sana_wm": block,
+    }
+
+
+def _identity_poses(count):
+    return [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]] for _ in range(count)]
+
+
+def _camera_prompt(**camera):
+    prompt = _prompt(num_frames=9)
+    prompt["sana_wm"].pop("action")
+    prompt["sana_wm"]["camera"] = {"poses": _identity_poses(9), **camera}
+    return prompt
+
+
+def test_defaults_are_latent_aligned():
+    payload = normalize_sana_wm_payload(_prompt())["additional_information"]["sana_wm"]
+    assert payload["num_frames"] == SANA_WM_DEFAULT_NUM_FRAMES
+    assert payload["height"] == SANA_WM_DEFAULT_HEIGHT
+    assert payload["width"] == SANA_WM_DEFAULT_WIDTH
+
+
+@pytest.mark.parametrize("num_frames", [9, 33, 161])
+def test_accepts_aligned_geometry(num_frames):
+    payload = normalize_sana_wm_payload(_prompt(num_frames=num_frames))["additional_information"]["sana_wm"]
+    assert payload["num_frames"] == num_frames
+
+
+@pytest.mark.parametrize("num_frames", [8, 10, 160])
+def test_rejects_unaligned_num_frames(num_frames):
+    # The VAE compresses 8x temporally; an unaligned count would be floored.
+    with pytest.raises(ValueError, match="num_frames"):
+        normalize_sana_wm_payload(_prompt(num_frames=num_frames))
+
+
+@pytest.mark.parametrize(("height", "width"), [(700, 1280), (704, 1270), (703, 1279)])
+def test_rejects_unaligned_resolution(height, width):
+    with pytest.raises(ValueError, match="divisible by 32"):
+        normalize_sana_wm_payload(_prompt(height=height, width=width))
+
+
+def test_normalization_is_idempotent():
+    once = normalize_sana_wm_payload(_prompt(num_frames=33))
+    twice = normalize_sana_wm_payload(once)
+    assert once["additional_information"]["sana_wm"] == twice["additional_information"]["sana_wm"]
+
+
+def test_rejects_action_shorter_than_num_frames():
+    # Used to shrink the camera tensor silently, then fail deep in the
+    # transformer on a camera/latent frame mismatch.
+    with pytest.raises(ValueError, match="rolls out 33 frames but num_frames is 161"):
+        normalize_sana_wm_payload(_prompt(action="w-32"))
+
+
+def test_rejects_action_longer_than_num_frames():
+    with pytest.raises(ValueError, match="rolls out 161 frames but num_frames is 9"):
+        normalize_sana_wm_payload(_prompt(num_frames=9, action="w-160"))
+
+
+def test_rejects_oversized_action_without_expanding_it():
+    # A 12-byte request body: expanding this rollout before checking the frame
+    # count allocated ~16 GB. Rejection has to be arithmetic, so this test also
+    # fails (by hanging/OOM) if the duration is ever expanded again.
+    with pytest.raises(ValueError, match="rolls out 1000000001 frames"):
+        normalize_sana_wm_payload(_prompt(action="w-1000000000"))
+
+
+@pytest.mark.parametrize("action", ["", "   ", "w-0", "w-abc", "w", "x-4"])
+def test_rejects_malformed_action(action):
+    with pytest.raises(ValueError):
+        normalize_sana_wm_payload(_prompt(action=action))
+
+
+@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")])
+def test_rejects_unusable_focal_length(value):
+    # compute_raymap divides by fx/fy; a zero or non-finite focal length turns
+    # the whole Plucker map into NaN and the response into a NaN video.
+    intrinsics = {"fx": value, "fy": 640, "cx": 640, "cy": 352}
+    with pytest.raises(ValueError, match="fx"):
+        normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_principal_point(value):
+    intrinsics = {"fx": 640, "fy": 640, "cx": value, "cy": 352}
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_speeds(value):
+    # `parsed <= 0` is False for NaN, so the positivity check alone let it pass.
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(_prompt(translation_speed=value))
+
+
+def test_rejects_non_finite_pose_entry():
+    prompt = _camera_prompt()
+    prompt["sana_wm"]["camera"]["poses"][3][2][1] = float("nan")
+    with pytest.raises(ValueError, match="must be finite"):
+        normalize_sana_wm_payload(prompt)
+
+
+def test_accepts_default_camera_metadata():
+    payload = normalize_sana_wm_payload(_camera_prompt())["additional_information"]["sana_wm"]
+    assert payload["camera"]["poses"]
+    # The unread fields are not carried into the canonical payload.
+    assert set(payload["camera"]) == {"poses"}
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [("format", "w2c_4x4"), ("coordinate_system", "opengl")],
+)
+def test_rejects_unsupported_camera_convention(field, value):
+    # Nothing downstream reads these, so accepting one would mean interpreting
+    # the poses under the wrong convention.
+    with pytest.raises(ValueError, match=field.split("_")[0]):
+        normalize_sana_wm_payload(_camera_prompt(**{field: value}))

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -1,346 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import ctypes
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
-from vllm_omni.diffusion.offloader.cuda_host_registration import (
-    CudaHostRegistration,
-    CudaHostRegistrationBudgetError,
-    CudaHostRegistrationCleanupError,
-    CudaHostRegistrationError,
-    _coalesce_ranges,
-)
-from vllm_omni.diffusion.offloader.host_registration import (
-    HostRegistrationError,
-    register_host_mappings,
-)
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
-
-
-class _FakeStorage:
-    def __init__(self, pointer: int, size: int) -> None:
-        self._pointer = pointer
-        self._size = size
-
-    def data_ptr(self) -> int:
-        return self._pointer
-
-    def nbytes(self) -> int:
-        return self._size
-
-
-class _FakeTensor:
-    device = SimpleNamespace(type="cpu")
-
-    def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
-        self._storage = _FakeStorage(pointer, size)
-        self._pinned = pinned
-
-    def untyped_storage(self) -> _FakeStorage:
-        return self._storage
-
-    def numel(self) -> int:
-        return self._storage.nbytes()
-
-    def is_pinned(self) -> bool:
-        return self._pinned
-
-
-class _FakeRuntime:
-    def __init__(
-        self,
-        register_results: list[int | Exception],
-        unregister_results: list[int | Exception] | None = None,
-        *,
-        read_only_registration_supported: bool = True,
-        attribute_result: int = 0,
-    ) -> None:
-        self._register_results = iter(register_results)
-        self._unregister_results = iter(unregister_results or [])
-        self._read_only_registration_supported = read_only_registration_supported
-        self._attribute_result = attribute_result
-        self.attribute_queries: list[tuple[int, int]] = []
-        self.registered: list[tuple[int, int, int]] = []
-        self.unregistered: list[int] = []
-
-    def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int:
-        self.attribute_queries.append((attribute, device))
-        ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_registration_supported)
-        return self._attribute_result
-
-    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
-        self.registered.append((pointer, size, flags))
-        result = next(self._register_results)
-        if isinstance(result, Exception):
-            raise result
-        return result
-
-    def cudaHostUnregister(self, pointer: int) -> int:
-        self.unregistered.append(pointer)
-        result = next(self._unregister_results, 0)
-        if isinstance(result, Exception):
-            raise result
-        return result
-
-    @staticmethod
-    def cudaGetErrorString(error: int) -> str:
-        return f"error-{error}"
-
-
-@pytest.fixture(autouse=True)
-def _current_cuda_device(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0)
-
-
-def test_platform_factory_reports_unsupported_registration() -> None:
-    with pytest.raises(HostRegistrationError, match="not supported on cpu"):
-        register_host_mappings({}, device=torch.device("cpu"), max_bytes=4096)
-
-
-def test_platform_factory_dispatches_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
-    sentinel = object()
-    sources = {"weights": []}
-
-    def create(actual_sources, *, max_bytes):
-        assert actual_sources is sources
-        assert max_bytes == 4096
-        return sentinel
-
-    monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
-
-    assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=4096) is sentinel
-
-
-def test_platform_factory_allows_registration_without_additional_budget(monkeypatch: pytest.MonkeyPatch) -> None:
-    sentinel = object()
-    sources = {"weights": []}
-
-    def create(actual_sources, *, max_bytes):
-        assert actual_sources is sources
-        assert max_bytes is None
-        return sentinel
-
-    monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
-
-    assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=None) is sentinel
-
-
-def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
-    assert _coalesce_ranges(
-        [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
-        page_size=4096,
-    ) == (
-        registration_module._AddressRange(0x1000, 0x3000),
-        registration_module._AddressRange(0x9000, 0xA000),
-    )
-
-
-def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-
-    with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
-        CudaHostRegistration.create(
-            {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
-            max_bytes=4096,
-        )
-
-    assert runtime.registered == []
-
-
-def test_registration_without_additional_budget_registers_complete_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-
-    registration = CudaHostRegistration.create(
-        {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
-        max_bytes=None,
-    )
-
-    assert registration.total_bytes == 8192
-    assert len(runtime.registered) == 1
-    assert registration.close() == []
-
-
-def test_registration_rejects_unsupported_read_only_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0], read_only_registration_supported=False)
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-
-    with pytest.raises(CudaHostRegistrationError, match="does not support read-only host registration"):
-        CudaHostRegistration.create(
-            {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
-            max_bytes=4096,
-        )
-
-    assert runtime.attribute_queries == [
-        (registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0)
-    ]
-    assert runtime.registered == []
-
-
-def test_registration_consumes_capability_query_error(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0], attribute_result=7)
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-    consumed_errors: list[int] = []
-    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
-
-    with pytest.raises(CudaHostRegistrationError, match="error-7"):
-        CudaHostRegistration.create(
-            {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
-            max_bytes=4096,
-        )
-
-    assert consumed_errors == [7]
-    assert runtime.registered == []
-
-
-def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0, 7])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-    consumed_errors: list[int] = []
-    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
-
-    with pytest.raises(CudaHostRegistrationError, match="error-7"):
-        CudaHostRegistration.create(
-            {
-                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
-                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
-            },
-            max_bytes=8192,
-        )
-
-    assert runtime.unregistered == [0x1000]
-    assert consumed_errors == [7]
-
-
-def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-
-    with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
-        CudaHostRegistration.create(
-            {
-                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
-                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
-            },
-            max_bytes=8192,
-        )
-
-    assert runtime.unregistered == [0x1000]
-
-
-def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0, 7], unregister_results=[9])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-    consumed_errors: list[int] = []
-    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
-
-    with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
-        CudaHostRegistration.create(
-            {
-                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
-                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
-            },
-            max_bytes=8192,
-        )
-
-    assert consumed_errors == [7, 9]
-
-
-@pytest.mark.parametrize("pending_error", [0, 801])
-def test_consume_last_cuda_error_accepts_cleared_or_matching_state(
-    monkeypatch: pytest.MonkeyPatch,
-    pending_error: int,
-) -> None:
-    class _GetLastError:
-        def __call__(self) -> int:
-            return pending_error
-
-    runtime = SimpleNamespace(cudaGetLastError=_GetLastError())
-    monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime)
-
-    registration_module._consume_last_cuda_error(801)
-
-
-def test_consume_last_cuda_error_rejects_unrelated_pending_error(monkeypatch: pytest.MonkeyPatch) -> None:
-    class _GetLastError:
-        def __call__(self) -> int:
-            return 700
-
-    runtime = SimpleNamespace(cudaGetLastError=_GetLastError())
-    monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime)
-
-    with pytest.raises(CudaHostRegistrationCleanupError, match="cudaGetLastError reported 700"):
-        registration_module._consume_last_cuda_error(801)
-
-
-def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0, 0])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-    registration = CudaHostRegistration.create(
-        {
-            "first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)],  # type: ignore[list-item]
-            "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
-        },
-        max_bytes=12288,
-    )
-
-    assert registration.total_bytes == 12288
-    assert registration.region_count == 2
-    assert registration.close() == []
-    assert runtime.unregistered == [0x9000, 0x1000]
-    assert registration.close() == []
-
-
-def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0, 0])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-
-    registration = CudaHostRegistration.create(
-        {
-            "first.safetensors": [_FakeTensor(0x1000, 4096)],  # type: ignore[list-item]
-            "second.safetensors": [_FakeTensor(0x2000, 4096)],  # type: ignore[list-item]
-        },
-        max_bytes=8192,
-    )
-
-    assert registration.region_count == 2
-    assert runtime.registered == [
-        (0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
-        (0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
-    ]
-    assert registration.close() == []
-
-
-def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
-    runtime = _FakeRuntime([0], unregister_results=[9, 0])
-    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
-    consumed_errors: list[int] = []
-    monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
-    registration = CudaHostRegistration.create(
-        {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
-        max_bytes=4096,
-    )
-
-    assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
-    assert registration.close() == []
-    assert runtime.unregistered == [0x1000, 0x1000]
-    assert consumed_errors == [9]

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -18,7 +18,6 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
-from vllm_omni.diffusion.model_loader.host_weight_cache_utils import file_digest
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     TensorBinding,
@@ -547,31 +546,6 @@ class TestDistributedLayerwiseOffloadHook:
         assert staging[0][torch.float32].numel() == 12
         assert staging[1][torch.float32].numel() == 12
 
-    def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
-        current_block = nn.Linear(2, 2, bias=False)
-        next_block = nn.Linear(2, 2, bias=False)
-        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
-        next_block.weight.data.copy_(expected)
-        hook = DistributedLayerwiseOffloadHook(
-            next_block=next_block,
-            device=torch.device("cpu"),
-            dp_group=None,
-            dp_size=1,
-            rank=0,
-            pin_memory=False,
-            rank_local_mmap=True,
-        )
-        hook.initialize_hook(current_block)
-        hook.registered_mmap = True
-
-        def fail_staging(_slot):
-            raise AssertionError("registered mmap must bypass host staging")
-
-        hook._stage_mmap_sources = fail_staging  # type: ignore[method-assign]
-        hook.prefetch_layer(slot=0, non_blocking=True)
-
-        assert torch.equal(next_block.weight, expected)
-
 
 class TestPinnedResidentLayerGroup:
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -662,24 +636,6 @@ class TestPinnedResidentLayerGroup:
         group.offload()
         assert all(block.weight.numel() == 0 for block in blocks)
 
-    def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
-        block = nn.Linear(2, 2, bias=False)
-        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
-        block.weight.data.copy_(expected)
-        group = PinnedResidentLayerGroup(
-            [block],
-            device=torch.device("cpu"),
-            copy_stream=DummyStream(),
-            pin_memory=False,
-            rank_local_mmap=True,
-        )
-        group.registered_mmap = True
-        group._cpu_staging_buffers.clear()
-
-        group.load()
-
-        assert torch.equal(block.weight, expected)
-
     def test_shared_allgather_output_is_narrowed_to_current_block(
         self, dist_group, patched_offload_runtime, monkeypatch
     ):
@@ -809,140 +765,6 @@ class _MmapPostLoadPipeline(nn.Module):
 
 
 class TestMmapWeightLoading:
-    def test_host_weight_cache_maps_final_weights_without_rerunning_post_load(
-        self,
-        tmp_path,
-        patched_offload_runtime,
-    ):
-        class Transformer(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.proj = nn.Linear(2, 2, bias=False)
-                self.post_load_calls = 0
-
-            def post_load_weights(self):
-                self.post_load_calls += 1
-
-        class Pipeline(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.transformer = Transformer()
-
-        pipeline = Pipeline()
-        old_parameter = pipeline.transformer.proj.weight
-        old_parameter.runtime_marker = "preserved"
-        weight_file = tmp_path / "runtime.safetensors"
-        save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file))
-        plan = HostWeightPlan(
-            backing_kind="host_weight_cache",
-            bindings={
-                "transformer.proj.weight": TensorBinding(
-                    storage_key="transformer.proj.weight",
-                    file_path=str(weight_file),
-                )
-            },
-            runtime_layout_key="runtime-layout",
-            expected_file_digests={str(weight_file): file_digest(weight_file)},
-        )
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dlo_use_allgather=False,
-            ),
-            torch.device("cpu"),
-            host_weight_plan=plan,
-        )
-        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
-
-        backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
-
-        assert pipeline.transformer.post_load_calls == 0
-        assert pipeline.transformer.proj.weight is old_parameter
-        assert pipeline.transformer.proj.weight.runtime_marker == "preserved"
-        assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0))
-        backend._release_mmap_handles()
-
-    def test_host_weight_cache_mapping_failure_leaves_ordinary_tensors_intact(
-        self,
-        tmp_path,
-        patched_offload_runtime,
-    ):
-        pipeline = nn.Module()
-        pipeline.transformer = nn.Linear(2, 2, bias=False)
-        old_parameter = pipeline.transformer.weight
-        plan = HostWeightPlan(
-            backing_kind="host_weight_cache",
-            bindings={
-                "transformer.weight": TensorBinding(
-                    storage_key="transformer.weight",
-                    file_path=str(tmp_path / "missing.safetensors"),
-                )
-            },
-            runtime_layout_key="runtime-layout",
-            expected_file_digests={str(tmp_path / "missing.safetensors"): "unused"},
-        )
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dlo_use_allgather=False,
-            ),
-            torch.device("cpu"),
-            host_weight_plan=plan,
-        )
-        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
-
-        with pytest.raises(Exception, match="missing.safetensors"):
-            backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
-
-        assert pipeline.transformer.weight is old_parameter
-
-    def test_host_weight_cache_digest_change_leaves_ordinary_tensors_intact(
-        self,
-        tmp_path,
-        patched_offload_runtime,
-    ):
-        pipeline = nn.Module()
-        pipeline.transformer = nn.Linear(2, 2, bias=False)
-        old_parameter = pipeline.transformer.weight
-        old_value = old_parameter.detach().clone()
-        weight_file = tmp_path / "runtime.safetensors"
-        save_file({"transformer.weight": torch.full((2, 2), 9.0)}, str(weight_file))
-        expected_digest = file_digest(weight_file)
-        with weight_file.open("r+b") as handle:
-            handle.seek(-1, os.SEEK_END)
-            original = handle.read(1)
-            handle.seek(-1, os.SEEK_END)
-            handle.write(bytes([original[0] ^ 0xFF]))
-        plan = HostWeightPlan(
-            backing_kind="host_weight_cache",
-            bindings={
-                "transformer.weight": TensorBinding(
-                    storage_key="transformer.weight",
-                    file_path=str(weight_file),
-                )
-            },
-            runtime_layout_key="runtime-layout",
-            expected_file_digests={str(weight_file): expected_digest},
-        )
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dlo_use_allgather=False,
-            ),
-            torch.device("cpu"),
-            host_weight_plan=plan,
-        )
-        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
-
-        with pytest.raises(RuntimeError, match="changed after loader validation"):
-            backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
-
-        assert pipeline.transformer.weight is old_parameter
-        assert torch.equal(pipeline.transformer.weight, old_value)
-
     def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
         pipeline = _MmapPostLoadPipeline()
         weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
@@ -966,7 +788,7 @@ class TestMmapWeightLoading:
             backing_kind="checkpoint_mmap",
             bindings={
                 name: TensorBinding(
-                    storage_key=name,
+                    checkpoint_key=name,
                     file_path=str(weight_file),
                 )
                 for name in weights
@@ -1037,292 +859,6 @@ class TestMmapWeightLoading:
         assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
         backend.disable()
 
-    def test_host_weight_cache_enable_reclaims_private_weights_before_hook_setup(
-        self,
-        tmp_path,
-        patched_offload_runtime,
-    ):
-        class Transformer(nn.Module):
-            _layerwise_offload_blocks_attrs = ["blocks"]
-
-            def __init__(self):
-                super().__init__()
-                self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
-
-        class Pipeline(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.transformer = Transformer()
-
-        pipeline = Pipeline()
-        weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
-        weight_file = tmp_path / "runtime.safetensors"
-        save_file(weights, str(weight_file))
-        plan = HostWeightPlan(
-            backing_kind="host_weight_cache",
-            bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
-            runtime_layout_key="runtime-layout",
-            expected_file_digests={str(weight_file): file_digest(weight_file)},
-        )
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dlo_use_allgather=False,
-            ),
-            torch.device("cpu"),
-            host_weight_plan=plan,
-        )
-        cleanup_calls: list[None] = []
-        backend._cleanup_after_loading = lambda: cleanup_calls.append(None)  # type: ignore[method-assign]
-
-        backend.enable(pipeline)
-
-        assert backend._using_rank_local_mmap
-        # One reclamation immediately after remap, plus the existing final
-        # post-hook cleanup.
-        assert len(cleanup_calls) == 2
-        backend.disable()
-
-    def test_host_weight_cache_enable_falls_back_to_ordinary_tensors(
-        self,
-        tmp_path,
-        patched_offload_runtime,
-    ):
-        class Transformer(nn.Module):
-            _layerwise_offload_blocks_attrs = ["blocks"]
-
-            def __init__(self):
-                super().__init__()
-                self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
-
-        class Pipeline(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.transformer = Transformer()
-
-        pipeline = Pipeline()
-        weight_file = tmp_path / "runtime.safetensors"
-        weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
-        save_file(weights, str(weight_file))
-        plan = HostWeightPlan(
-            backing_kind="host_weight_cache",
-            bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
-            runtime_layout_key="runtime-layout",
-            expected_file_digests={str(weight_file): "stale-digest"},
-        )
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dlo_use_allgather=False,
-            ),
-            torch.device("cpu"),
-            host_weight_plan=plan,
-        )
-
-        backend.enable(pipeline)
-
-        assert backend.enabled
-        assert backend.host_weight_plan is None
-        assert not backend._using_mmap
-        assert not backend._using_rank_local_mmap
-        backend.disable()
-
-    def test_disable_synchronizes_allgather_device_work(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-        patched_offload_runtime,
-    ):
-        synchronize_calls: list[None] = []
-        monkeypatch.setattr(
-            dist_backend_module.current_omni_platform,
-            "synchronize",
-            lambda: synchronize_calls.append(None),
-        )
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dlo_use_allgather=True,
-            ),
-            torch.device("cpu"),
-        )
-        backend.enabled = True
-        backend._using_rank_local_mmap = False
-
-        backend.disable()
-
-        assert synchronize_calls == [None]
-
-    def test_host_weight_cache_registration_uses_budget_and_releases_before_mmap(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-        patched_offload_runtime,
-    ):
-        events: list[str] = []
-
-        class Registration:
-            total_bytes = 4096
-            region_count = 1
-
-            @staticmethod
-            def close() -> list[str]:
-                events.append("unregister")
-                return []
-
-        create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
-
-        def register(sources, *, device, max_bytes):
-            assert device == torch.device("cuda")
-            create_calls.append((sources, max_bytes))
-            return Registration()
-
-        monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=True,
-                dlo_use_allgather=False,
-                dlo_host_weight_cache_pin_limit_gib=1.5,
-            ),
-            torch.device("cuda"),
-        )
-        sources = {"runtime.safetensors": [torch.ones(1)]}
-        backend._host_weight_cache_mapped_sources = sources
-
-        assert backend._try_register_host_weight_cache_mmap()
-        assert create_calls == [(sources, int(1.5 * 1024**3))]
-
-        backend.enabled = True
-        backend._using_rank_local_mmap = True
-        backend._mmap_file_cache = {}
-        backend._release_mmap_handles = lambda: events.append("mmap")  # type: ignore[method-assign]
-        backend.disable()
-
-        assert events == ["unregister", "mmap"]
-
-    def test_host_weight_cache_registration_defaults_to_pinned_memory_policy(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-        patched_offload_runtime,
-    ):
-        class Registration:
-            total_bytes = 4096
-            region_count = 1
-
-            @staticmethod
-            def close() -> list[str]:
-                return []
-
-        registration = Registration()
-        create_calls: list[tuple[dict[str, list[torch.Tensor]], int | None]] = []
-
-        def register(sources, *, device, max_bytes):
-            assert device == torch.device("cuda")
-            create_calls.append((sources, max_bytes))
-            return registration
-
-        monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=True,
-                dlo_use_allgather=False,
-                dlo_host_weight_cache_pin_limit_gib=0.0,
-            ),
-            torch.device("cuda"),
-        )
-        sources = {"runtime.safetensors": [torch.ones(1)]}
-        backend._host_weight_cache_mapped_sources = sources
-
-        assert backend._try_register_host_weight_cache_mmap()
-        assert create_calls == [(sources, None)]
-        assert backend._host_registration is registration
-
-    def test_host_weight_cache_registration_respects_disabled_pinned_memory(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-        patched_offload_runtime,
-    ):
-        monkeypatch.setattr(
-            dist_backend_module,
-            "register_host_mappings",
-            lambda *_args, **_kwargs: pytest.fail("registration must follow pin_cpu_memory=False"),
-        )
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dlo_use_allgather=False,
-                dlo_host_weight_cache_pin_limit_gib=0.0,
-            ),
-            torch.device("cuda"),
-        )
-        backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
-
-        assert not backend._try_register_host_weight_cache_mmap()
-        assert backend._host_registration is None
-
-    def test_host_weight_cache_transfer_is_configured_once_for_all_consumers(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-        patched_offload_runtime,
-    ):
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=True,
-                dlo_use_allgather=False,
-            ),
-            torch.device("cuda"),
-            host_weight_plan=HostWeightPlan(
-                backing_kind="host_weight_cache",
-                bindings={},
-            ),
-        )
-        backend._using_rank_local_mmap = True
-        backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
-        streamed_hook = SimpleNamespace(registered_mmap=False)
-        resident_group = SimpleNamespace(
-            registered_mmap=False,
-            _cpu_staging_buffers=[{"buffer": torch.ones(1)}],
-        )
-        backend._resident_layer_group = resident_group
-        register_calls: list[None] = []
-
-        def register() -> bool:
-            register_calls.append(None)
-            return True
-
-        monkeypatch.setattr(backend, "_try_register_host_weight_cache_mmap", register)
-
-        backend._configure_host_weight_cache_transfer([streamed_hook])  # type: ignore[list-item]
-
-        assert register_calls == [None]
-        assert streamed_hook.registered_mmap
-        assert resident_group.registered_mmap
-        assert resident_group._cpu_staging_buffers == []
-        assert backend._host_weight_cache_mapped_sources == {}
-
-    def test_host_weight_cache_registration_falls_back_on_unsupported_platform(
-        self,
-        patched_offload_runtime,
-    ):
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=True,
-                dlo_use_allgather=False,
-                dlo_host_weight_cache_pin_limit_gib=1.0,
-            ),
-            torch.device("cpu"),
-        )
-        backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
-
-        assert not backend._try_register_host_weight_cache_mmap()
-        assert backend._host_registration is None
-
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
@@ -1628,7 +1164,7 @@ class TestMmapValidation:
         assert result.fallback_reason is None
         assert result.plan is not None
         assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
-            storage_key="blocks.0.weight",
+            checkpoint_key="blocks.0.weight",
             file_path=str(checkpoint_file),
         )
         assert result.plan.planned_source_prefixes == frozenset({"transformer."})
@@ -1918,7 +1454,7 @@ class TestMmapValidation:
 
         assert result.plan is not None
         assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
-            storage_key="proj_in.weight",
+            checkpoint_key="proj_in.weight",
             file_path=str(checkpoint_file),
         )
 

tests/diffusion/test_diffusion_config_propagation.py

@@ -92,6 +92,30 @@ class TestParallelConfigPropagation:
         od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
         assert od.num_gpus == 2
 
+    def test_dp_is_inferred_from_num_gpus(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8)
+        assert od.parallel_config.data_parallel_size == 2
+        assert od.parallel_config.world_size == 8
+
+    def test_explicit_dp_is_validated_against_num_gpus(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+        assert od.parallel_config.data_parallel_size == 2
+
+        with pytest.raises(ValueError, match="does not match WORLD-derived value"):
+            OmniDiffusionConfig.from_kwargs(
+                model="x",
+                parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2),
+                num_gpus=8,
+            )
+
+    def test_hsdp_does_not_infer_ordinary_dp(self):
+        pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4)
+        od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+        assert od.parallel_config.data_parallel_size == 1
+        assert od.parallel_config.world_size == 4
+
 
 class TestCreateDefaultDiffusion:
     """Verify engine_args structure from create_default_diffusion."""

tests/diffusion/test_worker_wrapper_base.py

@@ -12,13 +12,11 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
-from types import SimpleNamespace
 from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
-import vllm_omni.diffusion.worker.diffusion_worker as worker_module
 from vllm_omni.diffusion.worker.diffusion_worker import (
     CustomPipelineWorkerExtension,
     DiffusionWorker,
@@ -206,29 +204,6 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
-    def test_worker_shutdown_disables_offloader_before_distributed_teardown(
-        self,
-        mocker: MockerFixture,
-        monkeypatch: pytest.MonkeyPatch,
-    ):
-        worker = object.__new__(DiffusionWorker)
-        offload_backend = mocker.Mock()
-        kv_transfer_manager = mocker.Mock()
-        worker.model_runner = SimpleNamespace(
-            offload_backend=offload_backend,
-            kv_transfer_manager=kv_transfer_manager,
-        )
-        events: list[str] = []
-        offload_backend.disable.side_effect = lambda: events.append("offloader")
-        kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
-        monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
-
-        worker.shutdown()
-
-        offload_backend.disable.assert_called_once_with()
-        kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
-        assert events == ["offloader", "kv_transfer", "distributed"]
-
 
 # -------------------------------------------------------------------------
 # Tests: execute_method

tests/e2e/offline_inference/test_sana_wm.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+import pytest
+
+pytestmark = [
+    pytest.mark.advanced_model,
+    pytest.mark.diffusion,
+    pytest.mark.gpu,
+]
+
+
+def _make_sana_wm_e2e_image() -> Any:
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.sana_wm import SANA_WM_OUTPUT_HEIGHT, SANA_WM_OUTPUT_WIDTH
+
+    return Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+
+
+def _sana_wm_e2e_tensor_parallel_size() -> int:
+    return int(
+        os.environ.get(
+            "SANA_WM_E2E_TENSOR_PARALLEL_SIZE",
+            os.environ.get("SANA_WM_E2E_TP", "1"),
+        )
+    )
+
+
+def _coerce_video_array(video: Any) -> np.ndarray:
+    import torch
+
+    if isinstance(video, list):
+        assert video, "SANA-WM e2e produced an empty video list."
+        video = video[0]
+    if isinstance(video, torch.Tensor):
+        video = video.detach().cpu().float().numpy()
+    video_array = np.asarray(video)
+    if video_array.ndim == 5:
+        assert video_array.shape[0] == 1
+        video_array = video_array[0]
+    assert video_array.ndim == 4
+    return video_array
+
+
+def _assert_sana_wm_e2e_shape(
+    video: np.ndarray,
+    *,
+    output_type: str,
+    num_frames: int,
+) -> None:
+    if output_type == "latent":
+        # Latent output is channel-first: (C, T, H, W) after removing the
+        # batch dimension. Stage-1 emits LTX-2 latents with 128 channels.
+        assert video.shape[0] == 128
+        assert 0 < video.shape[1] <= num_frames
+        return
+    assert 0 < video.shape[0] <= num_frames
+
+
+def _run_sana_wm_e2e(*, output_type: str) -> np.ndarray:
+    import torch
+
+    from vllm_omni.diffusion.models.sana_wm import (
+        SANA_WM_MODEL_ID,
+        SANA_WM_OUTPUT_HEIGHT,
+        SANA_WM_OUTPUT_WIDTH,
+    )
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
+
+    if not torch.cuda.is_available():
+        pytest.skip("Sana-WM e2e requires CUDA.")
+
+    model = os.environ.get("SANA_WM_E2E_MODEL", SANA_WM_MODEL_ID)
+    model_class_name = os.environ.get("SANA_WM_E2E_MODEL_CLASS", "SanaWmPipeline")
+    num_frames = int(os.environ.get("SANA_WM_E2E_NUM_FRAMES", "9"))
+    stage1_steps = int(
+        os.environ.get(
+            "SANA_WM_E2E_STAGE1_STEPS",
+            os.environ.get("SANA_WM_E2E_NUM_INFERENCE_STEPS", "1"),
+        )
+    )
+    action = os.environ.get("SANA_WM_E2E_ACTION", f"w-{max(num_frames - 1, 1)}")
+    tensor_parallel_size = _sana_wm_e2e_tensor_parallel_size()
+    image = _make_sana_wm_e2e_image()
+    omni = Omni(
+        model=model,
+        model_class_name=model_class_name,
+        enforce_eager=True,
+        tensor_parallel_size=tensor_parallel_size,
+    )
+    extra_args: dict[str, Any] = {}
+    # The native Stage-1 path is latent-token capped; raise it for full-size
+    # e2e runs (704x1280 exceeds the default 4096-token cap).
+    native_max_tokens = os.environ.get("SANA_WM_E2E_NATIVE_MAX_TOKENS", "").strip()
+    if native_max_tokens:
+        extra_args["sana_wm_native_max_tokens"] = int(native_max_tokens)
+    # Return the requested output type so the shape assertion matches the
+    # pipeline output instead of the raw Stage-1 latent default.
+
+    output = omni.generate(
+        {
+            "prompt": "A slow forward camera move through a quiet city street.",
+            "multi_modal_data": {"image": image},
+            "sana_wm": {
+                "action": action,
+                "num_frames": num_frames,
+                "translation_speed": 0.055,
+                "rotation_speed_deg": 1.2,
+                # Avoid the optional Pi3X dependency by passing deterministic
+                # camera intrinsics directly.
+                "intrinsics": {
+                    "fx": SANA_WM_OUTPUT_WIDTH / 2,
+                    "fy": SANA_WM_OUTPUT_WIDTH / 2,
+                    "cx": SANA_WM_OUTPUT_WIDTH / 2,
+                    "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+                },
+            },
+        },
+        OmniDiffusionSamplingParams(
+            height=SANA_WM_OUTPUT_HEIGHT,
+            width=SANA_WM_OUTPUT_WIDTH,
+            num_frames=num_frames,
+            seed=0,
+            fps=16,
+            num_inference_steps=stage1_steps,
+            guidance_scale=1.0,
+            guidance_scale_provided=True,
+            output_type=output_type,
+            extra_args=extra_args,
+        ),
+    )
+
+    request_output = output[0] if isinstance(output, list) else output
+    assert isinstance(request_output, OmniRequestOutput)
+    assert request_output.error is None
+    assert request_output.images
+
+    frames = request_output.images[0]
+    video = _coerce_video_array(frames)
+    _assert_sana_wm_e2e_shape(video, output_type=output_type, num_frames=num_frames)
+    return video
+
+
+def test_sana_wm_native_generates_video() -> None:
+    output_type = os.environ.get("SANA_WM_E2E_OUTPUT_TYPE", "np")
+    video = _run_sana_wm_e2e(output_type=output_type)
+    if output_type == "latent":
+        assert video.ndim == 4
+    else:
+        assert video.shape[-1] in (3, 4)

tests/e2e/online_serving/test_sana_wm.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Online serving smoke for SANA-WM image-to-video (first-frame I2V via ``/v1/videos``).
+
+Boots the server straight from the model id — the Stage-1 repo's
+``model_index.json`` names ``SanaWmPipeline``, so the class resolves on its own
+— submits one async ``/v1/videos`` job with a first-frame reference image, and
+asserts the job completes and returns video bytes.
+
+From ``tests/``::
+
+    pytest -s -v e2e/online_serving/test_sana_wm.py -m "advanced_model and diffusion" --run-level=advanced_model
+"""
+
+import base64
+import json
+import os
+from io import BytesIO
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from vllm_omni.diffusion.models.sana_wm import (
+    SANA_WM_MODEL_ID,
+    SANA_WM_OUTPUT_HEIGHT,
+    SANA_WM_OUTPUT_WIDTH,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = SANA_WM_MODEL_ID
+PROMPT = "A slow forward camera move through a quiet city street."
+NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
+
+# Smoke knobs: a handful of frames / steps so the job finishes quickly.
+SMOKE_NUM_FRAMES = 9
+SMOKE_NUM_INFERENCE_STEPS = 2
+# First-frame I2V needs a camera trajectory; "w-<n>" is the forward-move action DSL.
+# Passed via the ``sana_wm`` form field (a JSON string parsed server-side); explicit
+# intrinsics avoid the optional Pi3X camera-calibration dependency.
+SANA_WM_PARAMS = {
+    "action": f"w-{SMOKE_NUM_FRAMES - 1}",
+    "translation_speed": 0.055,
+    "rotation_speed_deg": 1.2,
+    "intrinsics": {
+        "fx": SANA_WM_OUTPUT_WIDTH / 2,
+        "fy": SANA_WM_OUTPUT_WIDTH / 2,
+        "cx": SANA_WM_OUTPUT_WIDTH / 2,
+        "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+    },
+}
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _first_frame_data_url() -> str:
+    """A solid-color first frame encoded as a PNG ``data:`` URL for I2V."""
+    from PIL import Image
+
+    image = Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+    return f"data:image/png;base64,{encoded}"
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Single default server row."""
+    return [
+        pytest.param(
+            OmniServerParams(model=model),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_image_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Default SANA-WM I2V smoke: async ``/v1/videos`` job completes and returns video bytes."""
+    request_config = {
+        "model": omni_server.model,
+        "image_reference": _first_frame_data_url(),
+        "form_data": {
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            # The camera payload rides inside extra_params; there is no
+            # top-level sana_wm form field on the video endpoints, so sending
+            # one is silently dropped and the request fails in the
+            # preprocessor for missing camera control.
+            "extra_params": json.dumps({"sana_wm": SANA_WM_PARAMS}),
+            "height": SANA_WM_OUTPUT_HEIGHT,
+            "width": SANA_WM_OUTPUT_WIDTH,
+            "num_frames": SMOKE_NUM_FRAMES,
+            "fps": 16,
+            "num_inference_steps": SMOKE_NUM_INFERENCE_STEPS,
+            "guidance_scale": 5.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -0,0 +1,810 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""API-server surface guards for the OpenAI Omni entrypoint.
+
+Purpose
+-------
+Freeze the *observable* server contract so helper moves (P0.2), route
+extraction (P0.3), and later endpoint-family PRs cannot silently change:
+
+* which routes exist,
+* how upstream routes are overridden at assembly time,
+* unavailable / missing-handler envelopes,
+* shared engine-error JSON fields,
+* app-state keys wired by ``omni_init_app_state``.
+
+CPU-only (fakes / TestClient). No checked-in OpenAPI golden JSON.
+
+Common wrong changes these tests are meant to catch
+----------------------------------------------------
+* Route accidentally removed, renamed, method-changed, or double-registered
+  after a move (census / assembled-app checks).
+* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse).
+* Assembly forgets to strip upstream ``/v1/chat/completions`` (batch) /
+  ``/v1/models``, or also deletes unrelated upstream routes.
+* Profiler / Omni router not mounted; storage not started; Omni engine
+  exception handlers not registered.
+* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes.
+* Missing-handler WS/HTTP responses change status, JSON shape, or wording
+  (clients treat these as capability probes).
+* Engine error payloads lose ``request_id`` / ``error_stage_id``.
+* Pure-diffusion vs multi-stage init drops an ``app.state`` key, or leaves a
+  handler as ``None`` while the route is still registered.
+
+Not covered here (add elsewhere if needed)
+------------------------------------------
+* Deep per-endpoint success behavior (family UTs).
+* Byte-identical helper bodies after a move (P0.2 AST integrity).
+* Live GPU / speech byte-identity checks.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from fastapi.routing import APIRoute, APIWebSocketRoute
+from fastapi.testclient import TestClient
+from starlette.datastructures import State
+from starlette.requests import Request
+from starlette.websockets import WebSocketDisconnect
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
+
+from vllm_omni.entrypoints.openai import api_server
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
+# would invent ("HEAD", path) entries for every GET and fail the current census.
+# TODO: extend this set and the expected census if @router.head / @router.options appear.
+_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
+
+# Exact Omni router census on main. Update intentionally when routes change.
+_EXPECTED_ROUTER_ROUTES = {
+    ("POST", "/v1/chat/completions"),
+    ("POST", "/v1/chat/completions/batch"),
+    ("POST", "/v1/audio/speech"),
+    ("POST", "/v1/audio/speech/batch"),
+    ("POST", "/v1/audio/generate"),
+    ("GET", "/v1/audio/voices"),
+    ("POST", "/v1/audio/voices"),
+    ("DELETE", "/v1/audio/voices/{name}"),
+    ("WEBSOCKET", "/v1/audio/speech/stream"),
+    ("WEBSOCKET", "/v1/video/chat/stream"),
+    ("WEBSOCKET", "/v1/realtime/video"),
+    ("WEBSOCKET", "/v1/realtime"),
+    ("WEBSOCKET", "/v1/realtime/robot/openpi"),
+    ("WEBSOCKET", "/v1/duplex"),
+    ("GET", "/health"),
+    ("GET", "/v1/models"),
+    ("POST", "/v1/images/generations"),
+    ("POST", "/v1/images/edits"),
+    ("POST", "/v1/videos"),
+    ("POST", "/v1/videos/sync"),
+    ("GET", "/v1/videos"),
+    ("GET", "/v1/videos/{video_id}"),
+    ("DELETE", "/v1/videos/{video_id}"),
+    ("GET", "/v1/videos/{video_id}/content"),
+    ("POST", "/v1/omni/sleep"),
+    ("POST", "/v1/omni/wakeup"),
+}
+_EXPECTED_PROFILER_ROUTES = {
+    ("POST", "/start_profile"),
+    ("POST", "/stop_profile"),
+}
+
+# App-state keys that must remain available across refactor phases.
+# Presence is not enough: "not wired" is written as ``state.X = None``.
+_DIFFUSION_APP_STATE_KEYS = {
+    "engine_client",
+    "log_stats",
+    "args",
+    "sleeping_stages",
+    "stage_configs",
+    "vllm_config",
+    "diffusion_engine",
+    "openai_serving_models",
+    "serving_tokenization",
+    "openai_serving_chat",
+    "openai_serving_chat_batch",
+    "openai_serving_speech",
+    "openai_serving_audio_generate",
+    "openai_serving_video",
+    "openai_streaming_video_output",
+    "openai_serving_duplex",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+    "openai_serving_realtime_robot",
+    "enable_server_load_tracking",
+    "server_load_metrics",
+}
+_DIFFUSION_MUST_BE_NONE = {
+    "vllm_config",
+    "serving_tokenization",
+    "openai_serving_duplex",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+}
+_DIFFUSION_MUST_BE_WIRED = _DIFFUSION_APP_STATE_KEYS - _DIFFUSION_MUST_BE_NONE
+_MULTISTAGE_APP_STATE_KEYS = {
+    "engine_client",
+    "log_stats",
+    "args",
+    "sleeping_stages",
+    "stage_configs",
+    "vllm_config",
+    "openai_serving_models",
+    "serving_tokenization",
+    "online_renderer",
+    "openai_serving_chat",
+    "openai_serving_chat_batch",
+    "openai_serving_speech",
+    "openai_serving_audio_generate",
+    "openai_streaming_speech",
+    "openai_streaming_video",
+    "openai_serving_duplex",
+    "openai_serving_realtime",
+    "openai_serving_video",
+    "openai_serving_realtime_robot",
+    "enable_server_load_tracking",
+    "server_load_metrics",
+}
+_MULTISTAGE_MUST_BE_NONE = {
+    "openai_serving_duplex",
+    "openai_serving_realtime_robot",
+}
+_MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE
+
+
+def _route_entries(routes) -> list[tuple[str, str]]:
+    entries: list[tuple[str, str]] = []
+    for route in routes:
+        if isinstance(route, APIRoute):
+            for method in sorted((route.methods or set()) & _HTTP_METHODS):
+                entries.append((method, route.path))
+        elif isinstance(route, APIWebSocketRoute):
+            entries.append(("WEBSOCKET", route.path))
+    return entries
+
+
+def _route_manifest(routes) -> set[tuple[str, str]]:
+    return set(_route_entries(routes))
+
+
+def _assert_unique_route_entries(routes) -> None:
+    entries = _route_entries(routes)
+    duplicates = sorted({item for item in entries if entries.count(item) > 1})
+    assert not duplicates, f"duplicate route registrations: {duplicates}"
+
+
+def _matching_http_routes(routes, method: str, path: str) -> list[APIRoute]:
+    return [
+        route
+        for route in routes
+        if isinstance(route, APIRoute) and route.path == path and method in (route.methods or set())
+    ]
+
+
+def _single_http_route(routes, method: str, path: str) -> APIRoute:
+    matches = _matching_http_routes(routes, method, path)
+    assert len(matches) == 1, f"expected exactly one {method} {path}, got {len(matches)}"
+    return matches[0]
+
+
+def _assert_app_state_snapshot(
+    state,
+    *,
+    expected_keys: set[str],
+    must_be_wired: set[str],
+    must_be_none: set[str],
+) -> None:
+    assert must_be_wired.isdisjoint(must_be_none)
+    assert must_be_wired | must_be_none <= expected_keys
+    present = {key for key in expected_keys if hasattr(state, key)}
+    assert present == expected_keys
+    not_wired = sorted(key for key in must_be_wired if getattr(state, key) is None)
+    assert not not_wired, f"app.state keys registered but not wired: {not_wired}"
+    unexpectedly_set = sorted(key for key in must_be_none if getattr(state, key) is not None)
+    assert not unexpectedly_set, f"app.state keys expected to stay None: {unexpectedly_set}"
+
+
+def _minimal_args(**overrides) -> SimpleNamespace:
+    args = SimpleNamespace(
+        tool_parser_plugin=None,
+        reasoning_parser_plugin=None,
+        reasoning_parser=None,
+        structured_outputs_config=SimpleNamespace(reasoning_parser=None),
+        enable_ssl_refresh=False,
+        host="127.0.0.1",
+        port=8000,
+        uvicorn_log_level="info",
+        disable_uvicorn_access_log=True,
+        ssl_keyfile=None,
+        ssl_certfile=None,
+        ssl_ca_certs=None,
+        ssl_cert_reqs=0,
+        ssl_ciphers=None,
+        h11_max_incomplete_event_size=None,
+        h11_max_header_count=None,
+        log_config_file=None,
+        disable_access_log_for_endpoints=None,
+        disable_log_stats=True,
+        enable_log_requests=False,
+        max_log_len=None,
+        served_model_name=None,
+        model="demo-model",
+        enable_server_load_tracking=False,
+        log_error_stack=False,
+        chat_template=None,
+        chat_template_content_format="auto",
+        trust_request_chat_template=False,
+        tool_server=None,
+        lora_modules=None,
+        response_role="assistant",
+        default_chat_template_kwargs=None,
+        return_tokens_as_token_ids=False,
+        enable_auto_tool_choice=False,
+        exclude_tools_when_tool_choice_none=False,
+        tool_call_parser=None,
+        enable_prompt_tokens_details=False,
+        enable_force_include_usage=False,
+        enable_log_outputs=False,
+        enable_log_deltas=False,
+        tokens_only=False,
+        shutdown_timeout=0,
+        stage_configs_path=None,
+        deploy_config=None,
+    )
+    for key, value in overrides.items():
+        setattr(args, key, value)
+    return args
+
+
+def _request_for(app: FastAPI, *, method: str = "GET", path: str = "/") -> Request:
+    return Request(
+        {
+            "type": "http",
+            "method": method,
+            "path": path,
+            "headers": [],
+            "app": app,
+            "state": {},
+        }
+    )
+
+
+class _FakeSocket:
+    def __init__(self) -> None:
+        self.closed = False
+
+    def close(self) -> None:
+        self.closed = True
+
+
+class _FakeEngineClient:
+    def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None:
+        self.stage_configs = stage_configs if stage_configs is not None else []
+        self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {}
+        self.model_config = SimpleNamespace()
+        self.input_processor = object()
+        self.renderer = object()
+        self.errored = False
+        self.vllm_config = vllm_config
+
+    async def get_vllm_config(self):
+        return self.vllm_config
+
+    async def get_supported_tasks(self) -> tuple[str, ...]:
+        return ("generate",)
+
+    async def get_tokenizer(self):
+        return SimpleNamespace(chat_template="dummy")
+
+    async def check_health(self) -> None:
+        return None
+
+
+def _marker(name: str):
+    return SimpleNamespace(name=name)
+
+
+def test_router_route_manifest_is_explicit_and_complete() -> None:
+    """Lock the Omni + profiler router census.
+
+    Fails if a route is added/removed/renamed, a method changes, or a
+    WebSocket path disappears during helper/router moves. Also fails if the
+    same public route is registered more than once, which a set manifest alone
+    would collapse.
+    """
+    # TODO(P0.3): retarget this object-level census when routes move onto family
+    # routers. The assembled-app census is the client-facing lock that must survive.
+    _assert_unique_route_entries(api_server.router.routes)
+    _assert_unique_route_entries(api_server.profiler_router.routes)
+    assert _route_manifest(api_server.router.routes) == _EXPECTED_ROUTER_ROUTES
+    assert _route_manifest(api_server.profiler_router.routes) == _EXPECTED_PROFILER_ROUTES
+
+
+def test_router_openapi_paths_cover_http_manifest() -> None:
+    """Lock that OpenAPI advertises every HTTP census route/method.
+
+    Fails if a decorator still registers a route but OpenAPI loses the path
+    (or method), e.g. after moving route bodies / response models.
+    """
+    app = FastAPI()
+    app.include_router(api_server.router)
+    app.include_router(api_server.profiler_router)
+    paths = app.openapi().get("paths", {})
+
+    http_routes = {
+        (method, path) for method, path in _EXPECTED_ROUTER_ROUTES | _EXPECTED_PROFILER_ROUTES if method != "WEBSOCKET"
+    }
+    for method, path in sorted(http_routes):
+        assert path in paths, f"missing OpenAPI path {path}"
+        assert method.lower() in paths[path], f"missing OpenAPI method {method} {path}"
+
+
+@pytest.mark.asyncio
+async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None:
+    """Lock worker assembly: override, mount, storage, handlers, full census.
+
+    Fails if assembly stops replacing upstream chat/batch/models, deletes
+    unrelated upstream routes, forgets Omni/profiler mounts, skips storage
+    start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError`` handlers.
+    """
+    captured: dict[str, object] = {}
+
+    def fake_build_openai_app(args, supported_tasks):
+        app = FastAPI()
+
+        @app.post("/v1/chat/completions")
+        async def upstream_chat():
+            return {"owner": "upstream"}
+
+        @app.post("/v1/chat/completions/batch")
+        async def upstream_batch():
+            return {"owner": "upstream"}
+
+        @app.get("/v1/models")
+        async def upstream_models():
+            return {"owner": "upstream"}
+
+        @app.get("/v1/upstream-only")
+        async def upstream_only():
+            return {"owner": "upstream"}
+
+        captured["supported_tasks"] = supported_tasks
+        return app
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **_kwargs):
+        yield _FakeEngineClient(
+            stage_configs=[{"engine_args": {"profiler_config": {"profiler": "torch"}}}],
+            vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(_api_process_rank=0)),
+        )
+
+    async def fake_omni_init_app_state(engine_client, state, args):
+        state.engine_client = engine_client
+        state.initialized_by_omni = True
+
+    async def fake_storage_start():
+        captured["storage_started"] = True
+
+    async def fake_serve_http(app, **_kwargs):
+        captured["served_app"] = app
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+    monkeypatch.setattr(
+        api_server,
+        "shutdown_unsupported_routes",
+        lambda app, restrictions: captured.setdefault("restrictions", restrictions),
+    )
+    monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=fake_storage_start))
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    sock = _FakeSocket()
+    await api_server.omni_run_server_worker("127.0.0.1:8000", sock, _minimal_args())
+
+    served_wrapper = captured["served_app"]
+    # Timestamp middleware wraps the FastAPI app.
+    assert hasattr(served_wrapper, "_inner")
+    served_app = served_wrapper._inner
+    routes = served_app.routes
+    _assert_unique_route_entries(routes)
+    manifest = _route_manifest(routes)
+
+    assert captured["supported_tasks"] == ("generate",)
+    assert captured["storage_started"] is True
+    assert captured["restrictions"] == {}
+    assert sock.closed is True
+    assert served_app.state.initialized_by_omni is True
+
+    # Upstream override: unrelated route kept; Omni owns overridden pairs once.
+    # Endpoint identity ensures the surviving route is the Omni handler, not
+    # the fake upstream handler with the same method/path.
+    assert ("GET", "/v1/upstream-only") in manifest
+    assert _single_http_route(routes, "POST", "/v1/chat/completions").endpoint is api_server.create_chat_completion
+    assert (
+        _single_http_route(routes, "POST", "/v1/chat/completions/batch").endpoint
+        is api_server.create_batch_chat_completion
+    )
+    assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models
+
+    # Assembled app includes Omni router HTTP/WS surface + profiler.
+    for item in _EXPECTED_ROUTER_ROUTES:
+        assert item in manifest, f"assembled app missing {item}"
+    assert ("POST", "/start_profile") in manifest
+    assert ("POST", "/stop_profile") in manifest
+
+    assert served_app.exception_handlers[EngineGenerateError]
+    assert served_app.exception_handlers[EngineDeadError]
+
+    # Light OpenAPI check on the assembled FastAPI app (not the ASGI wrapper).
+    openapi_paths = served_app.openapi().get("paths", {})
+    assert "/v1/chat/completions" in openapi_paths
+    assert "/v1/models" in openapi_paths
+    assert "/health" in openapi_paths
+
+
+@pytest.mark.asyncio
+async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch) -> None:
+    """Lock outermost timestamp middleware behavior.
+
+    Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes
+    start getting stamped / mutated incorrectly.
+    """
+    captured: dict[str, object] = {}
+
+    def fake_build_openai_app(args, supported_tasks):
+        return FastAPI()
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **_kwargs):
+        yield _FakeEngineClient()
+
+    async def fake_omni_init_app_state(engine_client, state, args):
+        state.engine_client = engine_client
+
+    async def fake_serve_http(app, **_kwargs):
+        captured["served_app"] = app
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+    monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_a, **_k: None)
+    monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=lambda: asyncio.sleep(0)))
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args())
+    middleware = captured["served_app"]
+
+    http_scope = {"type": "http", "state": {}}
+    ws_scope = {"type": "websocket"}
+
+    async def _receive():
+        return {"type": "http.disconnect"}
+
+    async def _send(_message):
+        return None
+
+    called = {"http": False, "websocket": False}
+
+    async def inner(scope, receive, send):
+        called[scope["type"]] = True
+
+    middleware._inner = inner
+    await middleware(http_scope, _receive, _send)
+    await middleware(ws_scope, _receive, _send)
+
+    assert called["http"] is True
+    assert called["websocket"] is True
+    assert "request_timestamp" in http_scope["state"]
+    assert isinstance(http_scope["state"]["request_timestamp"], float)
+    assert "request_timestamp" not in ws_scope.get("state", {})
+
+
+@pytest.mark.parametrize(
+    ("path", "payload"),
+    [
+        ("/v1/audio/speech/stream", {"type": "error", "message": "Streaming speech is not available"}),
+        ("/v1/video/chat/stream", {"type": "error", "message": "Streaming video chat is not available"}),
+        ("/v1/realtime/video", {"type": "error", "message": "Streaming video generation is not available"}),
+        ("/v1/realtime", {"type": "error", "error": "Realtime API is not available", "code": "unsupported"}),
+        (
+            "/v1/realtime/robot/openpi",
+            {"type": "error", "error": "Robot policy not available", "code": "unsupported"},
+        ),
+        ("/v1/duplex", {"type": "error", "error": "Duplex API is not available", "code": "unsupported"}),
+    ],
+)
+def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, payload: dict[str, str]) -> None:
+    """Lock missing-handler WebSocket capability probes.
+
+    Fails if unavailable JSON changes (keys/wording/code) or the socket stays
+    open after the error frame — common when WS wrappers move between modules.
+    """
+    app = FastAPI()
+    app.include_router(api_server.router)
+
+    with TestClient(app) as client:
+        with client.websocket_connect(path) as websocket:
+            assert websocket.receive_json() == payload
+            with pytest.raises(WebSocketDisconnect):
+                websocket.receive_text()
+
+
+def test_health_without_engine_returns_stable_unhealthy_response() -> None:
+    """Lock ``/health`` with no engine initialized.
+
+    Fails if status drifts from 503 or the unhealthy payload shape/reason
+    string changes during health-route extraction.
+    """
+    app = FastAPI()
+    response = asyncio.run(api_server.health(_request_for(app, path="/health")))
+
+    assert response.status_code == 503
+    assert json.loads(response.body) == {
+        "status": "unhealthy",
+        "reason": "No engine initialized",
+    }
+
+
+def test_models_without_handler_returns_empty_openai_list() -> None:
+    """Lock ``/v1/models`` when ``openai_serving_models`` is unset.
+
+    Fails if the empty listing becomes an error, or the OpenAI list envelope
+    changes (``object`` / ``data``).
+    """
+    app = FastAPI()
+    response = asyncio.run(api_server.show_available_models(_request_for(app, path="/v1/models")))
+
+    assert response.status_code == 200
+    assert json.loads(response.body) == {"object": "list", "data": []}
+
+
+def test_images_generation_without_engine_preserves_service_unavailable_error() -> None:
+    """Lock images generation when the multi-stage engine is missing.
+
+    Fails if the 503 detail string changes — clients/tests use it as a
+    capability signal when bootstrap failed.
+    """
+    app = FastAPI()
+    raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+    request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.generate_images(request, raw_request))
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model."
+
+
+def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None:
+    """Lock images generation when chat serving was not wired.
+
+    Fails if multi-stage images stop requiring ``openai_serving_chat``, or the
+    503 detail changes after images/chat ownership splits.
+    """
+    app = FastAPI()
+    app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()])
+    app.state.stage_configs = app.state.engine_client.stage_configs
+    app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")])
+    app.state.openai_serving_chat = None
+    app.state.args = SimpleNamespace(max_generated_image_size=None)
+
+    raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+    request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+    monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.generate_images(request, raw_request))
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == "openai_serving_chat is not initialized for multi-stage image generation."
+
+
+def test_speech_without_handler_preserves_not_found_http_error() -> None:
+    """Lock speech HTTP when ``openai_serving_speech`` is unset.
+
+    Fails if missing speech becomes a different status/detail after TTS route
+    wrappers move out of ``api_server``.
+    """
+    app = FastAPI()
+    app.state.openai_serving_speech = None
+    raw_request = _request_for(app, method="POST", path="/v1/audio/speech")
+    request = api_server.OpenAICreateSpeechRequest(input="hello", model="demo-model", voice="alloy")
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(api_server.create_speech(request, raw_request))
+
+    assert exc_info.value.status_code == 404
+    assert exc_info.value.detail == "The model does not support Speech API"
+
+
+def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
+    """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
+
+    Fails if ``request_id`` / ``error_stage_id`` are dropped when error helpers
+    move (P0.2) or exception-handler wiring changes.
+
+    Exercises the envelope through the registered handler, not the private
+    ``_create_engine_error_json_response`` helper, so relocating that helper
+    in P0.2 is not a false red.
+    """
+    app = FastAPI()
+    # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+    # leaves api_server (red here is a rename, not a surface regression).
+    api_server._register_omni_exception_handlers(app)
+    app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+    app.state.server = object()
+    app.state.args = SimpleNamespace(log_error_stack=False)
+
+    req = _request_for(app, method="POST", path="/v1/chat/completions")
+    req.state.request_metadata = SimpleNamespace(request_id="req-123")
+
+    monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+    exc = EngineGenerateError("boom")
+    exc.error_stage_id = 2  # type: ignore[attr-defined]
+    handler = app.exception_handlers[EngineGenerateError]
+    response = asyncio.run(handler(req, exc))
+    payload = json.loads(response.body)
+
+    assert "error" in payload
+    assert payload["error"]["request_id"] == "req-123"
+    assert payload["error"]["error_stage_id"] == 2
+    assert response.status_code >= 400
+
+
+def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None:
+    """Lock Omni ``EngineDeadError`` handler registration + JSON envelope.
+
+    Fails if registration is skipped or the handler stops returning OpenAI-style
+    JSON with ``request_id``.
+    """
+    app = FastAPI()
+    # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+    # leaves api_server (red here is a rename, not a surface regression).
+    api_server._register_omni_exception_handlers(app)
+    app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+    app.state.server = object()
+    app.state.args = SimpleNamespace(log_error_stack=False)
+
+    monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+    handler = app.exception_handlers[EngineDeadError]
+    req = _request_for(app)
+    req.state.request_metadata = SimpleNamespace(request_id="dead-1")
+    response = asyncio.run(handler(req, EngineDeadError()))
+    payload = json.loads(response.body)
+    assert payload["error"]["request_id"] == "dead-1"
+    assert response.status_code >= 400
+
+
+@pytest.mark.asyncio
+async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
+    """Lock pure-diffusion ``app.state`` keys after init, including live vs None.
+
+    Fails if bootstrap drops keys route owners still read (e.g. video/speech/
+    models), wires a handler that must stay None, or leaves a required
+    handler as None.
+    """
+    stage = SimpleNamespace(engine_args={})
+    engine = _FakeEngineClient(stage_configs=[stage])
+
+    def _for_diffusion_factory(label: str):
+        @classmethod
+        def _factory(cls, *args, **kwargs):
+            return _marker(label)
+
+        return _factory
+
+    monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
+    monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
+    monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
+    monkeypatch.setattr(
+        api_server.OmniOpenAIServingAudioGenerate,
+        "for_diffusion",
+        _for_diffusion_factory("audio_generate"),
+    )
+    monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
+    monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
+    monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _for_diffusion_factory("speech"))
+    monkeypatch.setattr(
+        api_server.ServingRealtimeRobotOpenPI,
+        "create_policy_server",
+        classmethod(lambda cls, *a, **k: _marker("openpi")),
+    )
+
+    state = State()
+    await api_server.omni_init_app_state(engine, state, _minimal_args())
+
+    _assert_app_state_snapshot(
+        state,
+        expected_keys=_DIFFUSION_APP_STATE_KEYS,
+        must_be_wired=_DIFFUSION_MUST_BE_WIRED,
+        must_be_none=_DIFFUSION_MUST_BE_NONE,
+    )
+    assert state.diffusion_engine is engine
+
+
+@pytest.mark.asyncio
+async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
+    """Lock multi-stage ``app.state`` keys after init, including live vs None.
+
+    Fails if chat/speech/video/realtime/tokenization keys disappear or are
+    left as None after init refactors — classic “route still registered,
+    handler never wired” failure mode.
+    """
+    engine = _FakeEngineClient(
+        stage_configs=[object(), object()],
+        vllm_config=SimpleNamespace(
+            lora_config=None,
+            model_config=SimpleNamespace(),
+            parallel_config=SimpleNamespace(_api_process_rank=0),
+        ),
+    )
+
+    class _FakeModels:
+        def __init__(self, *args, **kwargs):
+            self.base_model_paths = kwargs.get("base_model_paths") or []
+
+        async def init_static_loras(self):
+            return None
+
+    class _FakeCtor:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def warmup(self):
+            return None
+
+    class _FakeSpeech(_FakeCtor):
+        async def warmup(self):
+            return None
+
+    monkeypatch.setattr(api_server, "load_chat_template", lambda *_a, **_k: None)
+    monkeypatch.setattr(api_server, "process_lora_modules", lambda modules, _defaults: modules or [])
+    monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
+    monkeypatch.setattr(api_server, "OnlineRenderer", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingResponses", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingChat", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingChatBatch", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingCompletion", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingPooling", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingEmbedding", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingClassification", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingScores", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingTokenization", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingTranscription", _FakeCtor)
+    monkeypatch.setattr(api_server, "OpenAIServingTranslation", _FakeCtor)
+    monkeypatch.setattr(api_server, "AnthropicServingMessages", _FakeCtor)
+    monkeypatch.setattr(api_server, "ServingTokens", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingSpeech", _FakeSpeech)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingAudioGenerate", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniStreamingSpeechHandler", _FakeCtor)
+    monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video"))
+    monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor)
+    monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor)
+    monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False)
+
+    state = State()
+    await api_server.omni_init_app_state(engine, state, _minimal_args())
+
+    _assert_app_state_snapshot(
+        state,
+        expected_keys=_MULTISTAGE_APP_STATE_KEYS,
+        must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
+        must_be_none=_MULTISTAGE_MUST_BE_NONE,
+    )

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -22,6 +22,10 @@ from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
     IndexTTS25Adapter,
     indextts2_conditioning_cache_salt,
 )
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
+    MossTTSAdapter,
+    MossTTSNanoAdapter,
+)
 from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
 from vllm_omni.model_executor.models.indextts2 import prompt_utils
 from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
@@ -94,6 +98,20 @@ def test_all_adapters_are_ar_or_diffusion():
         assert cls.backend in ("ar", "diffusion")
 
 
+@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
+def test_moss_tts_applies_request_max_new_tokens(adapter_cls):
+    adapter = adapter_cls(SimpleNamespace(server=object()))
+    stage_defaults = [SimpleNamespace(max_tokens=4096)]
+
+    overridden = adapter.apply_sampling_overrides(
+        stage_defaults,
+        SimpleNamespace(max_new_tokens=512),
+    )
+
+    assert overridden[0].max_tokens == 512
+    assert stage_defaults[0].max_tokens == 4096
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/entrypoints/openai_api/test_video_server.py

@@ -2168,6 +2168,31 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
     assert captured.quality == "high"
 
 
+def test_sync_sana_wm_extra_params_payload_passes_to_engine_prompt(test_client, mocker: MockerFixture):
+    _mock_encode_video_bytes(mocker)
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={
+            "prompt": "drive forward",
+            "extra_params": json.dumps(
+                {
+                    "sana_wm": {"action": "d-4", "rotation_speed_deg": 1.5},
+                    "sana_wm_native_max_tokens": 30000,
+                }
+            ),
+        },
+    )
+
+    assert response.status_code == 200
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured_params = engine.captured_sampling_params_list[0]
+    # The camera block reaches the model through extra_args; the Sana preprocess
+    # hook is what lifts it onto the prompt.
+    assert captured_params.extra_args["sana_wm"]["action"] == "d-4"
+    assert captured_params.extra_args["sana_wm"]["rotation_speed_deg"] == 1.5
+    assert captured_params.extra_args["sana_wm_native_max_tokens"] == 30000
+
+
 def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
     """Frame interpolation parameters should be forwarded on the sync path."""
     encode_mock = _mock_encode_video_bytes(mocker)

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -58,6 +58,23 @@ def test_default_stage_devices_from_sequence_parallel():
     assert stage_cfg["runtime"]["devices"] == "0,1,2,3"
 
 
+def test_default_stage_devices_and_dp_from_num_gpus():
+    """Resolve omitted DP before deriving devices from an explicit GPU count."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "num_gpus": 8,
+            "tensor_parallel_size": 2,
+            "ulysses_degree": 2,
+        }
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config.data_parallel_size == 2
+    assert parallel_config.world_size == 8
+    assert stage_cfg["engine_args"]["num_gpus"] == 8
+    assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
+
+
 def test_default_stage_config_uses_parallel_size_kwargs():
     """Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -76,8 +93,8 @@ def test_default_stage_config_uses_parallel_size_kwargs():
     assert parallel_config.enable_expert_parallel is True
 
 
-def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
-    """Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
+def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
+    """Keep omitted DP unresolved until the runtime WORLD size is known."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
         {
             "pipeline_parallel_size": None,
@@ -92,7 +109,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
     assert parallel_config.pipeline_parallel_size == 1
-    assert parallel_config.data_parallel_size == 1
+    assert parallel_config.data_parallel_size is None
     assert parallel_config.tensor_parallel_size == 1
     assert parallel_config.enable_expert_parallel is False
     assert stage_cfg["engine_args"]["enforce_eager"] is False
@@ -328,9 +345,6 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--omni",
             "--enable-distributed-layerwise-offload",
             "--dlo-no-use-allgather",
-            "--dlo-use-host-weight-cache",
-            "--dlo-host-weight-cache-pin-limit-gib",
-            "80",
             "--dlo-resident-layers",
             "20",
         ]
@@ -342,15 +356,9 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
-    assert args.dlo_use_host_weight_cache is True
-    assert args.dlo_host_weight_cache_pin_limit_gib == 80
     assert args.dlo_resident_layers == 20
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
-    assert engine_args["dlo_use_host_weight_cache"] is True
-    assert engine_args["dlo_host_weight_cache_dir"] is None
-    assert engine_args["dlo_host_weight_cache_lock_timeout"] == 600.0
-    assert engine_args["dlo_host_weight_cache_pin_limit_gib"] == 80
     assert engine_args["dlo_resident_layers"] == 20
 
 

tests/model_executor/model_loader/test_resolve_model_path.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for resolve_model_to_local_path."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import huggingface_hub
+import pytest
+from huggingface_hub.errors import LocalEntryNotFoundError
+
+from vllm_omni.model_executor.model_loader import weight_utils
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Both resolution paths are patched in every test, so nothing is ever fetched
+REPO_ID = "test-org/test-model"
+RESOLVED = "/cache/snapshots/deadbeef"
+
+
+@pytest.fixture
+def paths(mocker):
+    """The resolver's two ways out, both patched."""
+    return SimpleNamespace(
+        cache_only=mocker.patch.object(huggingface_hub, "snapshot_download", return_value=RESOLVED),
+        download=mocker.patch.object(weight_utils, "download_weights_from_hf_specific", return_value=RESOLVED),
+    )
+
+
+@pytest.mark.parametrize("allow_download", [False, True])
+def test_local_directory_short_circuits(tmp_path, paths, allow_download):
+    """A local dir resolves to itself without either path running."""
+    assert resolve_model_to_local_path(str(tmp_path), allow_download=allow_download) == str(tmp_path)
+
+    paths.cache_only.assert_not_called()
+    paths.download.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "allow_patterns", "cache_dir"),
+    [
+        ({}, None, None),
+        ({"allow_patterns": ["subdir/*"], "cache_dir": "/cache"}, ["subdir/*"], "/cache"),
+    ],
+)
+def test_repo_id_is_cache_only_by_default(paths, kwargs, allow_patterns, cache_dir):
+    """Without `allow_download` nothing is fetched, and narrowing is forwarded."""
+    assert resolve_model_to_local_path(REPO_ID, **kwargs) == RESOLVED
+
+    paths.cache_only.assert_called_once_with(
+        REPO_ID,
+        cache_dir=cache_dir,
+        allow_patterns=allow_patterns,
+        local_files_only=True,
+    )
+    paths.download.assert_not_called()
+
+
+def test_cache_miss_propagates(paths):
+    """A failed resolution raises instead of degrading into the raw repo id."""
+    paths.cache_only.side_effect = LocalEntryNotFoundError("cold cache")
+
+    with pytest.raises(LocalEntryNotFoundError):
+        resolve_model_to_local_path(REPO_ID)
+
+    paths.download.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "allow_patterns"),
+    [
+        ({}, ["*"]),
+        ({"allow_patterns": ["subdir/model*.safetensors"]}, ["subdir/model*.safetensors"]),
+    ],
+)
+def test_download_routes_through_the_locked_helper(paths, kwargs, allow_patterns):
+    """Downloads go through the vLLM helper holding the node-wide lock.
+
+    That helper is also what honors `VLLM_USE_MODELSCOPE`;
+    calling `snapshot_download` directly would drop both.
+    """
+    assert resolve_model_to_local_path(REPO_ID, allow_download=True, cache_dir="/cache", **kwargs) == RESOLVED
+
+    paths.download.assert_called_once_with(
+        model_name_or_path=REPO_ID,
+        cache_dir="/cache",
+        allow_patterns=allow_patterns,
+        require_all=True,
+    )
+    paths.cache_only.assert_not_called()

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -7,8 +7,14 @@ import pytest
 import torch
 
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
-from vllm_omni.model_executor.models.common.ming.fm import Solver
+from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler
+from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT
+from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss
+from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, get_epss_timesteps
+from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
+from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
 from vllm_omni.model_executor.models.ming_tts.constants import (
     AGGREGATOR_HIDDEN_SIZE,
     HISTORY_PATCH_SIZE,
@@ -24,7 +30,6 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
     SPEAKER_EMBEDDING_DIM,
     VAE_PATCH_SIZE,
 )
-from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
 from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
 from vllm_omni.model_executor.models.ming_tts.speaker_extractor import _resolve_speaker_embeddings
 from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
@@ -44,7 +49,87 @@ def test_ming_tts_audio_vae_uses_common_config():
 
 def test_ming_tts_cfm_solver_uses_common_implementation():
     """Ming dense imports the shared solver implementation directly."""
-    assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
+    assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.cfm_solver"
+
+
+@pytest.mark.parametrize(
+    ("symbol", "expected_module"),
+    [
+        (Aggregator, "vllm_omni.model_executor.models.common.ming.aggregator"),
+        (DiT, "vllm_omni.model_executor.models.common.ming.cfm_dit"),
+        (FlowLoss, "vllm_omni.model_executor.models.common.ming.cfm_head"),
+        (CFMGraphExecutor, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
+        (CFMSampler, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
+        (SpeakerEmbeddingExtractor, "vllm_omni.model_executor.models.common.ming.speaker_extractor"),
+        (SpkembExtractor, "vllm_omni.model_executor.models.common.ming.spk_embedding"),
+    ],
+)
+def test_ming_tts_reusable_modules_live_in_common(symbol, expected_module):
+    """The pieces a second Ming-family TTS model reuses live under common/ming.
+
+    Guards the promotion: if an implementation drifts back into `ming_tts`
+    (or into another model package), a second model would have to reach across
+    packages to get it.
+    """
+    assert symbol.__module__ == expected_module
+
+
+def test_ming_flash_talker_uses_the_shared_aggregator():
+    """Ming flash no longer carries its own copy of the aggregator.
+
+    Both Ming packages ran near-identical aggregators one directory apart,
+    which is the duplication this module layout exists to remove.
+    """
+    from vllm_omni.model_executor.models.ming_flash_omni import talker_module
+
+    assert talker_module.Aggregator is Aggregator
+    assert not any(
+        isinstance(obj, type) and obj.__module__ == talker_module.__name__ and name == "Aggregator"
+        for name, obj in vars(talker_module).items()
+    )
+
+
+def test_shared_aggregator_matches_the_ming_flash_construction():
+    """The shared aggregator reproduces Ming flash's former default geometry.
+
+    `talker_module.Aggregator` defaulted to `in_channels=64` while the
+    shared class defaults to 4, so the call site now passes `latent_dim`
+    explicitly. This pins the resulting shapes rather than the default.
+    """
+    latent_dim, llm_hidden = 64, 32
+    aggregator = Aggregator(
+        in_channels=latent_dim,
+        llm_input_dim=llm_hidden,
+        hidden_size=16,
+        depth=1,
+        num_heads=2,
+    )
+
+    assert aggregator.x_embedder.in_features == latent_dim
+    assert aggregator.final_layer.linear.out_features == llm_hidden
+
+    out = aggregator(torch.randn(1, 3, latent_dim))
+    # Keeps only the CLS projection: [B, T, D] -> [B, 1, llm_hidden].
+    assert out.shape == (1, 1, llm_hidden)
+
+
+def test_shared_aggregator_rejects_mismatched_latent_width():
+    """Wrong-width input fails loudly instead of silently projecting incorrect."""
+    aggregator = Aggregator(in_channels=64, llm_input_dim=32, hidden_size=16, depth=1, num_heads=2)
+
+    with pytest.raises(ValueError, match="feature dim mismatch"):
+        aggregator(torch.randn(1, 3, 4))
+
+
+def test_cfm_solver_owns_the_timestep_schedule():
+    """get_epss_timesteps sits with the solver, not with the DiT blocks."""
+    from vllm_omni.model_executor.models.common.ming import dit_blocks
+
+    assert not hasattr(dit_blocks, "get_epss_timesteps")
+
+    t = get_epss_timesteps(10, device="cpu", dtype=torch.float32)
+    assert t.shape == (11,)
+    assert float(t[0]) == 0.0 and float(t[-1]) == 1.0
 
 
 def test_ming_tts_flowloss_preserves_checkpoint_prefix():

tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniCPM-o 4.5 streaming audio cache compatibility."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.whisper.modeling_whisper import WhisperConfig
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    WHISPER_ATTENTION_CLASSES,
+    MiniCPMWhisperEncoder,
+    MiniCPMWhisperEncoderLayer,
+    _get_audio_cache_length,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _LegacyWhisperAttention(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.seen_cache = None
+
+    def forward(self, hidden_states, past_key_value=None, **kwargs):
+        self.seen_cache = past_key_value
+        return hidden_states, None, past_key_value
+
+
+class _CurrentWhisperAttention(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.seen_cache = None
+
+    def forward(self, hidden_states, past_key_values=None, **kwargs):
+        self.seen_cache = past_key_values
+        return hidden_states, None
+
+
+@pytest.mark.parametrize(
+    ("implementation", "attention_cls"),
+    [
+        ("test_legacy_cache", _LegacyWhisperAttention),
+        ("test_current_cache", _CurrentWhisperAttention),
+    ],
+)
+def test_whisper_attention_preserves_streaming_cache(
+    monkeypatch: pytest.MonkeyPatch,
+    implementation: str,
+    attention_cls: type[nn.Module],
+) -> None:
+    monkeypatch.setitem(WHISPER_ATTENTION_CLASSES, implementation, attention_cls)
+    config = WhisperConfig(
+        d_model=4,
+        encoder_attention_heads=1,
+        encoder_ffn_dim=8,
+        encoder_layers=1,
+    )
+    config._attn_implementation = implementation
+    layer = MiniCPMWhisperEncoderLayer(config, layer_idx=0)
+    cache = object()
+
+    outputs = layer(
+        torch.randn(1, 2, config.d_model),
+        attention_mask=None,
+        layer_head_mask=None,
+        past_key_values=cache,
+        use_cache=True,
+    )
+
+    assert layer.self_attn.seen_cache is cache
+    assert outputs[-1] is cache
+
+
+def test_audio_cache_length_uses_current_cache_api() -> None:
+    dynamic_cache = SimpleNamespace(get_seq_length=lambda: 7)
+    encoder_decoder_cache = SimpleNamespace(self_attention_cache=dynamic_cache)
+
+    assert _get_audio_cache_length(encoder_decoder_cache) == 7
+
+
+def test_audio_cache_length_supports_legacy_cache() -> None:
+    key = torch.zeros(1, 1, 5, 2)
+    value = torch.zeros_like(key)
+
+    assert _get_audio_cache_length(((key, value),)) == 5
+
+
+def _tiny_encoder_config(attn_implementation: str) -> WhisperConfig:
+    config = WhisperConfig(
+        num_mel_bins=16,
+        d_model=32,
+        encoder_layers=2,
+        encoder_attention_heads=2,
+        encoder_ffn_dim=64,
+        max_source_positions=100,
+        dropout=0.0,
+        attention_dropout=0.0,
+    )
+    config._attn_implementation = attn_implementation
+    return config
+
+
+@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"])
+def test_encoder_reuses_cache_across_streaming_chunks(attn_implementation: str) -> None:
+    """Encode two chunks against the installed Whisper attention and cache classes.
+
+    Unlike the stub-based tests above, this drives the real transformers cache
+    API that ``get_audio_embedding_streaming`` relies on, so it fails whenever a
+    removed cache method or renamed attention kwarg breaks streaming reuse.
+    """
+    torch.manual_seed(0)
+    config = _tiny_encoder_config(attn_implementation)
+    encoder = MiniCPMWhisperEncoder(config).eval()
+
+    first_frames, second_frames = 20, 14
+    first_len = (first_frames - 1) // 2 + 1
+    second_len = (second_frames - 1) // 2 + 1
+    first_mel = torch.randn(1, config.num_mel_bins, first_frames)
+    second_mel = torch.randn(1, config.num_mel_bins, second_frames)
+
+    with torch.no_grad():
+        first = encoder(
+            first_mel,
+            attention_mask=torch.zeros(1, 1, first_len, first_len),
+            past_key_values=None,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+        cache = first.past_key_values
+        assert cache is not None
+        assert _get_audio_cache_length(cache) == first_len
+
+        second = encoder(
+            second_mel,
+            attention_mask=torch.zeros(1, 1, second_len, first_len + second_len),
+            past_key_values=cache,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+        assert _get_audio_cache_length(second.past_key_values) == first_len + second_len
+
+        stateless = encoder(
+            second_mel,
+            attention_mask=torch.zeros(1, 1, second_len, second_len),
+            past_key_values=None,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+
+    assert torch.isfinite(second.last_hidden_state).all()
+    assert not torch.allclose(second.last_hidden_state, stateless.last_hidden_state)

tests/model_executor/models/moss_tts/test_compact_topk.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU regression tests for MOSS-TTS compact top-k sampling."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def test_compact_topk_only_samples_retained_candidates() -> None:
+    logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(7))
+
+    sampled = modeling_moss_tts_local._sample_token(
+        logits,
+        temperature=1.7,
+        top_k=25,
+        top_p=0.8,
+        do_sample=True,
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    retained = logits.topk(25, dim=-1).indices
+    assert torch.all((retained == sampled.unsqueeze(-1)).any(dim=-1))
+
+
+def test_compact_topk_tiny_nucleus_selects_argmax() -> None:
+    logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(17))
+
+    sampled = modeling_moss_tts_local._sample_token(
+        logits,
+        temperature=1.7,
+        top_k=25,
+        top_p=1e-9,
+        do_sample=True,
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(sampled, logits.argmax(dim=-1))

tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
+
+import pytest
+import torch
+from transformers import GPT2Config
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+    MossTTSLocalDepthTransformer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def _test_config() -> GPT2Config:
+    return GPT2Config(
+        n_embd=32,
+        n_head=4,
+        n_inner=64,
+        layer_norm_epsilon=1e-6,
+        rope_base=10000.0,
+    )
+
+
+def test_rope_cache_matches_reference_and_reuses_storage() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    attn = model.h[0].attn
+    device = torch.device("cpu")
+
+    attn.prepare_rope_cache(12, device, torch.float32)
+    cache_ptr = attn._rope_cos_cache.data_ptr()
+    cos, sin = attn._rope_cos_sin(4, device, torch.float32)
+
+    positions = torch.arange(4, dtype=torch.float32)
+    freqs = torch.einsum("s,d->sd", positions, attn.inv_freq.float())
+    expected_cos = freqs.cos().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+    expected_sin = freqs.sin().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+    torch.testing.assert_close(cos, expected_cos)
+    torch.testing.assert_close(sin, expected_sin)
+
+    # Smaller slices retain the fixed allocation used by CUDA Graph.
+    attn.prepare_rope_cache(6, device, torch.float32)
+    assert attn._rope_cos_cache.data_ptr() == cache_ptr

tests/model_executor/models/moss_tts/test_talker_logits_processor.py

@@ -23,6 +23,11 @@ class _RecordingLogitsProcessor:
         return torch.zeros((1, 4))
 
 
+class _EmbeddingBackbone(nn.Module):
+    def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return input_ids.to(dtype=torch.float32).unsqueeze(-1).repeat(1, 4)
+
+
 def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None:
     from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
         MossTTSDelayTalkerForGeneration,
@@ -41,3 +46,74 @@ def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> N
     assert logits.shape == (1, 4)
     assert model.logits_processor.args == (model.text_lm_head, hidden_states)
     assert model.logits_processor.kwargs == {}
+
+
+def test_moss_tts_local_keeps_fixed_shape_rows_and_routes_stop_logits() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    nn.Module.__init__(model)
+    model.audio_pad_token_id = 1024
+    model.audio_assistant_slot_token_id = 2
+    model.im_end_token_id = 3
+    model.text_vocab_size = 8
+    model.n_vq = 4
+    model._batch_state = None
+    model._batch_state_spans = None
+    model._batch_should_continue = None
+
+    valid = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)
+    stopped = torch.full((1, 4), model.audio_pad_token_id, dtype=torch.long)
+    infos = [
+        {"audio_state": {}, "audio_codes": {"current": valid}},
+        {"audio_state": {}, "audio_codes": {"current": stopped}},
+    ]
+    spans = [(0, 3), (3, 4)]
+
+    output = model.make_omni_output(
+        torch.randn(4, 8),
+        model_intermediate_buffer=infos,
+        request_token_spans=spans,
+    )
+
+    audio = output.multimodal_outputs["codes"]["audio"]
+    assert len(audio) == 2
+    torch.testing.assert_close(audio[0], valid)
+    # Keeping this row avoids CUDA boolean-index shape discovery; the
+    # talker2codec raw async-chunk processor drops all-pad rows on CPU.
+    torch.testing.assert_close(audio[1], stopped)
+    assert model._batch_should_continue.tolist() == [True, False]
+
+    logits = model.compute_logits(torch.randn(4, 8))
+
+    assert logits is not None
+    assert logits.argmax(dim=-1).tolist() == [2, 2, 2, 3]
+
+
+def test_moss_tts_local_single_token_prefill_does_not_advance_audio() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    nn.Module.__init__(model)
+    model.model = _EmbeddingBackbone()
+    model.n_vq = 4
+
+    input_ids = torch.tensor([7], dtype=torch.long)
+    _, embeds, update = model.preprocess(
+        input_ids,
+        None,
+        _omni_is_prefill=True,
+        audio_state={"is_stopping": True},
+        ref_offset=5,
+    )
+
+    torch.testing.assert_close(embeds, torch.full((1, 4), 7.0))
+    assert update == {
+        "audio_state": {"is_stopping": False},
+        "ref_offset": 6,
+    }
+    assert "mtp_inputs" not in update

tests/model_tests/diffusion/test_alignment.py

@@ -80,6 +80,7 @@ EXCLUDED_MODELS = [
     "Pi0Pipeline",
     "SoulXSingerPipeline",
     "SoulXSingerSVCPipeline",
+    "SanaWmPipeline",
 ]
 
 

vllm_omni/config/omni_config.py

@@ -11,7 +11,6 @@ later PRs cut consumers over to these classes.
 from __future__ import annotations
 
 import copy
-import math
 from collections.abc import Mapping
 from dataclasses import InitVar, dataclass, field, fields
 from functools import wraps
@@ -601,12 +600,17 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             * self.cfg_parallel_size
         )
         if self.use_hsdp:
-            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
-                raise ValueError(
-                    "HSDP (use_hsdp=True) cannot be used with TP or DP "
-                    f"(tensor_parallel_size={self.tensor_parallel_size}, "
-                    f"data_parallel_size={self.data_parallel_size})"
-                )
+            incompatible = []
+            if self.tensor_parallel_size > 1:
+                incompatible.append("TP")
+            if self.data_parallel_size > 1:
+                incompatible.append("DP")
+            if self.pipeline_parallel_size > 1:
+                incompatible.append("PP")
+            if self.enable_expert_parallel:
+                incompatible.append("EP")
+            if incompatible:
+                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
             if self.hsdp_shard_size == -1:
                 if other_parallel_world_size == 1:
                     raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1")
@@ -688,10 +692,6 @@ class _DiffusionConfigProjection:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
-    dlo_use_host_weight_cache: bool = False
-    dlo_host_weight_cache_dir: str | None = None
-    dlo_host_weight_cache_lock_timeout: float = Field(default=600.0, gt=0)
-    dlo_host_weight_cache_pin_limit_gib: float = Field(default=0.0, ge=0)
     dlo_resident_layers: int = Field(default=0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -754,20 +754,6 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
         from vllm_omni.quantization import build_quant_config
 
-        if not math.isfinite(self.dlo_host_weight_cache_lock_timeout):
-            raise ValueError("dlo_host_weight_cache_lock_timeout must be finite")
-        if not math.isfinite(self.dlo_host_weight_cache_pin_limit_gib):
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib must be finite")
-
-        if self.dlo_host_weight_cache_pin_limit_gib and (
-            not self.enable_distributed_layerwise_offload or self.dlo_use_allgather
-        ):
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
-        if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
-            raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
-        if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True")
-
         if self.tf_model_config is None:
             self.tf_model_config = TransformerConfig()
         elif isinstance(self.tf_model_config, Mapping):

vllm_omni/config/stage_config.py

@@ -382,10 +382,6 @@ class StageDeployConfig:
 
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
-    dlo_use_host_weight_cache: bool | None = None
-    dlo_host_weight_cache_dir: str | None = None
-    dlo_host_weight_cache_lock_timeout: float | None = None
-    dlo_host_weight_cache_pin_limit_gib: float | None = None
     dlo_resident_layers: int | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None

vllm_omni/diffusion/data.py

@@ -154,8 +154,8 @@ class DiffusionParallelConfig:
     pipeline_parallel_size: int = 1
     """Number of pipeline parallel stages."""
 
-    data_parallel_size: int = 1
-    """Number of data parallel groups."""
+    data_parallel_size: int | None = None
+    """Number of data parallel groups. Inferred from WORLD when omitted."""
 
     tensor_parallel_size: int = 1
     """Number of tensor parallel groups."""
@@ -243,7 +243,7 @@ class DiffusionParallelConfig:
     def _validate_parallel_config(self) -> Self:
         """Validates the config relationships among the parallel strategies."""
         assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0"
-        assert self.data_parallel_size > 0, "Data parallel size must be > 0"
+        assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0"
         assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0"
         assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
         assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
@@ -283,10 +283,12 @@ class DiffusionParallelConfig:
                 self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
             )
 
-        # Calculate world_size from other parallelism dimensions
+        # Until the runtime WORLD size is known, an omitted DP dimension means
+        # one replica. OmniDiffusionConfig resolves it against num_gpus below.
+        data_parallel_size = self.data_parallel_size or 1
         other_parallel_world_size = (
             self.pipeline_parallel_size
-            * self.data_parallel_size
+            * data_parallel_size
             * self.tensor_parallel_size
             * self.sequence_parallel_size
             * self.cfg_parallel_size
@@ -297,13 +299,17 @@ class DiffusionParallelConfig:
         # 1. Standalone: when other parallelism is all 1, HSDP determines world_size
         # 2. Combined: HSDP overlays on top of other parallelism
         if self.use_hsdp:
-            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
-                raise ValueError(
-                    "HSDP (use_hsdp=True) cannot be used with TP or DP "
-                    f"(tensor_parallel_size={self.tensor_parallel_size}, "
-                    f"data_parallel_size={self.data_parallel_size}). "
-                    "Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP."
-                )
+            incompatible = []
+            if self.tensor_parallel_size > 1:
+                incompatible.append("TP")
+            if data_parallel_size > 1:
+                incompatible.append("DP")
+            if self.pipeline_parallel_size > 1:
+                incompatible.append("PP")
+            if self.enable_expert_parallel:
+                incompatible.append("EP")
+            if incompatible:
+                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
             if self.hsdp_shard_size == -1:
                 # Auto-calculate: use other_parallel_world_size as shard_size
                 if self.hsdp_replicate_size <= 0:
@@ -338,6 +344,42 @@ class DiffusionParallelConfig:
         else:
             self.world_size = other_parallel_world_size
 
+    def resolve_data_parallel_size(self, world_size: int) -> int:
+        """Resolve or validate diffusion DP against the actual WORLD size."""
+        if world_size <= 0:
+            raise ValueError(f"WORLD size must be > 0, but got {world_size}")
+
+        if self.use_hsdp:
+            if self.data_parallel_size not in (None, 1):
+                raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1")
+            expected_world_size = self.hsdp_replicate_size * self.hsdp_shard_size
+            if world_size != expected_world_size:
+                raise ValueError(
+                    f"WORLD size ({world_size}) must equal HSDP size "
+                    f"({self.hsdp_replicate_size} x {self.hsdp_shard_size} = {expected_world_size})"
+                )
+            self.data_parallel_size = 1
+            self.world_size = world_size
+            return 1
+
+        non_dp_size = (
+            self.pipeline_parallel_size
+            * self.tensor_parallel_size
+            * self.sequence_parallel_size
+            * self.cfg_parallel_size
+        )
+        if world_size % non_dp_size != 0:
+            raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})")
+        inferred_data_parallel_size = world_size // non_dp_size
+        if self.data_parallel_size is not None and self.data_parallel_size != inferred_data_parallel_size:
+            raise ValueError(
+                f"data_parallel_size ({self.data_parallel_size}) does not match WORLD-derived value "
+                f"({inferred_data_parallel_size})"
+            )
+        self.data_parallel_size = inferred_data_parallel_size
+        self.world_size = world_size
+        return inferred_data_parallel_size
+
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig":
         """
@@ -708,16 +750,6 @@ class OmniDiffusionConfig:
     # This avoids AllGather synchronization, while host memory follows the
     # loader's existing rank-local layout instead of adding a second DP shard.
     dlo_use_allgather: bool = True
-    # Opt into publishing/joining the transform-complete final-layout host
-    # weight cache instead of direct-checkpoint mmap.
-    dlo_use_host_weight_cache: bool = False
-    # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-host-weights.
-    dlo_host_weight_cache_dir: str | None = None
-    # Maximum wait for the per-layout POSIX writer lock.
-    dlo_host_weight_cache_lock_timeout: float = 600.0
-    # Optional maximum file-backed memory a worker may register with CUDA.
-    # Zero applies no additional ceiling; pin_cpu_memory controls registration.
-    dlo_host_weight_cache_pin_limit_gib: float = 0.0
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
 
@@ -946,18 +978,6 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
-        if not math.isfinite(self.dlo_host_weight_cache_lock_timeout) or self.dlo_host_weight_cache_lock_timeout <= 0:
-            raise ValueError("dlo_host_weight_cache_lock_timeout must be positive")
-        if not math.isfinite(self.dlo_host_weight_cache_pin_limit_gib) or self.dlo_host_weight_cache_pin_limit_gib < 0:
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib must be non-negative")
-        if self.dlo_host_weight_cache_pin_limit_gib and (
-            not self.enable_distributed_layerwise_offload or self.dlo_use_allgather
-        ):
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
-        if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
-            raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
-        if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
-            raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True")
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "
@@ -1021,10 +1041,7 @@ class OmniDiffusionConfig:
             else:
                 self.num_gpus = 1
 
-        if self.num_gpus < self.parallel_config.world_size:
-            raise ValueError(
-                f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})"
-            )
+        self.parallel_config.resolve_data_parallel_size(self.num_gpus)
 
         if self.diffusion_compile_granularity == "full":
             incompatible_features = []

vllm_omni/diffusion/distributed/__init__.py

@@ -6,11 +6,6 @@ from vllm_omni.diffusion.distributed.hsdp import (
     HSDPInferenceConfig,
     apply_hsdp_to_model,
 )
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_fs_group,
-    get_fully_shard_rank,
-    get_fully_shard_world_size,
-)
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelConfig,
     SequenceParallelInput,
@@ -46,8 +41,4 @@ __all__ = [
     "get_sharding_validator",
     "HSDPInferenceConfig",
     "apply_hsdp_to_model",
-    # FS (Fully Shard) utilities
-    "get_fs_group",
-    "get_fully_shard_rank",
-    "get_fully_shard_world_size",
 ]

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -9,8 +9,6 @@ import torch
 import torch.distributed as dist
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
-
 logger = init_logger(__name__)
 
 
@@ -45,9 +43,10 @@ class DistributedVaeExecutor:
     """
 
     def __init__(self):
-        self.group = get_dit_group()
-        self.world_size = dist.get_world_size(self.group)
-        self.rank = dist.get_rank(self.group)
+        # VAE tiles are scheduled across the complete worker WORLD.
+        self.group = None
+        self.world_size = dist.get_world_size()
+        self.rank = dist.get_rank()
         self.parallel_size = 1
         self.parallel_mode = "tile"
 
@@ -189,7 +188,7 @@ class DistributedVaeMixin:
         if self.distributed_executor.parallel_size > pp_size:
             logger.warning(
                 f"vae_patch_parallel_size={self.distributed_executor.parallel_size} "
-                f"is greater than dit_group={world_size};"
-                f" using dit_group size={world_size}"
+                f"is greater than WORLD={world_size};"
+                f" using WORLD size={world_size}"
             )
         return True

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -991,6 +991,18 @@ class PipelineGroupCoordinator(GroupCoordinator):
     def _pipeline_isend_skip(self, tensor: torch.tensor):
         return torch.distributed.isend(tensor, dst=self.skip_rank, group=self.skip_device_group)
 
+    def destroy(self):
+        groups = [*self.device_groups, *self.cpu_groups, self.skip_device_group]
+        seen = {id(self.device_group), id(self.cpu_group)}
+        for group in groups:
+            if group is not None and id(group) not in seen:
+                torch.distributed.destroy_process_group(group)
+                seen.add(id(group))
+        self.device_groups = []
+        self.cpu_groups = []
+        self.skip_device_group = None
+        super().destroy()
+
 
 class SequenceParallelGroupCoordinator(GroupCoordinator):
     def __init__(
@@ -1031,3 +1043,15 @@ class SequenceParallelGroupCoordinator(GroupCoordinator):
         self.ring_rank = torch.distributed.get_rank(self.ring_group)
         self.allgather_world_size = torch.distributed.get_world_size(self.allgather_group)
         self.allgather_rank = torch.distributed.get_rank(self.allgather_group)
+
+    def destroy(self):
+        groups = [self.ulysses_group, self.ring_group, self.allgather_group]
+        seen = {id(self.device_group), id(self.cpu_group)}
+        for group in groups:
+            if group is not None and id(group) not in seen:
+                torch.distributed.destroy_process_group(group)
+                seen.add(id(group))
+        self.ulysses_group = None
+        self.ring_group = None
+        self.allgather_group = None
+        super().destroy()

vllm_omni/diffusion/distributed/hsdp.py

@@ -15,12 +15,7 @@ from torch.distributed.fsdp import (
 )
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_fs_group,
-    get_fully_shard_rank,
-    get_fully_shard_world_size,
-    get_world_group,
-)
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
@@ -54,36 +49,36 @@ class HSDPInferenceConfig:
 def _create_hsdp_mesh(
     device_type: str,
     replicate_size: int,
-    shard_pg: torch.distributed.ProcessGroup,
+    shard_size: int,
 ) -> DeviceMesh:
-    """Create a 2D DeviceMesh for HSDP using an existing ProcessGroup for the shard dimension.
+    """Create the FSDP2 DeviceMesh; it owns the HSDP process groups.
 
     Args:
         device_type: The device type (e.g., "cuda", "npu")
         replicate_size: Number of replica groups
-        shard_pg: The ProcessGroup for the shard dimension (from FS GroupCoordinator)
+        shard_size: Number of ranks in each FSDP shard group
 
     Returns:
         A 2D DeviceMesh with dimensions ("replicate", "shard")
     """
-    shard_size = torch.distributed.get_world_size(shard_pg)
     world_size = replicate_size * shard_size
+    actual_world_size = torch.distributed.get_world_size()
+    if world_size != actual_world_size:
+        raise ValueError(
+            f"HSDP mesh dimensions ({replicate_size} x {shard_size} = {world_size}) "
+            f"must equal WORLD size ({actual_world_size})"
+        )
 
     # Build 2D mesh tensor: shape (replicate_size, shard_size)
     # Ranks are arranged so that each row is a shard group
     mesh_tensor = torch.arange(world_size).reshape(replicate_size, shard_size)
 
-    # Create DeviceMesh with the shard ProcessGroup
-    # For the shard dimension, we reuse the existing FS ProcessGroup
     device_mesh = init_device_mesh(
         device_type,
         mesh_shape=(replicate_size, shard_size),
         mesh_dim_names=("replicate", "shard"),
     )
 
-    # Note: init_device_mesh creates new ProcessGroups internally.
-    # For consistency, we verify the mesh structure matches our FS group.
-    # In a future optimization, we could pass the existing ProcessGroups directly.
     logger.debug(
         "Created HSDP mesh: replicate_size=%d, shard_size=%d, mesh=%s",
         replicate_size,
@@ -119,34 +114,24 @@ def apply_hsdp_to_model(
     if not hsdp_config.enabled:
         raise ValueError("HSDP is not enabled in config")
 
-    # Use GroupCoordinator for distributed info
     world_group = get_world_group()
-    fs_group = get_fs_group()
-
     world_size = world_group.world_size
     rank = world_group.rank_in_group
-    fs_world_size = get_fully_shard_world_size()
-    fs_rank = get_fully_shard_rank()
 
     hsdp_replicate_size = hsdp_config.hsdp_replicate_size
     hsdp_shard_size = hsdp_config.hsdp_shard_size
 
-    # Validate that the FS group matches the HSDP shard size
-    if fs_world_size != hsdp_shard_size:
+    if hsdp_replicate_size * hsdp_shard_size != world_size:
         raise ValueError(
-            f"FS group world_size ({fs_world_size}) does not match "
-            f"HSDP shard_size ({hsdp_shard_size}). "
-            "Ensure fully_shard_degree is set correctly in initialize_model_parallel."
+            f"HSDP mesh dimensions ({hsdp_replicate_size} x {hsdp_shard_size}) must equal WORLD size ({world_size})"
         )
 
     logger.info(
-        "HSDP Inference: replicate_size=%d, shard_size=%d, world_size=%d, rank=%d, fs_world_size=%d, fs_rank=%d",
+        "HSDP Inference (FSDP2): replicate_size=%d, shard_size=%d, world_size=%d, rank=%d",
         hsdp_replicate_size,
         hsdp_shard_size,
         world_size,
         rank,
-        fs_world_size,
-        fs_rank,
     )
 
     # When the model contains FP8 parameters (online quantization), let FSDP
@@ -163,14 +148,14 @@ def apply_hsdp_to_model(
 
     device_type = current_omni_platform.device_type
 
-    # Create 2D DeviceMesh for HSDP using the FS group's ProcessGroup for shard dimension
+    # DeviceMesh is the single source of truth for FSDP2 shard/replicate groups
     # The mesh shape is (replicate, shard) where:
     # - replicate: groups of ranks that hold the same shard (for gradient all-reduce in training)
     # - shard: groups of ranks that each hold different shards (for parameter all-gather)
     device_mesh = _create_hsdp_mesh(
         device_type=device_type,
         replicate_size=hsdp_replicate_size,
-        shard_pg=fs_group.device_group,
+        shard_size=hsdp_shard_size,
     )
 
     hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)

vllm_omni/diffusion/distributed/parallel_state.py

@@ -29,11 +29,11 @@ If you only need to use the distributed environment without model parallelism,
 """
 
 import inspect
+from math import prod
 
 import torch
 import torch.distributed
 import vllm.distributed.parallel_state as vllm_parallel_state
-from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
@@ -59,8 +59,6 @@ _SP: SequenceParallelGroupCoordinator | None = None
 _PP: PipelineGroupCoordinator | None = None
 _CFG: GroupCoordinator | None = None
 _DP: GroupCoordinator | None = None
-_FS: GroupCoordinator | None = None  # Fully Sharded (HSDP shard dimension)
-_DIT: GroupCoordinator | None = None
 
 # Rank-layout metadata for expert parallelism. This is not a process group;
 # it is reused by platform-specific runtimes that must build companion groups
@@ -184,7 +182,6 @@ class RankGenerator:
         pp: int,
         cfg: int,
         dp: int,
-        fs: int = 1,
         order: str = "tp-sp-pp-cfg-dp",
         rank_offset: int = 0,
     ) -> None:
@@ -193,9 +190,7 @@ class RankGenerator:
         self.pp = pp
         self.cfg = cfg
         self.dp = dp
-        self.fs = fs
         self.rank_offset = rank_offset
-        self.world_size = tp * sp * pp * cfg * dp
 
         self.name_to_size = {
             "tp": self.tp,
@@ -203,15 +198,10 @@ class RankGenerator:
             "pp": self.pp,
             "cfg": self.cfg,
             "dp": self.dp,
-            "fs": self.fs,
         }
         order = order.lower()
 
         for name in self.name_to_size.keys():
-            # Skip 'fs' validation - it's handled separately with independent_ranks=True
-            # and doesn't participate in the main orthogonal rank generation
-            if name == "fs":
-                continue
             if name not in order and self.name_to_size[name] != 1:
                 raise RuntimeError(
                     f"The size of ({name}) is ({self.name_to_size[name]}), "
@@ -234,7 +224,7 @@ class RankGenerator:
             mask[ordered_token.index(t)] = True
         return mask
 
-    def get_ranks(self, token, independent_ranks: bool = False):
+    def get_ranks(self, token):
         """Get rank group by input token.
 
         Arguments:
@@ -243,27 +233,11 @@ class RankGenerator:
                 to obtain multiple parallel types, we can use a hyphen
                 '-' to separate them. For example, if we want to obtain
                 the TP_DP group, the token should be 'tp-dp'.
-            independent_ranks (bool):
-                If True, generate independent rank groups that divide the world
-                into groups of the specified size. Used for FS (fully shard) groups
-                which operate independently from the main parallelism hierarchy.
         """
-        if independent_ranks and token == "fs":
-            # FS groups divide world into groups of size fs
-            # e.g., world_size=8, fs=4 -> [[0,1,2,3], [4,5,6,7]]
-            ranks = []
-            num_groups = self.world_size // self.fs
-            for i in range(num_groups):
-                group = list(range(i * self.fs + self.rank_offset, (i + 1) * self.fs + self.rank_offset))
-                ranks.append(group)
-            return ranks
-
         mask = self.get_mask(self.order, token)
-        ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
-        if self.rank_offset > 0:
-            for rank_group in ranks:
-                for i in range(len(rank_group)):
-                    rank_group[i] += self.rank_offset
+        ranks = generate_masked_orthogonal_rank_groups(prod(self.ordered_size), self.ordered_size, mask)
+        if self.rank_offset:
+            return [[rank + self.rank_offset for rank in rank_group] for rank_group in ranks]
         return ranks
 
 
@@ -376,22 +350,6 @@ def get_data_parallel_rank():
     return get_dp_group().rank_in_group
 
 
-# FS (Fully Shard / HSDP shard dimension)
-def get_fs_group() -> GroupCoordinator:
-    assert _FS is not None, "fully shard group is not initialized"
-    return _FS
-
-
-def get_fully_shard_world_size():
-    """Return world size for the fully shard group."""
-    return get_fs_group().world_size
-
-
-def get_fully_shard_rank():
-    """Return my rank for the fully shard group."""
-    return get_fs_group().rank_in_group
-
-
 def is_dp_last_group():
     """Return True if in the last data parallel group, False otherwise."""
     return (
@@ -401,17 +359,6 @@ def is_dp_last_group():
     )
 
 
-def get_dit_world_size():
-    """Return world size for the DiT model."""
-    return (
-        get_data_parallel_world_size()
-        * get_classifier_free_guidance_world_size()
-        * get_sequence_parallel_world_size()
-        * get_pipeline_parallel_world_size()
-        * get_tensor_model_parallel_world_size()
-    )
-
-
 # * SET
 
 
@@ -498,7 +445,6 @@ def init_model_parallel_group(
         "expert",
         "sequence",
         "classifier_free_guidance",
-        "fully_shard",
     ], f"parallel_mode {parallel_mode} is not supported"
     if parallel_mode == "pipeline":
         return PipelineGroupCoordinator(
@@ -546,19 +492,6 @@ def init_vllm_model_parallel_group(
     )
 
 
-def init_dit_group(
-    dit_parallel_size: int,
-    backend: str,
-):
-    global _DIT
-    _DIT = torch.distributed.new_group(ranks=list(range(dit_parallel_size)), backend=backend)
-
-
-def get_dit_group():
-    assert _DIT is not None, "DIT group is not initialized"
-    return _DIT
-
-
 # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
 def set_seq_parallel_pg(
     sp_ulysses_degree: int,
@@ -737,8 +670,8 @@ def set_seq_parallel_pg(
     return ulysses_pg, ring_pg, allgather_pg
 
 
-def initialize_model_parallel(
-    data_parallel_size: int = 1,
+def _initialize_model_parallel(
+    data_parallel_size: int | None = None,
     cfg_parallel_size: int = 1,
     sequence_parallel_size: int | None = None,
     ulysses_degree: int = 1,
@@ -746,9 +679,8 @@ def initialize_model_parallel(
     allgather_degree: int = 1,
     tensor_parallel_size: int = 1,
     pipeline_parallel_size: int = 1,
-    fully_shard_degree: int = 1,
-    hsdp_replicate_size: int = 1,
     enable_expert_parallel: bool = False,
+    use_hsdp: bool = False,
     backend: str | None = None,
 ) -> None:
     if backend is None:
@@ -768,7 +700,6 @@ def initialize_model_parallel(
             (causal=False only). Mutually exclusive with ulysses/ring in v1.
         tensor_parallel_size: number of GPUs used for tensor parallelism.
         pipeline_parallel_size: number of GPUs used for pipeline parallelism.
-        fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension).
         backend: distributed backend of pytorch collective comm.
 
     Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
@@ -798,16 +729,20 @@ def initialize_model_parallel(
     ranks 8 to 15 belong to the second box.
     """
     # Get world size and rank. Ensure some consistencies.
-    assert torch.distributed.is_initialized()
+    if not torch.distributed.is_initialized():
+        raise RuntimeError("torch.distributed must be initialized before model parallel initialization")
+    if model_parallel_is_initialized():
+        raise RuntimeError("model parallel groups are already initialized")
     world_size: int = torch.distributed.get_world_size()
     backend = backend or torch.distributed.get_backend(get_world_group().device_group)
 
     if allgather_degree > 1:
-        assert ulysses_degree == 1 and ring_degree == 1, (
-            "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
-            f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
-            f"allgather_degree={allgather_degree}."
-        )
+        if ulysses_degree != 1 or ring_degree != 1:
+            raise ValueError(
+                "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
+                f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
+                f"allgather_degree={allgather_degree}."
+            )
 
     expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree
     if sequence_parallel_size is None:
@@ -820,42 +755,38 @@ def initialize_model_parallel(
             f"but got {sequence_parallel_size}"
         )
 
-    dit_parallel_size = (
-        data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
-    )
-
-    # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
-    is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
-
-    # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) as dit_parallel_size
-    # This ensures orthogonal rank generation works correctly for all HSDP workers
-    if is_standalone_hsdp:
-        dit_parallel_size = fully_shard_degree * hsdp_replicate_size
-
-    if world_size < dit_parallel_size:
-        raise RuntimeError(
-            f"world_size ({world_size}) is less than "
-            f"tensor_parallel_size ({tensor_parallel_size}) x "
-            f"pipeline_parallel_size ({pipeline_parallel_size}) x"
-            f"sequence_parallel_size ({sequence_parallel_size}) x"
-            f"cfg_parallel_size "
-            f"({cfg_parallel_size}) x"
-            f"data_parallel_size ({data_parallel_size})"
-        )
-
-    # For standalone HSDP, use (fully_shard_degree * hsdp_replicate_size) as data_parallel_size
-    # so that RankGenerator.world_size matches the actual number of workers
-    effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
+    non_dp_size = cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
+    if world_size % non_dp_size != 0:
+        raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})")
+    if use_hsdp:
+        if data_parallel_size not in (None, 1):
+            raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1")
+        if non_dp_size not in (1, world_size):
+            raise ValueError(f"HSDP non-DP parallel size must be 1 or WORLD size ({world_size}), but got {non_dp_size}")
+        data_parallel_size = 1
+    else:
+        inferred_data_parallel_size = world_size // non_dp_size
+        if data_parallel_size is not None and data_parallel_size != inferred_data_parallel_size:
+            raise ValueError(
+                f"data_parallel_size ({data_parallel_size}) does not match WORLD-derived value "
+                f"({inferred_data_parallel_size})"
+            )
+        data_parallel_size = inferred_data_parallel_size
 
     rank_generator: RankGenerator = RankGenerator(
         tensor_parallel_size,
         sequence_parallel_size,
         pipeline_parallel_size,
         cfg_parallel_size,
-        effective_dp_size,
-        fs=fully_shard_degree,
+        data_parallel_size,
         order="tp-sp-pp-cfg-dp",
     )
+
+    def get_rank_groups(token: str) -> list[list[int]]:
+        if use_hsdp and non_dp_size == 1:
+            return [[rank] for rank in range(world_size)]
+        return rank_generator.get_ranks(token)
+
     use_moe_parallel_mapping = False
     if enable_expert_parallel:
         od_config = get_forward_context().omni_diffusion_config
@@ -863,11 +794,11 @@ def initialize_model_parallel(
         if not use_moe_parallel_mapping:
             raise RuntimeError("Expert parallelism enabled for a non-MoE model")
 
-    sp_group_ranks = rank_generator.get_ranks("sp")
+    sp_group_ranks = get_rank_groups("sp")
     global _DP
     assert _DP is None, "data parallel group is already initialized"
     _DP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("dp"),
+        group_ranks=get_rank_groups("dp"),
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="data",
@@ -877,7 +808,7 @@ def initialize_model_parallel(
     global _CFG
     assert _CFG is None, "classifier_free_guidance group is already initialized"
     _CFG = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("cfg"),
+        group_ranks=get_rank_groups("cfg"),
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="classifier_free_guidance",
@@ -885,7 +816,7 @@ def initialize_model_parallel(
     global _PP
     assert _PP is None, "pipeline model parallel group is already initialized"
     _PP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("pp"),
+        group_ranks=get_rank_groups("pp"),
         local_rank=get_world_group().local_rank,
         backend=backend,
         parallel_mode="pipeline",
@@ -899,7 +830,7 @@ def initialize_model_parallel(
         sp_ring_degree=ring_degree,
         sp_allgather_degree=allgather_degree,
         rank=get_world_group().rank_in_group,
-        world_size=dit_parallel_size,
+        world_size=world_size,
         sp_group_ranks=sp_group_ranks,
     )
     _SP = init_model_parallel_group(
@@ -925,7 +856,7 @@ def initialize_model_parallel(
         )
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
-    tp_group_ranks = rank_generator.get_ranks("tp")
+    tp_group_ranks = get_rank_groups("tp")
     if use_moe_parallel_mapping:
         vllm_parallel_state._TP = init_vllm_model_parallel_group(
             group_ranks=tp_group_ranks,
@@ -944,22 +875,14 @@ def initialize_model_parallel(
         # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
         # only when constructing the vLLM EP layout for expert-parallel paths.
         vllm_parallel_state._DP = init_vllm_model_parallel_group(
-            group_ranks=rank_generator.get_ranks("cfg-dp"),
+            group_ranks=get_rank_groups("cfg-dp"),
             local_rank=get_world_group().local_rank,
             backend=backend,
             group_name="dp",
         )
 
-    global _FS, _EXPERT_PARALLEL_GROUP_RANKS
-    assert _FS is None, "fully shard group is already initialized"
-    _FS = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("fs", independent_ranks=True),
-        local_rank=get_world_group().local_rank,
-        backend=backend,
-        parallel_mode="fully_shard",
-    )
-
-    _EXPERT_PARALLEL_GROUP_RANKS = rank_generator.get_ranks("tp-sp-cfg-dp")
+    global _EXPERT_PARALLEL_GROUP_RANKS
+    _EXPERT_PARALLEL_GROUP_RANKS = get_rank_groups("tp-sp-cfg-dp")
     if use_moe_parallel_mapping:
         vllm_parallel_state._EP = init_vllm_model_parallel_group(
             group_ranks=_EXPERT_PARALLEL_GROUP_RANKS,
@@ -969,12 +892,63 @@ def initialize_model_parallel(
             use_all2all=True,
         )
 
-    init_dit_group(dit_parallel_size, backend)
+
+def initialize_model_parallel(
+    data_parallel_size: int | None = None,
+    cfg_parallel_size: int = 1,
+    sequence_parallel_size: int | None = None,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    allgather_degree: int = 1,
+    tensor_parallel_size: int = 1,
+    pipeline_parallel_size: int = 1,
+    enable_expert_parallel: bool = False,
+    use_hsdp: bool = False,
+    backend: str | None = None,
+) -> None:
+    """Atomically initialize diffusion parallel groups.
+
+    Configuration is validated before group creation. Any failure tears down
+    groups created by this attempt, including vLLM-bound coordinators.
+    """
+    existing_groups = {
+        "dp": _DP,
+        "cfg": _CFG,
+        "sp": _SP,
+        "pp": _PP,
+        "vllm_tp": vllm_parallel_state._TP,
+        "vllm_dp": vllm_parallel_state._DP,
+        "vllm_pp": vllm_parallel_state._PP,
+        "vllm_ep": vllm_parallel_state._EP,
+        "vllm_pcp": vllm_parallel_state._PCP,
+    }
+    initialized_names = [name for name, group in existing_groups.items() if group is not None]
+    if initialized_names:
+        raise RuntimeError(
+            "model parallel state must be empty before initialization; found " + ", ".join(initialized_names)
+        )
+    try:
+        _initialize_model_parallel(
+            data_parallel_size=data_parallel_size,
+            cfg_parallel_size=cfg_parallel_size,
+            sequence_parallel_size=sequence_parallel_size,
+            ulysses_degree=ulysses_degree,
+            ring_degree=ring_degree,
+            allgather_degree=allgather_degree,
+            tensor_parallel_size=tensor_parallel_size,
+            pipeline_parallel_size=pipeline_parallel_size,
+            enable_expert_parallel=enable_expert_parallel,
+            use_hsdp=use_hsdp,
+            backend=backend,
+        )
+    except BaseException:
+        destroy_model_parallel()
+        raise
 
 
 def destroy_model_parallel():
     """Set the groups to none and destroy them."""
-    global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS
+    global _DP, _CFG, _SP, _PP, _EXPERT_PARALLEL_GROUP_RANKS
 
     if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP:
         vllm_parallel_state._DP.destroy()
@@ -1005,14 +979,14 @@ def destroy_model_parallel():
     vllm_parallel_state._EP = None
     _EXPERT_PARALLEL_GROUP_RANKS = None
 
+    if vllm_parallel_state._PP and vllm_parallel_state._PP is not _PP:
+        vllm_parallel_state._PP.destroy()
+    vllm_parallel_state._PP = None
+
     if _PP:
         _PP.destroy()
     _PP = None
 
-    if _FS:
-        _FS.destroy()
-    _FS = None
-
 
 def destroy_distributed_environment():
     global _WORLD

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -14,7 +14,6 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
@@ -441,13 +440,10 @@ class DiffusersPipelineLoader:
                 _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
                 _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
                 _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size
-                _select_host_weight_cache = (
-                    _dist_offload and not _use_ag and bool(getattr(self.od_config, "dlo_use_host_weight_cache", False))
-                )
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
-                if _dist_offload and not _select_host_weight_cache:
+                if _dist_offload:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -459,16 +455,6 @@ class DiffusersPipelineLoader:
                         online_quantization=_has_online_quant,
                     )
                     self.host_weight_plan = plan_result.plan
-                elif _select_host_weight_cache:
-                    # Direct checkpoint bindings can carry deferred transforms
-                    # (for example grouped-QKV reordering), so registering the
-                    # raw checkpoint mapping would not make every recurrent H2D
-                    # source directly copyable. Materialize once, apply every
-                    # post-load mutation, then publish/register the final layout.
-                    logger.info(
-                        "DLO host weight cache selected; using the ordinary loader "
-                        "once to publish a transform-complete final mmap layout"
-                    )
 
                 _skip_load = self.host_weight_plan is not None
 
@@ -510,8 +496,7 @@ class DiffusersPipelineLoader:
                         )
                     if _dist_offload and plan_result is not None:
                         logger.info(
-                            "DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s",
-                            plan_result.fallback_code or "unknown",
+                            "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
                             plan_result.fallback_reason,
                         )
                     logger.debug("Loading weights on %s ...", load_device)
@@ -536,110 +521,7 @@ class DiffusersPipelineLoader:
                 logger.info("Quantization complete, offloaded model back to CPU")
 
         self._apply_skip_softmax_calibration(model)
-        model.eval()
-        self._maybe_build_host_weight_cache(model, load_format)
-        return model
-
-    def _maybe_build_host_weight_cache(self, model: nn.Module, load_format: str) -> None:
-        """Publish final ordinary-loader DiT tensors for no-AllGather DLO."""
-        if self.host_weight_plan is not None:
-            return
-        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
-            return
-        if getattr(self.od_config, "dlo_use_allgather", True):
-            return
-        if not getattr(self.od_config, "dlo_use_host_weight_cache", False):
-            logger.info(
-                "DLO host weight cache not selected; retaining ordinary host tensors after checkpoint mmap fallback"
-            )
-            return
-
-        parallel_config = self.parallel_config
-        tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1))
-        tp_rank = 0
-        if tp_size > 1:
-            try:
-                from vllm.distributed import get_tensor_model_parallel_rank
-
-                tp_rank = int(get_tensor_model_parallel_rank())
-            except Exception as exc:
-                logger.warning(
-                    "DLO host weight cache unavailable [parallel_rank_unavailable]: "
-                    "cannot resolve TP rank for TP=%d: %s",
-                    tp_size,
-                    exc,
-                )
-                return
-
-        modules = ModuleDiscovery.discover(model)
-        from vllm_omni.diffusion.model_loader.host_weight_cache import (
-            DEFAULT_LOCK_TIMEOUT_SECONDS,
-            DEFAULT_SHARD_SIZE_BYTES,
-            build_host_weight_cache_plan,
-        )
-
-        loader_inputs = {
-            "model_class_name": getattr(self.od_config, "model_class_name", None),
-            "model_config": getattr(self.od_config, "model_config", None),
-            "diffusers_load_kwargs": getattr(self.od_config, "diffusers_load_kwargs", None),
-            "custom_pipeline_args": getattr(self.od_config, "custom_pipeline_args", None),
-            "lora_path": getattr(self.od_config, "lora_path", None),
-            "lora_scale": getattr(self.od_config, "lora_scale", None),
-            "lora_backend": getattr(self.od_config, "lora_backend", None),
-        }
-        sp_guard = {
-            "sequence_parallel_size": int(getattr(parallel_config, "sequence_parallel_size", 1)),
-            "ulysses_degree": int(getattr(parallel_config, "ulysses_degree", 1)),
-            "ring_degree": int(getattr(parallel_config, "ring_degree", 1)),
-            "allgather_degree": int(getattr(parallel_config, "allgather_degree", 1)),
-            "ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"),
-        }
-        model_identity = getattr(self.od_config, "model", None)
-        runtime_quantization = self.quant_config
-        if runtime_quantization is None and any(
-            (quant_method := getattr(module, "quant_method", None)) is not None
-            and not isinstance(quant_method, UnquantizedLinearMethod)
-            for module in model.modules()
-        ):
-            runtime_quantization = "module_quant_method"
-        result = build_host_weight_cache_plan(
-            model,
-            dit_modules=tuple(zip(modules.dit_names, modules.dits)),
-            loader_type=type(self),
-            cache_root=getattr(self.od_config, "dlo_host_weight_cache_dir", None),
-            lock_timeout_seconds=float(
-                getattr(self.od_config, "dlo_host_weight_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS)
-            ),
-            max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES,
-            model_identity=str(model_identity) if model_identity is not None else None,
-            revision=getattr(self.od_config, "revision", None),
-            runtime_dtype=getattr(self.od_config, "dtype", None),
-            load_format=load_format,
-            loader_inputs=loader_inputs,
-            tensor_parallel_size=tp_size,
-            tensor_parallel_rank=tp_rank,
-            sequence_parallel_guard=sp_guard,
-            use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)),
-            enable_expert_parallel=bool(getattr(parallel_config, "enable_expert_parallel", False)),
-            quantization_config=runtime_quantization,
-            cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)),
-            pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)),
-        )
-        if result.plan is not None:
-            self.host_weight_plan = result.plan
-            logger.info(
-                "DLO host weight cache selected (layout=%s, TP=%d/%d, SP=%d)",
-                result.plan.runtime_layout_key[:12] if result.plan.runtime_layout_key else "unknown",
-                tp_rank,
-                tp_size,
-                sp_guard["sequence_parallel_size"],
-            )
-        else:
-            logger.warning(
-                "DLO host weight cache unavailable [%s]; retaining ordinary host tensors: %s",
-                result.fallback_code or "unknown",
-                result.fallback_reason or "unspecified reason",
-            )
+        return model.eval()
 
     @staticmethod
     def _request_offload_after_quant(model: nn.Module) -> int:

vllm_omni/diffusion/model_loader/host_weight_cache.py

@@ -1,574 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Build node-local host weight cache plans for diffusion DLO.
-
-``build_host_weight_cache_plan`` is the module's single entry point. It
-validates final ordinary-loader DiT tensors, joins or atomically publishes an
-immutable content-addressed entry, validates the mmap result, and returns a
-fail-closed host-weight plan. Reusable identity, hashing, locking, and file
-primitives live in :mod:`host_weight_cache_utils`.
-"""
-
-from __future__ import annotations
-
-import contextlib
-import hashlib
-import json
-import math
-import os
-import shutil
-import time
-import uuid
-from collections.abc import Mapping, Sequence
-from itertools import chain
-from pathlib import Path
-from typing import Any
-
-import torch
-from safetensors import safe_open
-from safetensors.torch import save_file
-from torch import nn
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.model_loader.host_weight_cache_utils import (
-    HostWeightCacheError,
-    IdentityNormalizationError,
-    RuntimeTensor,
-    canonical_json,
-    canonicalize_existing_local_path,
-    exclusive_lock,
-    file_digest,
-    fsync_directory,
-    fsync_file,
-    implementation_fingerprint,
-    normalize_identity,
-    remove_stale_temps,
-    runtime_content_digest,
-    split_shards,
-    tensor_metadata,
-)
-from vllm_omni.diffusion.model_loader.host_weight_plan import (
-    HostWeightPlan,
-    HostWeightPlanResult,
-    TensorBinding,
-)
-
-logger = init_logger(__name__)
-
-# Internal cache-format and I/O defaults, not environment variables. Public
-# configuration may override the root and lock timeout; chunk sizes stay private.
-# Bump the schema to isolate incompatible manifest layouts.
-HOST_WEIGHT_CACHE_SCHEMA_VERSION = 1
-# Keep the default shared across local stage replicas.
-DEFAULT_HOST_WEIGHT_CACHE_ROOT = "~/.cache/vllm-omni/dlo-host-weights"
-# Bound waits for another publisher instead of hanging initialization forever.
-DEFAULT_LOCK_TIMEOUT_SECONDS = 600.0
-# Avoid one model-sized output file and bound hashing/read working sets.
-DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3
-
-_SUPPORTED_DTYPES = {
-    torch.bool,
-    torch.uint8,
-    torch.int8,
-    torch.int16,
-    torch.int32,
-    torch.int64,
-    torch.float16,
-    torch.bfloat16,
-    torch.float32,
-    torch.float64,
-}
-for _dtype_name in ("float8_e4m3fn", "float8_e5m2"):
-    if (_dtype := getattr(torch, _dtype_name, None)) is not None:
-        _SUPPORTED_DTYPES.add(_dtype)
-
-
-def default_host_weight_cache_root() -> str:
-    """Return the default root, intentionally independent of VLLM_CACHE_ROOT."""
-    return os.path.expanduser(DEFAULT_HOST_WEIGHT_CACHE_ROOT)
-
-
-def _persistent_buffer(module: nn.Module, local_name: str) -> bool:
-    parent_path, _, leaf_name = local_name.rpartition(".")
-    owner = module.get_submodule(parent_path)
-    return leaf_name not in owner._non_persistent_buffers_set
-
-
-def _named_parameters_with_duplicates(module: nn.Module):
-    try:
-        return module.named_parameters(remove_duplicate=False)
-    except TypeError:  # pragma: no cover - compatibility with old torch
-        return module.named_parameters()
-
-
-def _named_buffers_with_duplicates(module: nn.Module):
-    try:
-        return module.named_buffers(remove_duplicate=False)
-    except TypeError:  # pragma: no cover - compatibility with old torch
-        return module.named_buffers()
-
-
-def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None:
-    parent_path, _, leaf_name = runtime_name.rpartition(".")
-    parent = pipeline.get_submodule(parent_path)
-    tensor = parent._parameters.get(leaf_name)
-    if tensor is None:
-        tensor = parent._buffers.get(leaf_name)
-    return tensor
-
-
-def _collect_runtime_tensors(
-    pipeline: nn.Module,
-    dit_modules: Sequence[tuple[str, nn.Module]],
-) -> list[RuntimeTensor]:
-    """Collect a complete, safely remappable view of final CPU DiT tensors.
-
-    Parameters and persistent buffers must resolve to the same objects through
-    the owning pipeline, occupy complete contiguous CPU storages, use supported
-    dtypes, and have no aliases inside or outside the discovered DiT modules.
-    The stable name ordering is reused by content hashing and publication.
-    """
-    records: dict[str, RuntimeTensor] = {}
-    for dit_name, dit_module in dit_modules:
-        for local_name, tensor in _named_parameters_with_duplicates(dit_module):
-            runtime_name = f"{dit_name}.{local_name}"
-            candidate = RuntimeTensor(runtime_name, tensor, "parameter")
-            existing = records.get(runtime_name)
-            if existing is not None and existing.tensor is not tensor:
-                raise HostWeightCacheError(
-                    "ambiguous_ownership",
-                    f"multiple DiT tensors resolve to {runtime_name!r}",
-                )
-            records[runtime_name] = candidate
-        for local_name, tensor in _named_buffers_with_duplicates(dit_module):
-            if not _persistent_buffer(dit_module, local_name):
-                continue
-            runtime_name = f"{dit_name}.{local_name}"
-            candidate = RuntimeTensor(runtime_name, tensor, "buffer")
-            existing = records.get(runtime_name)
-            if existing is not None and existing.tensor is not tensor:
-                raise HostWeightCacheError(
-                    "ambiguous_ownership",
-                    f"multiple DiT tensors resolve to {runtime_name!r}",
-                )
-            records[runtime_name] = candidate
-
-    if not records:
-        raise HostWeightCacheError(
-            "ambiguous_ownership",
-            "no DiT parameters or persistent buffers were discovered",
-        )
-
-    storage_owners: dict[tuple[int, int], str] = {}
-    for record in records.values():
-        tensor = record.tensor
-        try:
-            pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name)
-        except AttributeError as exc:
-            raise HostWeightCacheError(
-                "ambiguous_ownership",
-                f"DiT tensor {record.name!r} is not owned by the pipeline",
-            ) from exc
-        if pipeline_tensor is not tensor:
-            raise HostWeightCacheError(
-                "ambiguous_ownership",
-                f"DiT tensor {record.name!r} does not resolve to the discovered object",
-            )
-        if tensor.device.type != "cpu" or tensor.is_meta:
-            raise HostWeightCacheError(
-                "unsupported_tensor",
-                f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}",
-            )
-        if tensor.layout != torch.strided:
-            raise HostWeightCacheError(
-                "unsupported_tensor",
-                f"{record.name!r} uses unsupported layout {tensor.layout}",
-            )
-        if not tensor.is_contiguous():
-            raise HostWeightCacheError(
-                "unsupported_tensor",
-                f"{record.name!r} is non-contiguous with stride {tensor.stride()}",
-            )
-        if tensor.dtype not in _SUPPORTED_DTYPES:
-            raise HostWeightCacheError(
-                "unsupported_tensor",
-                f"{record.name!r} uses unsupported dtype {tensor.dtype}",
-            )
-        if hasattr(tensor, "to_local"):
-            raise HostWeightCacheError(
-                "unsupported_tensor",
-                f"{record.name!r} is a distributed tensor",
-            )
-
-        tensor_nbytes = tensor.numel() * tensor.element_size()
-        storage = tensor.untyped_storage()
-        if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes:
-            raise HostWeightCacheError(
-                "unsupported_tensor",
-                f"{record.name!r} is a view into a larger storage",
-            )
-        if tensor_nbytes:
-            storage_id = (storage.data_ptr(), storage.nbytes())
-            owner = storage_owners.setdefault(storage_id, record.name)
-            if owner != record.name:
-                raise HostWeightCacheError(
-                    "unsupported_alias",
-                    f"{record.name!r} shares storage with {owner!r}",
-                )
-
-    # The ownership boundary is the complete pipeline, not only the discovered
-    # DiT traversal. Reject a DiT storage that is also registered by an encoder,
-    # VAE, resident component, or non-persistent buffer.
-    for pipeline_name, tensor in chain(
-        _named_parameters_with_duplicates(pipeline),
-        _named_buffers_with_duplicates(pipeline),
-    ):
-        if tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0:
-            continue
-        storage = tensor.untyped_storage()
-        owner = storage_owners.get((storage.data_ptr(), storage.nbytes()))
-        if owner is not None and owner != pipeline_name:
-            raise HostWeightCacheError(
-                "unsupported_alias",
-                f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}",
-            )
-
-    return [records[name] for name in sorted(records)]
-
-
-def _publish_entry(
-    *,
-    cache_root: Path,
-    entry_dir: Path,
-    cache_key: str,
-    layout_identity: dict[str, Any],
-    content_digest: str,
-    records: Sequence[RuntimeTensor],
-    max_shard_bytes: int,
-) -> None:
-    total_nbytes = sum(record.tensor.numel() * record.tensor.element_size() for record in records)
-    free_bytes = shutil.disk_usage(cache_root).free
-    required_bytes = total_nbytes + max(64 * 1024**2, total_nbytes // 100)
-    if free_bytes < required_bytes:
-        raise HostWeightCacheError(
-            "insufficient_disk",
-            f"host weight cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}",
-        )
-
-    tmp_root = cache_root / ".tmp"
-    tmp_root.mkdir(parents=True, exist_ok=True)
-    temp_dir = tmp_root / f"{cache_key}.{os.getpid()}.{uuid.uuid4().hex}"
-    temp_dir.mkdir()
-    try:
-        shards = split_shards(records, max_shard_bytes)
-        tensor_manifest: dict[str, dict[str, Any]] = {}
-        file_manifest: dict[str, dict[str, Any]] = {}
-        shard_count = len(shards)
-        for index, shard in enumerate(shards, start=1):
-            filename = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
-            path = temp_dir / filename
-            save_file(
-                {record.name: record.tensor.detach() for record in shard},
-                str(path),
-                metadata={"format": "pt", "dlo_host_weight_cache_schema": str(HOST_WEIGHT_CACHE_SCHEMA_VERSION)},
-            )
-            fsync_file(path)
-            file_manifest[filename] = {
-                "size": path.stat().st_size,
-                "sha256": file_digest(path),
-            }
-            for record in shard:
-                tensor_manifest[record.name] = {
-                    **tensor_metadata(record),
-                    "file": filename,
-                    "storage_key": record.name,
-                }
-
-        manifest = {
-            "schema_version": HOST_WEIGHT_CACHE_SCHEMA_VERSION,
-            "cache_key": cache_key,
-            "layout_identity": layout_identity,
-            "runtime_content_sha256": content_digest,
-            "total_tensor_bytes": total_nbytes,
-            "created_at_unix_seconds": time.time(),
-            "writer_pid": os.getpid(),
-            "files": file_manifest,
-            "tensors": tensor_manifest,
-        }
-        manifest_path = temp_dir / "manifest.json"
-        with manifest_path.open("w", encoding="utf-8") as handle:
-            json.dump(manifest, handle, sort_keys=True, separators=(",", ":"))
-            handle.flush()
-            os.fsync(handle.fileno())
-        fsync_directory(temp_dir)
-
-        invalid_dir: Path | None = None
-        if entry_dir.exists():
-            invalid_dir = tmp_root / f"{cache_key}.invalid.{os.getpid()}.{uuid.uuid4().hex}"
-            os.replace(entry_dir, invalid_dir)
-        os.replace(temp_dir, entry_dir)
-        fsync_directory(entry_dir.parent)
-        if invalid_dir is not None:
-            shutil.rmtree(invalid_dir, ignore_errors=True)
-    except Exception:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-        raise
-
-
-def _validate_entry(
-    *,
-    entry_dir: Path,
-    cache_key: str,
-    layout_identity: dict[str, Any],
-    content_digest: str,
-    expected_records: Sequence[RuntimeTensor],
-) -> HostWeightPlan:
-    manifest_path = entry_dir / "manifest.json"
-    try:
-        with manifest_path.open(encoding="utf-8") as handle:
-            manifest = json.load(handle)
-    except (OSError, ValueError) as exc:
-        raise HostWeightCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc
-
-    if manifest.get("schema_version") != HOST_WEIGHT_CACHE_SCHEMA_VERSION:
-        raise HostWeightCacheError("entry_rejected", "host weight cache schema mismatch")
-    if manifest.get("cache_key") != cache_key:
-        raise HostWeightCacheError("entry_rejected", "host weight cache key mismatch")
-    if manifest.get("layout_identity") != layout_identity:
-        raise HostWeightCacheError("entry_rejected", "host weight cache layout identity mismatch")
-    if manifest.get("runtime_content_sha256") != content_digest:
-        raise HostWeightCacheError("entry_rejected", "host weight cache content identity mismatch")
-
-    expected_metadata = {record.name: tensor_metadata(record) for record in expected_records}
-    tensor_manifest = manifest.get("tensors")
-    file_manifest = manifest.get("files")
-    if not isinstance(tensor_manifest, dict) or not isinstance(file_manifest, dict):
-        raise HostWeightCacheError("entry_rejected", "host weight cache manifest has invalid sections")
-    if set(tensor_manifest) != set(expected_metadata):
-        raise HostWeightCacheError("entry_rejected", "host weight cache tensor names differ from the loaded DiT")
-
-    for filename, metadata in file_manifest.items():
-        if Path(filename).name != filename or not isinstance(metadata, dict):
-            raise HostWeightCacheError("entry_rejected", f"invalid cache filename {filename!r}")
-        path = entry_dir / filename
-        try:
-            stat = path.stat()
-        except OSError as exc:
-            raise HostWeightCacheError("entry_rejected", f"missing cache shard {path}: {exc}") from exc
-        if stat.st_size != metadata.get("size") or file_digest(path) != metadata.get("sha256"):
-            raise HostWeightCacheError("entry_rejected", f"content digest mismatch for cache shard {path}")
-
-    bindings: dict[str, TensorBinding] = {}
-    mapped_records: list[RuntimeTensor] = []
-    with contextlib.ExitStack() as stack:
-        handles = {
-            filename: stack.enter_context(safe_open(entry_dir / filename, framework="pt", device="cpu"))
-            for filename in file_manifest
-        }
-        expected_keys_by_file: dict[str, set[str]] = {filename: set() for filename in file_manifest}
-        for record in expected_records:
-            stored = tensor_manifest.get(record.name)
-            if not isinstance(stored, dict):
-                raise HostWeightCacheError("entry_rejected", f"invalid tensor metadata for {record.name!r}")
-            filename = stored.get("file")
-            storage_key = stored.get("storage_key")
-            if filename not in handles or not isinstance(storage_key, str):
-                raise HostWeightCacheError("entry_rejected", f"invalid storage binding for {record.name!r}")
-            actual_metadata = {key: stored.get(key) for key in expected_metadata[record.name]}
-            if actual_metadata != expected_metadata[record.name]:
-                raise HostWeightCacheError("entry_rejected", f"metadata mismatch for {record.name!r}")
-            expected_keys_by_file[filename].add(storage_key)
-            try:
-                tensor = handles[filename].get_tensor(storage_key)
-            except Exception as exc:
-                raise HostWeightCacheError(
-                    "entry_rejected",
-                    f"cannot map {record.name!r} from {filename}: {exc}",
-                ) from exc
-            if tuple(tensor.shape) != tuple(record.tensor.shape) or tensor.dtype != record.tensor.dtype:
-                raise HostWeightCacheError("entry_rejected", f"mapped metadata mismatch for {record.name!r}")
-            mapped_records.append(RuntimeTensor(record.name, tensor, record.kind))
-            bindings[record.name] = TensorBinding(storage_key=storage_key, file_path=str(entry_dir / filename))
-
-        for filename, expected_keys in expected_keys_by_file.items():
-            if set(handles[filename].keys()) != expected_keys:
-                raise HostWeightCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}")
-
-        if runtime_content_digest(mapped_records) != content_digest:
-            raise HostWeightCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher")
-
-    return HostWeightPlan(
-        backing_kind="host_weight_cache",
-        bindings=bindings,
-        runtime_layout_key=cache_key,
-        expected_file_digests={
-            str(entry_dir / filename): metadata["sha256"] for filename, metadata in file_manifest.items()
-        },
-    )
-
-
-def build_host_weight_cache_plan(
-    pipeline: nn.Module,
-    *,
-    dit_modules: Sequence[tuple[str, nn.Module]],
-    loader_type: type,
-    cache_root: str | os.PathLike[str] | None,
-    lock_timeout_seconds: float,
-    max_shard_bytes: int,
-    model_identity: str | None,
-    revision: str | None,
-    runtime_dtype: Any,
-    load_format: str,
-    loader_inputs: Mapping[str, Any],
-    tensor_parallel_size: int,
-    tensor_parallel_rank: int,
-    sequence_parallel_guard: Mapping[str, Any],
-    use_hsdp: bool,
-    enable_expert_parallel: bool,
-    quantization_config: Any,
-    cfg_parallel_size: int,
-    pipeline_parallel_size: int,
-) -> HostWeightPlanResult:
-    """Build a host-weight plan from the host weight cache after ordinary loading.
-
-    This is the module's main entry point. It first rejects unsupported runtime
-    layouts, then discovers and validates the final CPU DiT tensors and derives
-    a content-addressed identity from their metadata, bytes, loader inputs, and
-    parallelism guards. A valid existing entry is joined immediately; otherwise
-    one process publishes the entry under a per-identity file lock and peers
-    validate it after the lock is released.
-
-    Success returns bindings to immutable safetensors mmap files. Any expected
-    incompatibility, lock, publication, or validation failure returns a
-    fail-closed result with no plan and a stable fallback code, leaving the
-    already-loaded pipeline tensors intact.
-    """
-    if load_format == "dummy":
-        return HostWeightPlanResult(
-            None,
-            "dummy/random weights do not have a reusable runtime identity",
-            "unsupported_load_format",
-        )
-    if use_hsdp:
-        return HostWeightPlanResult(
-            None, "HSDP/DTensor layouts are not host weight cache compatible", "unsupported_hsdp"
-        )
-    if enable_expert_parallel:
-        return HostWeightPlanResult(
-            None,
-            "expert-parallel weight ownership is outside host weight cache v1",
-            "unsupported_expert_parallel",
-        )
-    if quantization_config is not None:
-        return HostWeightPlanResult(
-            None,
-            "quantized runtime layouts are not supported by the first host weight cache version",
-            "unsupported_quantization",
-        )
-    if cfg_parallel_size != 1:
-        return HostWeightPlanResult(None, "CFG-parallel component ownership is not proven", "unsupported_cfg_parallel")
-    if pipeline_parallel_size != 1:
-        return HostWeightPlanResult(
-            None, "pipeline-parallel ownership is outside host weight cache v1", "unsupported_pp"
-        )
-    if not math.isfinite(lock_timeout_seconds) or lock_timeout_seconds <= 0:
-        return HostWeightPlanResult(None, "host weight cache lock timeout must be positive", "invalid_configuration")
-    if max_shard_bytes <= 0:
-        return HostWeightPlanResult(None, "host weight cache shard size must be positive", "invalid_configuration")
-
-    try:
-        records = _collect_runtime_tensors(pipeline, dit_modules)
-        content_digest = runtime_content_digest(records)
-        # The cache stores loader-produced ownership: TP coordinates differ,
-        # while equivalent DP replicas and SP ranks may share one entry. The
-        # conservative SP guard prevents reuse across unproven implementations.
-        layout_identity = {
-            "schema_version": HOST_WEIGHT_CACHE_SCHEMA_VERSION,
-            "model": (canonicalize_existing_local_path(model_identity) if model_identity is not None else None),
-            "revision": revision,
-            "runtime_dtype": str(runtime_dtype),
-            "load_format": load_format,
-            "loader_inputs": normalize_identity(loader_inputs),
-            "loader_implementation_sha256": implementation_fingerprint(loader_type, pipeline, dit_modules),
-            "components": [name for name, _ in dit_modules],
-            "tensor_parallel_size": tensor_parallel_size,
-            "tensor_parallel_rank": tensor_parallel_rank,
-            "sequence_parallel_guard": normalize_identity(sequence_parallel_guard),
-            "runtime_content_sha256": content_digest,
-        }
-        cache_key = hashlib.sha256(canonical_json(layout_identity)).hexdigest()
-        root = Path(cache_root or default_host_weight_cache_root()).expanduser().resolve()
-        entries_root = root / f"v{HOST_WEIGHT_CACHE_SCHEMA_VERSION}"
-        entry_dir = entries_root / cache_key
-        entries_root.mkdir(parents=True, exist_ok=True)
-
-        if entry_dir.is_dir():
-            try:
-                plan = _validate_entry(
-                    entry_dir=entry_dir,
-                    cache_key=cache_key,
-                    layout_identity=layout_identity,
-                    content_digest=content_digest,
-                    expected_records=records,
-                )
-                logger.info("Reusing DLO host weight cache entry %s from %s", cache_key[:12], entry_dir)
-                return HostWeightPlanResult(plan)
-            except HostWeightCacheError as exc:
-                logger.warning("DLO host weight cache entry %s will be rebuilt: %s", cache_key[:12], exc)
-
-        lock_path = root / ".locks" / f"{cache_key}.lock"
-        with exclusive_lock(lock_path, lock_timeout_seconds):
-            remove_stale_temps(root, cache_key)
-            if entry_dir.is_dir():
-                try:
-                    plan = _validate_entry(
-                        entry_dir=entry_dir,
-                        cache_key=cache_key,
-                        layout_identity=layout_identity,
-                        content_digest=content_digest,
-                        expected_records=records,
-                    )
-                    logger.info("Joined DLO host weight cache entry %s from %s", cache_key[:12], entry_dir)
-                    return HostWeightPlanResult(plan)
-                except HostWeightCacheError:
-                    pass
-
-            _publish_entry(
-                cache_root=root,
-                entry_dir=entry_dir,
-                cache_key=cache_key,
-                layout_identity=layout_identity,
-                content_digest=content_digest,
-                records=records,
-                max_shard_bytes=max_shard_bytes,
-            )
-            plan = _validate_entry(
-                entry_dir=entry_dir,
-                cache_key=cache_key,
-                layout_identity=layout_identity,
-                content_digest=content_digest,
-                expected_records=records,
-            )
-            logger.info(
-                "Published DLO host weight cache entry %s with %d tensors at %s",
-                cache_key[:12],
-                len(records),
-                entry_dir,
-            )
-            return HostWeightPlanResult(plan)
-    except HostWeightCacheError as exc:
-        return HostWeightPlanResult(None, str(exc), exc.code)
-    except IdentityNormalizationError as exc:
-        return HostWeightPlanResult(None, str(exc), "unstable_identity")
-    except Exception as exc:
-        return HostWeightPlanResult(None, f"host weight cache operation failed: {exc}", "cache_operation_failed")
-
-
-__all__ = [
-    "DEFAULT_LOCK_TIMEOUT_SECONDS",
-    "DEFAULT_HOST_WEIGHT_CACHE_ROOT",
-    "DEFAULT_SHARD_SIZE_BYTES",
-    "HOST_WEIGHT_CACHE_SCHEMA_VERSION",
-    "build_host_weight_cache_plan",
-    "default_host_weight_cache_root",
-]

vllm_omni/diffusion/model_loader/host_weight_cache_utils.py

@@ -1,287 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Reusable identity, hashing, and file helpers for the DLO host weight cache."""
-
-from __future__ import annotations
-
-import contextlib
-import dataclasses
-import errno
-import fcntl
-import hashlib
-import inspect
-import json
-import os
-import shutil
-import time
-from collections.abc import Iterator, Mapping, Sequence
-from dataclasses import dataclass
-from enum import Enum
-from pathlib import Path
-from typing import Protocol, TypeAlias, TypeVar
-
-import torch
-from torch import nn
-
-JsonScalar: TypeAlias = None | bool | int | float | str
-JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]
-PathIdentity: TypeAlias = str | os.PathLike[str] | os.PathLike[bytes]
-_HASH_CHUNK_BYTES = 64 * 1024**2
-_FILE_HASH_CHUNK_BYTES = 8 * 1024**2
-
-
-class IdentityNormalizationError(ValueError):
-    """Raised when a cache identity contains a process-unstable value."""
-
-
-@dataclass(frozen=True)
-class RuntimeTensor:
-    """One final runtime tensor and its manifest identity."""
-
-    name: str
-    tensor: torch.Tensor
-    kind: str
-
-
-class HostWeightCacheError(RuntimeError):
-    """Expected cache failure that should fall back with a stable code."""
-
-    def __init__(self, code: str, message: str):
-        super().__init__(message)
-        self.code = code
-
-
-class _HashUpdater(Protocol):
-    def update(self, data: bytes | bytearray | memoryview) -> object: ...
-
-
-_RuntimeTensorT = TypeVar("_RuntimeTensorT", bound=RuntimeTensor)
-
-
-def _type_identity(value: object) -> str:
-    value_type = value if isinstance(value, type) else type(value)
-    return f"{value_type.__module__}.{value_type.__qualname__}"
-
-
-def canonicalize_existing_local_path(value: PathIdentity) -> str:
-    """Resolve equivalent local paths while leaving model repository IDs alone."""
-    original = os.fsdecode(os.fspath(value))
-    candidate = Path(original).expanduser()
-    try:
-        if candidate.exists():
-            return str(candidate.resolve())
-    except OSError:
-        pass
-    return original
-
-
-def normalize_identity(value: object) -> JsonValue:
-    """Convert supported loader inputs into deterministic JSON-compatible data."""
-    if value is None or isinstance(value, (bool, int, float, str)):
-        return value
-    if isinstance(value, Enum):
-        return {"enum": _type_identity(value), "name": value.name}
-    if isinstance(value, torch.dtype):
-        return str(value)
-    if isinstance(value, os.PathLike):
-        return canonicalize_existing_local_path(value)
-    if isinstance(value, type):
-        return _type_identity(value)
-    if dataclasses.is_dataclass(value):
-        return normalize_identity(dataclasses.asdict(value))
-
-    model_dump = getattr(value, "model_dump", None)
-    if callable(model_dump):
-        try:
-            dumped = model_dump(mode="json")
-        except (TypeError, ValueError):
-            dumped = model_dump()
-        return normalize_identity(dumped)
-
-    if isinstance(value, Mapping):
-        normalized: dict[str, JsonValue] = {}
-        for key, item in value.items():
-            if not isinstance(key, str):
-                raise IdentityNormalizationError("host weight cache identity mappings require string keys")
-            normalized[key] = normalize_identity(item)
-        return {key: normalized[key] for key in sorted(normalized)}
-    if isinstance(value, (set, frozenset)):
-        return sorted((normalize_identity(item) for item in value), key=canonical_json)
-    if isinstance(value, (list, tuple)):
-        return [normalize_identity(item) for item in value]
-
-    raise IdentityNormalizationError(
-        f"host weight cache identity does not support values of type {_type_identity(value)}"
-    )
-
-
-def canonical_json(value: object) -> bytes:
-    """Serialize a supported identity with stable ordering and separators."""
-    return json.dumps(
-        normalize_identity(value),
-        sort_keys=True,
-        separators=(",", ":"),
-        ensure_ascii=True,
-    ).encode()
-
-
-def implementation_fingerprint(
-    loader_type: type,
-    pipeline: nn.Module,
-    dit_modules: Sequence[tuple[str, nn.Module]],
-) -> str:
-    """Hash relevant loader and module implementations for cache identity."""
-    objects: list[object] = [loader_type, type(pipeline)]
-    load_weights = getattr(type(pipeline), "load_weights", None)
-    if load_weights is not None:
-        objects.append(load_weights)
-    for _, dit_module in dit_modules:
-        objects.append(type(dit_module))
-        post_load = getattr(type(dit_module), "post_load_weights", None)
-        if post_load is not None:
-            objects.append(post_load)
-
-    digest = hashlib.sha256()
-    identities: set[str] = set()
-    for obj in objects:
-        identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}"
-        if identity in identities:
-            continue
-        identities.add(identity)
-        digest.update(identity.encode())
-        try:
-            digest.update(inspect.getsource(obj).encode())
-        except (OSError, TypeError):
-            pass
-    return digest.hexdigest()
-
-
-def tensor_metadata(record: RuntimeTensor) -> dict[str, JsonValue]:
-    """Return stable manifest metadata for one runtime tensor."""
-    tensor = record.tensor
-    return {
-        "kind": record.kind,
-        "shape": list(tensor.shape),
-        "dtype": str(tensor.dtype),
-        "nbytes": tensor.numel() * tensor.element_size(),
-        "layout": "contiguous",
-    }
-
-
-def _update_hash_with_tensor(digest: _HashUpdater, tensor: torch.Tensor) -> None:
-    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
-    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
-        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
-        digest.update(memoryview(chunk.numpy()))
-
-
-def runtime_content_digest(records: Sequence[RuntimeTensor]) -> str:
-    """Hash ordered runtime tensor metadata and contents."""
-    digest = hashlib.sha256()
-    for record in records:
-        digest.update(canonical_json({"name": record.name, **tensor_metadata(record)}))
-        _update_hash_with_tensor(digest, record.tensor)
-    return digest.hexdigest()
-
-
-def file_digest(path: Path) -> str:
-    """Hash one cache file with a bounded read working set."""
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        while chunk := handle.read(_FILE_HASH_CHUNK_BYTES):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-@contextlib.contextmanager
-def exclusive_lock(path: Path, timeout_seconds: float) -> Iterator[None]:
-    """Acquire one non-blocking file lock within a bounded wait."""
-    path.parent.mkdir(parents=True, exist_ok=True)
-    deadline = time.monotonic() + timeout_seconds
-    with path.open("a+b") as handle:
-        while True:
-            try:
-                fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
-                break
-            except OSError as exc:
-                if exc.errno not in (errno.EACCES, errno.EAGAIN):
-                    raise HostWeightCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc
-                remaining = deadline - time.monotonic()
-                if remaining <= 0:
-                    raise HostWeightCacheError(
-                        "lock_timeout",
-                        f"timed out after {timeout_seconds:g}s waiting for host weight cache writer {path.name}",
-                    ) from exc
-                time.sleep(min(0.2, remaining))
-        try:
-            yield
-        finally:
-            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
-
-
-def fsync_file(path: Path) -> None:
-    """Persist one cache file before publication."""
-    with path.open("rb") as handle:
-        os.fsync(handle.fileno())
-
-
-def fsync_directory(path: Path) -> None:
-    """Persist directory entries needed for atomic publication."""
-    fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
-    try:
-        os.fsync(fd)
-    finally:
-        os.close(fd)
-
-
-def split_shards(
-    records: Sequence[_RuntimeTensorT],
-    max_shard_bytes: int,
-) -> list[list[_RuntimeTensorT]]:
-    """Split ordered tensors without exceeding the target when possible."""
-    shards: list[list[_RuntimeTensorT]] = []
-    current: list[_RuntimeTensorT] = []
-    current_bytes = 0
-    for record in records:
-        nbytes = record.tensor.numel() * record.tensor.element_size()
-        if current and current_bytes + nbytes > max_shard_bytes:
-            shards.append(current)
-            current = []
-            current_bytes = 0
-        current.append(record)
-        current_bytes += nbytes
-    if current:
-        shards.append(current)
-    return shards
-
-
-def remove_stale_temps(cache_root: Path, cache_key: str) -> None:
-    """Remove abandoned temporary entries for one cache identity."""
-    temp_root = cache_root / ".tmp"
-    if not temp_root.is_dir():
-        return
-    for path in temp_root.glob(f"{cache_key}.*"):
-        if path.is_dir():
-            shutil.rmtree(path, ignore_errors=True)
-        else:
-            path.unlink(missing_ok=True)
-
-
-__all__ = [
-    "IdentityNormalizationError",
-    "JsonValue",
-    "HostWeightCacheError",
-    "RuntimeTensor",
-    "canonical_json",
-    "canonicalize_existing_local_path",
-    "exclusive_lock",
-    "file_digest",
-    "fsync_directory",
-    "fsync_file",
-    "implementation_fingerprint",
-    "normalize_identity",
-    "remove_stale_temps",
-    "runtime_content_digest",
-    "split_shards",
-    "tensor_metadata",
-]

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -9,7 +9,6 @@ import json
 import os
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
-from typing import Literal
 
 import torch
 from safetensors import safe_open
@@ -29,7 +28,7 @@ TensorTransform = Callable[[torch.Tensor], torch.Tensor]
 class TensorBinding:
     """One runtime tensor backed by one safetensors entry."""
 
-    storage_key: str
+    checkpoint_key: str
     file_path: str
     transform: TensorTransform | None = None
 
@@ -38,16 +37,9 @@ class TensorBinding:
 class HostWeightPlan:
     """Complete, prevalidated host backing consumed by an offload backend."""
 
-    backing_kind: Literal["checkpoint_mmap", "host_weight_cache"]
+    backing_kind: str
     bindings: dict[str, TensorBinding]
     planned_source_prefixes: frozenset[str] = frozenset()
-    runtime_layout_key: str | None = None
-    expected_file_digests: dict[str, str] | None = None
-
-    @property
-    def requires_backend(self) -> bool:
-        """Whether ordinary loading was skipped and this plan must be consumed."""
-        return self.backing_kind == "checkpoint_mmap"
 
 
 @dataclass(frozen=True)
@@ -56,7 +48,6 @@ class HostWeightPlanResult:
 
     plan: HostWeightPlan | None
     fallback_reason: str | None = None
-    fallback_code: str | None = None
 
 
 class _PlanIncompatibleError(RuntimeError):
@@ -267,11 +258,11 @@ def _validate_source_metadata(
         with safe_open(file_path, framework="pt", device="cpu") as handle:
             available = set(handle.keys())
             for runtime_name, target, binding in entries:
-                if binding.storage_key not in available:
+                if binding.checkpoint_key not in available:
                     raise _PlanIncompatibleError(
-                        f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}"
+                        f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
                     )
-                source = handle.get_slice(binding.storage_key)
+                source = handle.get_slice(binding.checkpoint_key)
                 source_shape = tuple(source.get_shape())
                 source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
                 if source_shape != tuple(target.shape):
@@ -296,19 +287,11 @@ def build_checkpoint_mmap_plan(
 ) -> HostWeightPlanResult:
     """Build a complete direct-checkpoint plan or return a fallback reason."""
     if tensor_parallel_size != 1:
-        return HostWeightPlanResult(
-            None,
-            f"TP={tensor_parallel_size} requires the ordinary loader",
-            "unsupported_tp",
-        )
+        return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
     if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp")
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
     if online_quantization:
-        return HostWeightPlanResult(
-            None,
-            "online quantization requires the ordinary loader",
-            "unsupported_online_quantization",
-        )
+        return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
     if not callable(remap_fn):
@@ -345,14 +328,14 @@ def build_checkpoint_mmap_plan(
                 )
             checkpoint_key, file_path = model_to_ckpt[runtime_name]
             bindings[runtime_name] = TensorBinding(
-                storage_key=checkpoint_key,
+                checkpoint_key=checkpoint_key,
                 file_path=file_path,
                 transform=policy.transform if policy is not None else None,
             )
 
         _validate_source_metadata(required, bindings)
     except (OSError, ValueError, _PlanIncompatibleError) as exc:
-        return HostWeightPlanResult(None, str(exc), "checkpoint_plan_incompatible")
+        return HostWeightPlanResult(None, str(exc))
 
     return HostWeightPlanResult(
         HostWeightPlan(

vllm_omni/diffusion/model_metadata.py

@@ -56,6 +56,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
     "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True),
     "WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True),
     "WanS2VPipeline": DiffusionModelMetadata(attention_mask_free=True),
+    "SanaWmPipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=1,
+    ),
 }
 
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -29,7 +29,7 @@ from vllm_omni.diffusion.cache.cachedit import (
 )
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
-    get_dit_group,
+    get_world_group,
     init_world_group,
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -423,7 +423,7 @@ def _validate_reference_image(image: Image.Image) -> None:
 def _dit_rank_world() -> tuple[Any, int, int]:
     if not dist.is_initialized():
         return None, 0, 1
-    group = get_dit_group()
+    group = get_world_group().device_group
     return group, dist.get_rank(group), dist.get_world_size(group)
 
 
@@ -1715,7 +1715,7 @@ class MiniMaxH3Pipeline(
                     dist.broadcast(
                         has_audio_tensor,
                         src=0,
-                        group=get_dit_group(),
+                        group=get_world_group().device_group,
                     )
                 has_audio = [bool(value) for value in has_audio_tensor.tolist()]
 

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -18,7 +18,7 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedVaeMixin,
 )
-from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group
 from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 
 from .packed_tokens import minimax_h3_patchify_video_latent
@@ -157,7 +157,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
     ) -> None:
         if mode != "tile":
             raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
-        group = get_dit_group()
+        group = get_world_group().device_group
         world_size = dist.get_world_size(group)
         rank = dist.get_rank(group)
         parallel_size = int(parallel_size)

vllm_omni/diffusion/models/sana_wm/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SANA-WM diffusion model integration."""
+
+from vllm_omni.diffusion.models.sana_wm.config import SanaWmConfig
+from vllm_omni.diffusion.models.sana_wm.pipeline_sana_wm import (
+    SANA_WM_MODEL_ID,
+    SANA_WM_OUTPUT_HEIGHT,
+    SANA_WM_OUTPUT_WIDTH,
+    SanaWmPipeline,
+    get_sana_wm_pre_process_func,
+)
+from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmTransformer3DModel
+
+__all__ = [
+    "SANA_WM_MODEL_ID",
+    "SANA_WM_OUTPUT_HEIGHT",
+    "SANA_WM_OUTPUT_WIDTH",
+    "SanaWmConfig",
+    "SanaWmPipeline",
+    "SanaWmTransformer3DModel",
+    "get_sana_wm_pre_process_func",
+    "normalize_sana_wm_payload",
+]

vllm_omni/diffusion/models/sana_wm/camera_control.py

@@ -0,0 +1,356 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Camera-control helpers for SANA-WM.
+
+This module mirrors the public NVlabs/Sana SANA-WM camera preparation path:
+WASD/IJKL action rollout or explicit camera-to-world poses are converted into
+relative poses, per-latent-frame ray metadata, and per-VAE-chunk Plucker maps.
+Model-local projection layers consume these tensors in the Stage-1 DiT.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+SANA_WM_DEFAULT_TRANSLATION_SPEED = 0.05
+SANA_WM_DEFAULT_ROTATION_SPEED_DEG = 1.2
+SANA_WM_DEFAULT_PITCH_LIMIT_DEG = 85.0
+SANA_WM_ALLOWED_ACTION_KEYS = frozenset("wasdijkl")
+SANA_WM_DEFAULT_VAE_STRIDE = (8, 32, 32)
+
+
+@dataclass(frozen=True)
+class SanaWmCameraCondition:
+    poses: Any | None = None
+    intrinsics: Any | None = None
+    action: str | None = None
+    num_frames: int | None = None
+    height: int = 704
+    width: int = 1280
+    translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED
+    rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG
+    pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG
+
+
+def _rot_x(angle_rad: float) -> np.ndarray:
+    c, s = np.cos(angle_rad), np.sin(angle_rad)
+    return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64)
+
+
+def _rot_y(angle_rad: float) -> np.ndarray:
+    c, s = np.cos(angle_rad), np.sin(angle_rad)
+    return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64)
+
+
+def _parse_action_segments(action: str) -> list[tuple[list[str], int]]:
+    """Parse an action string into ``(keys, duration)`` pairs without expanding it.
+
+    A single segment names an unbounded rollout (``"w-1000000000"`` is valid
+    syntax), so expanding here would allocate one list entry per frame before
+    anything checks the frame count. Callers total the durations and validate
+    against ``num_frames`` first; only :func:`_parse_action_string` expands.
+    """
+    cleaned = "".join(action.replace("，", ",").split())
+    if not cleaned:
+        raise ValueError("Sana-WM action string is empty.")
+
+    segments: list[tuple[list[str], int]] = []
+    for segment in cleaned.split(","):
+        if not segment or "-" not in segment:
+            raise ValueError(f"Invalid Sana-WM action segment {segment!r}: expected '<keys>-<duration>'.")
+        keys_part, duration = segment.rsplit("-", 1)
+        if not duration.isdigit() or int(duration) <= 0:
+            raise ValueError(f"Sana-WM action segment {segment!r} has a non-positive duration {duration!r}.")
+        keys_lower = keys_part.lower()
+        if keys_lower == "none":
+            keys: list[str] = []
+        else:
+            bad = sorted({key for key in keys_lower if key not in SANA_WM_ALLOWED_ACTION_KEYS})
+            if bad:
+                allowed = "".join(sorted(SANA_WM_ALLOWED_ACTION_KEYS))
+                raise ValueError(f"Sana-WM action segment {segment!r} contains unknown keys {bad}; allowed: {allowed}.")
+            keys = sorted(set(keys_lower))
+        segments.append((keys, int(duration)))
+    return segments
+
+
+def action_rollout_num_frames(action: str) -> int:
+    """Frames ``action`` rolls out: the identity start pose plus one per step.
+
+    Totals the durations arithmetically, so an oversized request is rejected
+    before anything allocates per-frame state.
+    """
+    return sum(duration for _, duration in _parse_action_segments(action)) + 1
+
+
+def _parse_action_string(action: str) -> list[list[str]]:
+    per_frame: list[list[str]] = []
+    for keys, duration in _parse_action_segments(action):
+        per_frame.extend([keys] * duration)
+    return per_frame
+
+
+def action_string_to_c2w(
+    action: str,
+    *,
+    translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED,
+    rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG,
+    pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG,
+) -> np.ndarray:
+    """Roll out an OpenCV-convention ``(N + 1, 4, 4)`` c2w trajectory."""
+
+    per_frame = _parse_action_string(action)
+    rotate_rad = math.radians(rotation_speed_deg)
+    pitch_limit_rad = math.radians(pitch_limit_deg)
+    current = np.eye(4, dtype=np.float64)
+    current_pitch = 0.0
+    poses = [current.copy()]
+
+    for keys in per_frame:
+        held = set(keys)
+        rotation = current[:3, :3]
+        translation = current[:3, 3]
+
+        pitch_delta = (rotate_rad if "i" in held else 0.0) - (rotate_rad if "k" in held else 0.0)
+        new_pitch = current_pitch + pitch_delta
+        if -pitch_limit_rad <= new_pitch <= pitch_limit_rad:
+            current_pitch = new_pitch
+        else:
+            pitch_delta = 0.0
+
+        yaw_delta = (rotate_rad if "l" in held else 0.0) - (rotate_rad if "j" in held else 0.0)
+        new_rotation = _rot_y(yaw_delta) @ rotation @ _rot_x(pitch_delta)
+
+        forward = new_rotation[:, 2].copy()
+        forward[1] = 0.0
+        right = new_rotation[:, 0].copy()
+        right[1] = 0.0
+        forward_norm = float(np.linalg.norm(forward))
+        right_norm = float(np.linalg.norm(right))
+        if forward_norm > 0:
+            forward /= forward_norm + 1e-6
+        if right_norm > 0:
+            right /= right_norm + 1e-6
+
+        move = np.zeros(3, dtype=np.float64)
+        if "w" in held:
+            move += forward * translation_speed
+        if "s" in held:
+            move -= forward * translation_speed
+        if "d" in held:
+            move += right * translation_speed
+        if "a" in held:
+            move -= right * translation_speed
+
+        current = np.eye(4, dtype=np.float64)
+        current[:3, :3] = new_rotation
+        current[:3, 3] = translation + move
+        poses.append(current.copy())
+
+    return np.stack(poses, axis=0).astype(np.float32)
+
+
+def get_pose_inverse(transform: torch.Tensor) -> torch.Tensor:
+    """Invert homogeneous rigid transforms with shape ``(..., 4, 4)``."""
+
+    rotation = transform[..., :3, :3]
+    translation = transform[..., :3, 3]
+    rotation_inv = rotation.transpose(-1, -2)
+    translation_inv = -torch.matmul(rotation_inv, translation.unsqueeze(-1)).squeeze(-1)
+    output = torch.eye(4, dtype=transform.dtype, device=transform.device).repeat(transform.shape[:-2] + (1, 1))
+    output[..., :3, :3] = rotation_inv
+    output[..., :3, 3] = translation_inv
+    return output
+
+
+def compute_raymap(
+    intrinsics: torch.Tensor,
+    poses: torch.Tensor,
+    height: int,
+    width: int,
+    *,
+    use_plucker: bool = True,
+) -> torch.Tensor:
+    """Compute SANA-WM geometry ray maps.
+
+    Args:
+        intrinsics: ``(T, 4)`` tensor in ``[fx, fy, cx, cy]`` order.
+        poses: ``(T, 4, 4)`` camera-to-world matrices.
+        height: raymap height in latent pixels.
+        width: raymap width in latent pixels.
+        use_plucker: return ``[direction, moment]`` if true, otherwise
+            ``[origin, direction]``.
+    """
+
+    if intrinsics.ndim != 2 or intrinsics.shape[-1] != 4:
+        raise ValueError(f"Sana-WM intrinsics must have shape (T, 4), got {tuple(intrinsics.shape)}.")
+    if poses.ndim != 3 or poses.shape[-2:] != (4, 4):
+        raise ValueError(f"Sana-WM poses must have shape (T, 4, 4), got {tuple(poses.shape)}.")
+    if intrinsics.shape[0] != poses.shape[0]:
+        raise ValueError("Sana-WM intrinsics and poses must have the same temporal length.")
+
+    frames = intrinsics.shape[0]
+    device = intrinsics.device
+    dtype = intrinsics.dtype
+    y_grid, x_grid = torch.meshgrid(
+        torch.arange(height, device=device, dtype=dtype),
+        torch.arange(width, device=device, dtype=dtype),
+        indexing="ij",
+    )
+    x_grid = x_grid[None].expand(frames, -1, -1)
+    y_grid = y_grid[None].expand(frames, -1, -1)
+
+    fx = intrinsics[:, 0].view(frames, 1, 1)
+    fy = intrinsics[:, 1].view(frames, 1, 1)
+    cx = intrinsics[:, 2].view(frames, 1, 1)
+    cy = intrinsics[:, 3].view(frames, 1, 1)
+    dirs_cam = torch.stack([(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], dim=-1)
+
+    rotation = poses[:, :3, :3]
+    translation = poses[:, :3, 3]
+    dirs_world = torch.einsum("tij,thwj->thwi", rotation, dirs_cam)
+    dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True).clamp_min(1e-12)
+    origins = translation.view(frames, 1, 1, 3).expand_as(dirs_world)
+
+    if use_plucker:
+        moments = torch.cross(origins, dirs_world, dim=-1)
+        return torch.cat([dirs_world, moments], dim=-1)
+    return torch.cat([origins, dirs_world], dim=-1)
+
+
+def intrinsics_to_vec4_array(intrinsics: Any, *, num_frames: int, height: int, width: int) -> np.ndarray:
+    """Normalize the intrinsics payload to an ``(F, 4)`` ``[fx, fy, cx, cy]`` array.
+
+    Accepts either ``None`` (derive from the output resolution) or the
+    ``{fx, fy, cx, cy}`` mapping — the only form the request contract exposes.
+    """
+
+    if intrinsics is None:
+        focal = float(max(height, width))
+        vec = np.array([focal, focal, width / 2.0, height / 2.0], dtype=np.float32)
+        return np.broadcast_to(vec, (num_frames, 4)).copy()
+
+    if isinstance(intrinsics, dict):
+        vec = np.array(
+            [intrinsics["fx"], intrinsics["fy"], intrinsics["cx"], intrinsics["cy"]],
+            dtype=np.float32,
+        )
+        return np.broadcast_to(vec, (num_frames, 4)).copy()
+
+    raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.")
+
+
+def _pack_camera_conditions(
+    poses: torch.Tensor,
+    intrinsics_latent: torch.Tensor,
+    *,
+    num_frames: int,
+    latent_height: int,
+    latent_width: int,
+    vae_time_stride: int,
+) -> dict[str, torch.Tensor]:
+    # This yields exactly the latent frame count the VAE produces:
+    # ``len(arange(0, n, s)) == ceil(n / s) == (n - 1) // s + 1``.
+    time_indices = torch.arange(0, num_frames, vae_time_stride)
+
+    raymap = torch.cat([poses[time_indices].reshape(len(time_indices), -1), intrinsics_latent[time_indices]], dim=-1)
+
+    chunks: list[torch.Tensor] = []
+    spatial_raymap_frames: list[torch.Tensor] = []
+    chunk_starts = time_indices - (vae_time_stride - 1)
+    for start, t_idx in zip(chunk_starts, time_indices):
+        start_idx = max(0, int(start))
+        end_idx = start_idx + vae_time_stride
+        chunk_poses = poses[start_idx:end_idx]
+        chunk_intrinsics = intrinsics_latent[start_idx:end_idx]
+        if chunk_poses.shape[0] < vae_time_stride:
+            pad = vae_time_stride - chunk_poses.shape[0]
+            chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)], dim=0)
+            chunk_intrinsics = torch.cat([chunk_intrinsics, chunk_intrinsics[-1:].repeat(pad, 1)], dim=0)
+        plucker = compute_raymap(chunk_intrinsics, chunk_poses, latent_height, latent_width, use_plucker=True)
+        chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_height, latent_width))
+
+        # Per-frame spatial ray-direction map for raymap_embedder: take the
+        # representative frame in each chunk (the anchor at t_idx) and use only
+        # the direction channels (last 3 of 6 in the origin+direction layout).
+        t = int(t_idx)
+        od = compute_raymap(
+            intrinsics_latent[t : t + 1],
+            poses[t : t + 1],
+            latent_height,
+            latent_width,
+            use_plucker=False,
+        )  # [1, H, W, 6]
+        spatial_raymap_frames.append(od[0, :, :, 3:])  # [H, W, 3] — direction only
+
+    chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3)
+    # spatial_raymap: [3, F_latent, H, W] — consumed by raymap_embedder
+    spatial_raymap = torch.stack(spatial_raymap_frames).permute(3, 0, 1, 2)
+    return {"raymap": raymap, "chunk_plucker": chunk_plucker, "spatial_raymap": spatial_raymap}
+
+
+def build_plucker_condition(
+    condition: SanaWmCameraCondition,
+    *,
+    vae_stride: tuple[int, int, int] | list[int] = SANA_WM_DEFAULT_VAE_STRIDE,
+) -> dict[str, torch.Tensor]:
+    """Build native Stage-1 camera tensors from normalized request metadata."""
+
+    if bool(condition.action) == (condition.poses is not None):
+        raise ValueError("Sana-WM camera condition requires exactly one of action or poses.")
+
+    if condition.action:
+        poses_c2w = action_string_to_c2w(
+            condition.action,
+            translation_speed=condition.translation_speed,
+            rotation_speed_deg=condition.rotation_speed_deg,
+            pitch_limit_deg=condition.pitch_limit_deg,
+        )
+    else:
+        poses_c2w = np.asarray(condition.poses, dtype=np.float32)
+
+    if poses_c2w.ndim != 3 or poses_c2w.shape[-2:] != (4, 4):
+        raise ValueError(f"Sana-WM c2w poses must have shape (F, 4, 4), got {poses_c2w.shape}.")
+
+    num_frames = condition.num_frames or int(poses_c2w.shape[0])
+    if poses_c2w.shape[0] < num_frames:
+        if condition.action:
+            num_frames = int(poses_c2w.shape[0])
+        else:
+            raise ValueError(f"Sana-WM c2w poses length {poses_c2w.shape[0]} is shorter than num_frames {num_frames}.")
+    poses_c2w = poses_c2w[:num_frames]
+
+    intrinsics = intrinsics_to_vec4_array(
+        condition.intrinsics,
+        num_frames=num_frames,
+        height=condition.height,
+        width=condition.width,
+    )
+
+    vae_time_stride = int(vae_stride[0])
+    vae_spatial_stride = int(vae_stride[-1])
+    latent_height = condition.height // vae_spatial_stride
+    latent_width = condition.width // vae_spatial_stride
+
+    poses = torch.from_numpy(poses_c2w).float()
+    first_inv = get_pose_inverse(poses[0:1]).squeeze(0)
+    poses_rel = torch.matmul(first_inv, poses[1:])
+    poses = torch.cat([torch.eye(4, dtype=poses.dtype).unsqueeze(0), poses_rel], dim=0)
+
+    intrinsics_latent = torch.from_numpy(intrinsics).float()
+    intrinsics_latent[:, [0, 2]] *= latent_width / float(condition.width)
+    intrinsics_latent[:, [1, 3]] *= latent_height / float(condition.height)
+
+    return _pack_camera_conditions(
+        poses,
+        intrinsics_latent,
+        num_frames=num_frames,
+        latent_height=latent_height,
+        latent_width=latent_width,
+        vae_time_stride=vae_time_stride,
+    )

vllm_omni/diffusion/models/sana_wm/config.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Config for the SANA-WM Stage-1 transformer.
+
+Loaded from the standard diffusers ``transformer/config.json`` (a flat dict
+keyed by the dataclass field names). The bespoke ``config.yaml`` release format
+is handled once, offline, at checkpoint-conversion time — the runtime only ever
+sees the converted, diffusers-standard config.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import Any
+
+
+def _as_tuple3(value: Any, default: tuple[int, int, int]) -> tuple[int, int, int]:
+    if value is None:
+        return default
+    if isinstance(value, int):
+        return (value, value, value)
+    if isinstance(value, (list, tuple)):
+        if len(value) == 2:
+            return (1, int(value[0]), int(value[1]))
+        if len(value) == 3:
+            return (int(value[0]), int(value[1]), int(value[2]))
+    return default
+
+
+# LTX-2 VAE compression ratios (SANA-WM ships ``AutoencoderKLLTX2Video``).
+# They live here rather than in the pipeline so the request-normalisation module
+# can validate latent geometry without importing the pipeline.
+SANA_WM_VAE_SPATIAL_COMPRESSION = 32
+SANA_WM_VAE_TEMPORAL_COMPRESSION = 8
+
+
+@dataclass(frozen=True)
+class SanaWmConfig:
+    """SANA-WM Stage-1 architecture and runtime config.
+
+    Defaults mirror the first public HF release. ``from_json`` (reading
+    ``transformer/config.json``) is the normal construction path.
+    """
+
+    architecture_name: str | None = "SanaMSVideoCamCtrl_1600M_P1_D20"
+    num_blocks: int = 20
+    hidden_size: int = 2240
+    mlp_ratio: float = 3.0
+    attn_type: str = "BidirectionalGDNTriton"
+    softmax_every_n: int = 4
+    linear_head_dim: int = 112
+    conv_kernel_size: int = 4
+    t_kernel_size: int = 3
+    k_conv_only: bool = True
+    ffn_type: str = "GLUMBConvTemp"
+    pos_embed_type: str = "wan_rope"
+    patch_size: tuple[int, int, int] = (1, 1, 1)
+    qk_norm: bool = True
+    cross_norm: bool = True
+    mixed_precision: str = "bf16"
+    fp32_attention: bool = True
+    image_size: int = 720
+    cam_attn_compress: int = 1
+    use_chunk_plucker_post_attn: bool = True
+    chunk_plucker_channels: int = 48
+    chunk_plucker_post_attn_blocks: int = 20
+    inference_flow_shift: float = 9.8
+    scheduler_type: str = "flow_dpm-solver"
+    chi_prompt: list[str] = field(default_factory=list)
+    model_max_length: int = 300
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> SanaWmConfig:
+        """Build from a flat mapping of field name -> value.
+
+        Unknown keys (e.g. the diffusers ``_class_name`` marker or the
+        ``latent_channels`` / ``prompt_channels`` constructor kwargs) are
+        ignored. ``patch_size`` is coerced from a JSON list back to a tuple.
+        """
+        field_names = {f.name for f in fields(cls)}
+        known = {k: v for k, v in data.items() if k in field_names}
+        if "patch_size" in known:
+            known["patch_size"] = _as_tuple3(known["patch_size"], cls.patch_size)
+        return cls(**known)
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> SanaWmConfig:
+        with open(path, encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.from_dict(data if isinstance(data, Mapping) else {})

vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py

@@ -0,0 +1,942 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sana-WM pipeline integration.
+
+This module wires the registry-visible surface, release-layout validation, and
+the Stage-1 sampling loop. The DiT, Gated-DeltaNet and camera stack are
+vLLM-Omni-native layers (``sana_wm_transformer.py`` / ``ucpe.py``); nothing here
+calls into NVlabs code at runtime.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, ClassVar
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import (
+    denormalize_latents,
+    normalize_latents,
+    pack_latents,
+    resolve_video_latent_shape,
+    unpack_latents,
+)
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.models.sana_wm.camera_control import (
+    SanaWmCameraCondition,
+    build_plucker_condition,
+)
+from vllm_omni.diffusion.models.sana_wm.config import (
+    SANA_WM_VAE_SPATIAL_COMPRESSION,
+    SANA_WM_VAE_TEMPORAL_COMPRESSION,
+    SanaWmConfig,
+)
+from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+    SANA_WM_STAGE1_PROMPT_CHANNELS,
+    SanaWmTransformer3DModel,
+)
+from vllm_omni.diffusion.models.schedulers import FlowMatchEulerDiscreteScheduler
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = init_logger(__name__)
+
+# The loader expects the standard diffusers layout (transformer/config.json +
+# transformer/diffusion_pytorch_model.safetensors + model_index.json), so this
+# repo is an offline conversion of the NVlabs release into that layout. It holds
+# Stage-1 only — transformer/ + vae/, with model_index.json naming
+# SanaWmPipeline — because startup prefetches the whole model repo
+# (allow_patterns=["*"]); a repo that also carried the refiner would pull tens of
+# GB this path never loads.
+SANA_WM_MODEL_ID = "BBBBruce/SANA-WM_bidirectional-stage1-diffusers"
+
+SANA_WM_STAGE1_DIT_FILE = "transformer/diffusion_pytorch_model.safetensors"
+SANA_WM_CONFIG_FILE = "transformer/config.json"
+SANA_WM_VAE_CONFIG_FILE = "vae/config.json"
+SANA_WM_VAE_WEIGHT_FILE = "vae/diffusion_pytorch_model.safetensors"
+SANA_WM_STAGE1_TEXT_ENCODER_ID = "google/gemma-2-2b-it"
+SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID = "Efficient-Large-Model/gemma-2-2b-it"
+SANA_WM_STAGE1_TEXT_ENCODER_ENV = "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER"
+SANA_WM_OUTPUT_HEIGHT = 704
+SANA_WM_OUTPUT_WIDTH = 1280
+SANA_WM_NUM_TRAIN_TIMESTEPS = 1000
+# Maximum supported clip length (20s @16fps) — the native envelope, not a
+# request default. Request-side defaults live in
+# ``models/sana_wm/request.py``.
+SANA_WM_NATIVE_NUM_FRAMES = 321
+# Fail-fast latent-token cap for the native Stage-1 path, sized to the native
+# 704x1280 output at the maximum clip length so it only rejects genuinely
+# oversized requests. Overridable per-request via ``sana_wm_native_max_tokens``.
+SANA_WM_NATIVE_MAX_TOKENS = (
+    ((SANA_WM_NATIVE_NUM_FRAMES - 1) // SANA_WM_VAE_TEMPORAL_COMPRESSION + 1)
+    * (SANA_WM_OUTPUT_HEIGHT // SANA_WM_VAE_SPATIAL_COMPRESSION)
+    * (SANA_WM_OUTPUT_WIDTH // SANA_WM_VAE_SPATIAL_COMPRESSION)
+)
+# The NVlabs production inference settings for this release. They belong in the
+# model rather than in a deploy YAML: single-stage diffusion models resolve
+# through ``create_default_diffusion``, which never reads a YAML's
+# ``default_sampling_params``, so a caller that omits either field lands here.
+SANA_WM_DEFAULT_NUM_INFERENCE_STEPS = 60
+SANA_WM_DEFAULT_GUIDANCE_SCALE = 5.0
+
+SANA_WM_STAGE1_PATTERNS = (
+    SANA_WM_CONFIG_FILE,
+    SANA_WM_STAGE1_DIT_FILE,
+    SANA_WM_VAE_CONFIG_FILE,
+    SANA_WM_VAE_WEIGHT_FILE,
+)
+SANA_WM_STAGE1_DIT_SUBFOLDER = "transformer"
+SANA_WM_STAGE1_DIT_BASENAME = Path(SANA_WM_STAGE1_DIT_FILE).name
+
+
+@dataclass(frozen=True)
+class SanaWmLocalPaths:
+    """Resolved local file paths for a SANA-WM snapshot."""
+
+    root: Path
+    config: Path
+    stage1_dit: Path
+    vae_config: Path
+    vae_weights: Path
+
+
+@dataclass(frozen=True)
+class SanaWmNativeParams:
+    """Resolved generation settings for one Stage-1 request.
+
+    Built by ``_native_params`` from the request payload and sampling params,
+    and consumed by ``_run_native_backend`` — this is the production path.
+    """
+
+    height: int
+    width: int
+    num_frames: int
+    num_inference_steps: int
+    seed: int
+    # Defaults to the NVlabs production scale, so a caller that never mentions
+    # guidance still gets the two-branch CFG forward the model was tuned for.
+    # An explicit ``guidance_scale <= 1.0`` still selects the single-branch path.
+    cfg_scale: float = SANA_WM_DEFAULT_GUIDANCE_SCALE
+
+
+def build_sana_wm_download_patterns() -> tuple[str, ...]:
+    """Return the minimal HF allow-patterns needed for SANA-WM."""
+
+    return tuple(SANA_WM_STAGE1_PATTERNS)
+
+
+def resolve_sana_wm_local_paths(snapshot_dir: str | Path) -> SanaWmLocalPaths:
+    root = Path(snapshot_dir)
+    return SanaWmLocalPaths(
+        root=root,
+        config=root / SANA_WM_CONFIG_FILE,
+        stage1_dit=root / SANA_WM_STAGE1_DIT_FILE,
+        vae_config=root / SANA_WM_VAE_CONFIG_FILE,
+        vae_weights=root / SANA_WM_VAE_WEIGHT_FILE,
+    )
+
+
+def validate_sana_wm_local_paths(paths: SanaWmLocalPaths) -> None:
+    required = [
+        paths.config,
+        paths.stage1_dit,
+        paths.vae_config,
+        paths.vae_weights,
+    ]
+    missing = [path for path in required if not path.exists()]
+    if missing:
+        joined = ", ".join(str(path) for path in missing)
+        raise FileNotFoundError(f"Missing SANA-WM checkpoint files: {joined}")
+
+
+def resolve_or_download_sana_wm_checkpoint(
+    model: str = SANA_WM_MODEL_ID,
+    *,
+    revision: str | None = None,
+    cache_dir: str | None = None,
+) -> SanaWmLocalPaths:
+    """Resolve a local SANA-WM tree or download the required HF files."""
+
+    root = Path(model)
+    if root.is_dir():
+        snapshot_dir = root
+    else:
+        snapshot_dir = Path(
+            download_weights_from_hf_specific(
+                model,
+                cache_dir,
+                list(build_sana_wm_download_patterns()),
+                revision=revision,
+                require_all=True,
+            )
+        )
+    paths = resolve_sana_wm_local_paths(snapshot_dir)
+    validate_sana_wm_local_paths(paths)
+    return paths
+
+
+def build_sana_wm_output_envelope(
+    *,
+    output: Any,
+    output_type: str,
+    metadata: dict[str, Any],
+) -> dict[str, Any]:
+    """Wrap a Stage-1 result in the canonical output envelope.
+
+    ``normalize_diffusion_postprocess_output`` splits ``{"payload": ...,
+    "metadata": ...}`` into the API-facing payload and the metadata groups, so
+    the model-specific diagnostics ride along under a ``sana_wm`` group instead
+    of the removed ``DiffusionOutput.custom_output`` field.
+    """
+    payload_key = "latents" if output_type == "latent" else "video"
+    return {"payload": {payload_key: output}, "metadata": {"sana_wm": metadata}}
+
+
+def get_sana_wm_pre_process_func(od_config: OmniDiffusionConfig):
+    del od_config
+
+    def pre_process_func(request: Any) -> Any:
+        from vllm_omni.inputs.data import OmniTextPrompt
+
+        sampling_params = getattr(request, "sampling_params", None)
+        prompt = request.prompt
+        prompt_mapping = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt
+        prompt_mapping = dict(prompt_mapping)
+
+        # The camera block travels as extra_params["sana_wm"], which serving
+        # merges into extra_args. Offline callers may also put it straight on
+        # the prompt.
+        if sampling_params is not None:
+            extra_args = getattr(sampling_params, "extra_args", None)
+            if isinstance(extra_args, Mapping):
+                block = extra_args.get("sana_wm")
+                if isinstance(block, Mapping):
+                    merged = dict(block)
+                    merged.update(prompt_mapping.get("sana_wm") or {})
+                    prompt_mapping["sana_wm"] = merged
+            # Geometry always comes from the standard sampling fields.
+            for field in ("num_frames", "height", "width"):
+                value = getattr(sampling_params, field, None)
+                if value is not None and not (field == "num_frames" and value == 1):
+                    prompt_mapping[field] = value
+
+        request.prompt = normalize_sana_wm_payload(prompt_mapping)
+
+        # Preprocess is the only place the request is normalised, so publish the
+        # canonical geometry back onto sampling_params for the rest of the stack.
+        if sampling_params is not None:
+            payload = request.prompt["additional_information"]["sana_wm"]
+            sampling_params.num_frames = payload["num_frames"]
+            sampling_params.height = payload["height"]
+            sampling_params.width = payload["width"]
+        return request
+
+    return pre_process_func
+
+
+class SanaWmPipeline(
+    nn.Module,
+    CFGParallelMixin,
+    SupportImageInput,
+    SupportsComponentDiscovery,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
+):
+    """Stage-1 SANA-WM image-to-video pipeline."""
+
+    support_image_input: ClassVar[bool] = True
+    color_format: ClassVar[str] = "RGB"
+    # Skip the startup warmup: the generic dummy request carries a blank image
+    # with no camera poses and no action string, which SANA-WM's request
+    # validation always rejects. The first real request does the warming.
+    dummy_run_num_frames: ClassVar[int] = 0
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
+
+    def __init__(self, *, od_config: OmniDiffusionConfig | None = None, prefix: str = "") -> None:
+        super().__init__()
+        self.od_config = od_config
+        self.prefix = prefix
+
+        self.sana_wm_config = SanaWmConfig()
+        self.quant_config = getattr(od_config, "quantization_config", None) if od_config is not None else None
+
+        self.tokenizer: Any | None = None
+        self.text_encoder: nn.Module | None = None
+        self.vae: nn.Module | None = None
+        self.release_paths: SanaWmLocalPaths | None = None
+        self.weights_sources = []
+        if od_config is not None and od_config.model is not None:
+            # Resolve the checkpoint once at construction (the standard
+            # startup-time flow, like every other diffusion pipeline) so the
+            # transformer below is built with the transformer/config.json shape
+            # rather than defaults.
+            self.resolve_checkpoint()
+            self.weights_sources = [
+                DiffusersPipelineLoader.ComponentSource(
+                    model_or_path=od_config.model,
+                    subfolder=SANA_WM_STAGE1_DIT_SUBFOLDER,
+                    revision=od_config.revision,
+                    # The subfolder file carries module-local keys (standard
+                    # diffusers layout); prefix them into the pipeline
+                    # namespace for the loader's strict-coverage check.
+                    prefix="transformer.",
+                    fall_back_to_pt=False,
+                    allow_patterns_overrides=[SANA_WM_STAGE1_DIT_BASENAME],
+                )
+            ]
+        # Built eagerly: the loader constructs the pipeline under the target
+        # device / default-dtype context, and the weight loader then streams
+        # checkpoint tensors directly into these modules.
+        self.transformer = SanaWmTransformer3DModel(
+            config=self.sana_wm_config,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.transformer" if prefix else "transformer",
+        )
+        self.device = get_local_device()
+        if od_config is not None and od_config.model is not None:
+            self._build_aux_components()
+
+    def resolve_checkpoint(self) -> SanaWmLocalPaths:
+        if self.release_paths is not None:
+            return self.release_paths
+        if self.od_config is None or self.od_config.model is None:
+            raise ValueError("Sana-WM checkpoint resolution requires od_config.model.")
+        self.release_paths = resolve_or_download_sana_wm_checkpoint(
+            self.od_config.model,
+            revision=self.od_config.revision,
+        )
+        self.sana_wm_config = SanaWmConfig.from_json(self.release_paths.config)
+        return self.release_paths
+
+    @staticmethod
+    def _extra_args(sampling_params: Any | None) -> dict[str, Any]:
+        extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
+        return dict(extra_args or {})
+
+    def _native_params(self, payload: dict[str, Any], sampling_params: Any | None) -> SanaWmNativeParams:
+        extra_args = self._extra_args(sampling_params)
+        height = int(getattr(sampling_params, "height", None) or payload["height"])
+        width = int(getattr(sampling_params, "width", None) or payload["width"])
+        num_frames = int(getattr(sampling_params, "num_frames", None) or payload["num_frames"])
+        steps = int(
+            getattr(sampling_params, "num_inference_steps", None)
+            or extra_args.get("num_inference_steps")
+            or SANA_WM_DEFAULT_NUM_INFERENCE_STEPS
+        )
+        seed = int(getattr(sampling_params, "seed", None) or extra_args.get("seed", 0))
+        if min(height, width, num_frames, steps) <= 0:
+            raise ValueError("Sana-WM native height, width, num_frames, and steps must be positive.")
+        cfg_scale = SANA_WM_DEFAULT_GUIDANCE_SCALE
+        if sampling_params is not None and getattr(sampling_params, "guidance_scale_provided", False):
+            cfg_scale = float(getattr(sampling_params, "guidance_scale", 1.0) or 1.0)
+        return SanaWmNativeParams(
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=steps,
+            seed=seed,
+            cfg_scale=cfg_scale,
+        )
+
+    def _native_dtype(self, device: torch.device) -> torch.dtype:
+        dtype = getattr(self.od_config, "dtype", None) if self.od_config is not None else None
+        if isinstance(dtype, torch.dtype):
+            return dtype
+        return torch.float32 if device.type == "cpu" else torch.bfloat16
+
+    def _runtime_device_dtype(self) -> tuple[torch.device, torch.dtype]:
+        return self.device, self._native_dtype(self.device)
+
+    def _build_aux_components(self) -> None:
+        """Construct the text encoder and VAE at startup, on CPU.
+
+        ``ModuleDiscovery`` skips attributes that are ``None``, so components
+        created lazily on the first request are invisible to sequential
+        offload, HSDP auxiliary placement, and the shared device/memory
+        lifecycle. Building them here and deferring placement to the common
+        logic mirrors LTX-2's ``initialize_pipeline_components``.
+        """
+        dtype = self._native_dtype(self.device)
+        self._load_stage1_text_encoder(dtype=dtype)
+        self._load_vae(dtype=dtype)
+        self._place_aux_components()
+
+    def _place_aux_components(self) -> None:
+        """Move the discovered auxiliary components onto the local device.
+
+        Skipped when offload or HSDP owns placement — those paths expect the
+        components to still be on CPU when they take over.
+        """
+        parallel_config = getattr(self.od_config, "parallel_config", None)
+        if (
+            getattr(self.od_config, "enable_cpu_offload", False)
+            or getattr(self.od_config, "enable_layerwise_offload", False)
+            or getattr(parallel_config, "use_hsdp", False)
+        ):
+            return
+        modules = ModuleDiscovery.discover(self)
+        for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):
+            module.to(self.device)
+
+    @staticmethod
+    def _component_local_files_only(model_id: Any) -> bool:
+        """Decide download policy per component, from its own model id.
+
+        The Stage-1 release directory is local once ``resolve_checkpoint`` has
+        run, but the Gemma text encoder lives in a separate Hub repo. Deriving
+        locality from the release paths forced ``local_files_only=True`` onto
+        Gemma as well, which forbids the download on a box that has not cached
+        it yet.
+        """
+        return Path(str(model_id)).expanduser().exists()
+
+    def _load_stage1_text_encoder(self, *, dtype: torch.dtype) -> None:
+        from transformers import AutoModel, AutoTokenizer
+
+        model_ids = [os.environ.get(SANA_WM_STAGE1_TEXT_ENCODER_ENV, "").strip()]
+        model_ids.extend([SANA_WM_STAGE1_TEXT_ENCODER_ID, SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID])
+        errors: list[str] = []
+        for model_id in dict.fromkeys(model_id for model_id in model_ids if model_id):
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    model_id,
+                    local_files_only=self._component_local_files_only(model_id),
+                )
+                text_encoder_model_id = model_id
+                break
+            except OSError as exc:
+                errors.append(f"{model_id}: {exc}")
+        else:
+            raise OSError("Could not load SANA-WM Stage-1 text encoder. Tried: " + "; ".join(errors))
+        if getattr(self.tokenizer, "pad_token", None) is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "right"
+        with torch.device("cpu"):
+            # AutoModel, not AutoModelForCausalLM: only the final hidden state is
+            # used, so the LM head would compute a [seq, 256k] logits tensor
+            # that is thrown away on every prompt encode.
+            self.text_encoder = AutoModel.from_pretrained(
+                text_encoder_model_id,
+                torch_dtype=dtype,
+                local_files_only=self._component_local_files_only(text_encoder_model_id),
+            )
+
+    def _ensure_stage1_text_encoder(self, *, device: torch.device, dtype: torch.dtype) -> None:
+        if self.text_encoder is None:
+            self._load_stage1_text_encoder(dtype=dtype)
+        self.text_encoder.to(device=device, dtype=dtype)
+
+    def _load_vae(self, *, dtype: torch.dtype) -> None:
+        if self.release_paths is None and self.od_config is not None and self.od_config.model is not None:
+            self.resolve_checkpoint()
+        if self.release_paths is None:
+            raise ValueError("Sana-WM VAE loading requires a resolved checkpoint.")
+
+        from diffusers import AutoencoderKLLTX2Video
+
+        with torch.device("cpu"):
+            self.vae = AutoencoderKLLTX2Video.from_pretrained(
+                str(self.release_paths.root),
+                subfolder="vae",
+                torch_dtype=dtype,
+                local_files_only=self._component_local_files_only(self.release_paths.root),
+            )
+        self.vae.use_framewise_encoding = True
+        self.vae.use_framewise_decoding = True
+        # NVlabs reads these from its YAML VAE config with 64/96 defaults; the
+        # diffusers VAE config.json only ships scaling_factor, so the defaults
+        # are load-bearing (the keys are genuinely absent here).
+        self.vae.tile_sample_stride_num_frames = int(getattr(self.vae.config, "tile_sample_stride_num_frames", 64))
+        self.vae.tile_sample_min_num_frames = int(getattr(self.vae.config, "tile_sample_min_num_frames", 96))
+
+    def _ensure_vae(self, *, device: torch.device, dtype: torch.dtype) -> None:
+        if self.vae is None:
+            self._load_vae(dtype=dtype)
+        # Tiling is re-asserted here, not in _load_vae: registry.py assigns
+        # ``model.vae.use_tiling = od_config.vae_use_tiling`` after the pipeline
+        # is constructed, and that flag defaults to False. Matching NVlabs'
+        # inference_video_scripts/inference_sana_wm.py, the SANA VAE has to
+        # tile — a 161-frame decode costs about 9 GiB more without it, and a
+        # 321-frame request (41 latent frames) decodes as one 3D volume and
+        # OOMs. Overriding is deliberate, but say so rather than dropping the
+        # caller's setting silently.
+        if not self.vae.use_tiling:
+            logger.warning("Sana-WM forces VAE tiling on; ignoring vae_use_tiling=False (the VAE OOMs without it).")
+        self.vae.use_tiling = True
+        self.vae.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _preprocess_first_frame(
+        image: Any,
+        *,
+        height: int,
+        width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Return the first-frame image as ``[1, 3, 1, H, W]`` in ``[-1, 1]``.
+
+        ``VideoProcessor.preprocess`` covers the PIL / ndarray / tensor cases.
+        ``resample="bicubic"`` keeps the filter PIL's ``Image.resize`` defaulted
+        to; the processor's own default is lanczos, which would move the
+        conditioning latent.
+        """
+        import numpy as np
+        from diffusers.video_processor import VideoProcessor
+
+        # The processor normalises to [-1, 1] assuming float inputs in [0, 1].
+        if isinstance(image, np.ndarray) and image.dtype == np.uint8:
+            image = image.astype(np.float32) / 255.0
+
+        processor = VideoProcessor(vae_scale_factor=SANA_WM_VAE_SPATIAL_COMPRESSION, resample="bicubic")
+        frame = processor.preprocess(image, height=height, width=width)
+        return frame.unsqueeze(2).to(device=device, dtype=dtype)
+
+    def _vae_normalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        """LTX-2 VAE per-channel latent normalisation matching NVlabs.
+
+        ``z_norm = (z_raw - latents_mean) * scaling_factor / latents_std``
+
+        The LTX-2 VAE (``AutoencoderKLLTX2Video`` in diffusers) ships with
+        per-channel ``latents_mean`` and ``latents_std`` tensors that must be
+        applied alongside ``scaling_factor`` for the round-trip to be
+        identity. Applying `* scaling_factor` alone gives decoded videos a
+        systematic G/B colour shift (pred RGB means
+        ~(110, 220, 208) vs ref ~(94, 114, 138)).
+
+        Returns the normalised latent in the same dtype as the input.
+        """
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        return normalize_latents(
+            latent,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+            float(self.vae.config.scaling_factor),
+        )
+
+    def _vae_denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        """Inverse of :meth:`_vae_normalize_latent` — match NVlabs decode."""
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        return denormalize_latents(
+            latent,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+            float(self.vae.config.scaling_factor),
+        )
+
+    def _vae_encode_first_frame(
+        self,
+        image: Any,
+        *,
+        height: int,
+        width: int,
+        latent_height: int,
+        latent_width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Encode the first-frame image to latent space via the LTX-2 VAE.
+
+        Returns shape ``[1, 128, 1, latent_height, latent_width]``.
+        """
+        self._ensure_vae(device=device, dtype=dtype)
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        frame = self._preprocess_first_frame(image, height=height, width=width, device=device, dtype=dtype)
+        with torch.inference_mode():
+            encoded = self.vae.encode(frame)
+            first_latent = encoded.latent_dist.mean
+            first_latent = self._vae_normalize_latent(first_latent.to(dtype=dtype))
+        # The encoder's own rounding can land a pixel off the latent grid the
+        # noise tensor was built on, and the two are concatenated below.
+        if first_latent.shape[-2:] != (latent_height, latent_width):
+            first_latent = F.interpolate(
+                first_latent.squeeze(2),
+                size=(latent_height, latent_width),
+                mode="bilinear",
+                align_corners=False,
+            ).unsqueeze(2)
+        return first_latent  # [1, 128, 1, lh, lw]
+
+    def _stage1_prompt_text(self, prompt: dict[str, Any]) -> str:
+        user_prompt = str(prompt.get("prompt") or "")
+        chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part)
+        if chi_prompt:
+            return chi_prompt + user_prompt
+        return user_prompt
+
+    def _native_prompt_embeds(
+        self,
+        prompt: dict[str, Any],
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        prompt_text = self._stage1_prompt_text(prompt)
+        self._ensure_stage1_text_encoder(device=device, dtype=dtype)
+        tokenizer = self.tokenizer
+        if tokenizer is None or self.text_encoder is None:
+            raise RuntimeError("Sana-WM Stage-1 text encoder did not initialize.")
+        chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part)
+        model_max_length = int(self.sana_wm_config.model_max_length)
+        max_length_all = model_max_length
+        if chi_prompt:
+            # NVlabs encodes the chi-prompt-prefixed text with extra room for
+            # the chi tokens, then keeps BOS plus the final model window.
+            max_length_all = len(tokenizer.encode(chi_prompt)) + model_max_length - 2
+        encoded = tokenizer(
+            [prompt_text],
+            padding="max_length",
+            max_length=max_length_all,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        ).to(device)
+        # last_hidden_state is the same tensor the CausalLM wrapper exposed as
+        # hidden_states[-1] (both are post-final-norm), without retaining every
+        # intermediate layer.
+        hidden_states = self.text_encoder(**encoded).last_hidden_state
+        attention_mask = encoded.attention_mask
+        if chi_prompt:
+            select_index = [0] + list(range(-model_max_length + 1, 0))
+            hidden_states = hidden_states[:, select_index]
+            attention_mask = attention_mask[:, select_index]
+        if hidden_states.shape[-1] != SANA_WM_STAGE1_PROMPT_CHANNELS:
+            raise ValueError(
+                "Sana-WM Stage-1 Gemma hidden size mismatch: expected "
+                f"{SANA_WM_STAGE1_PROMPT_CHANNELS}, got {hidden_states.shape[-1]}."
+            )
+        # Return RAW Gemma hidden states: the transformer's ``attention_y_norm``
+        # RMSNorm does the normalisation, so normalising here double-normalises.
+        prompt_attention_mask = attention_mask.to(device=device, dtype=torch.float32)
+        # An all-ones mask carries no information. Keep it ``None`` so
+        # cross-attention takes the mask-free attention path instead of handing
+        # every block, on every step, a mask the backend must reshape to no
+        # effect. Evaluated once per prompt encode, so the device sync is paid
+        # twice per request (cond + uncond), not once per denoising step.
+        if bool(prompt_attention_mask.all()):
+            prompt_attention_mask = None
+        self._last_prompt_attention_mask = prompt_attention_mask
+        return hidden_states.to(device=device, dtype=dtype)
+
+    def predict_noise(self, **kwargs: Any) -> torch.Tensor:
+        """Single transformer forward.
+
+        Overrides the mixin default, which assumes a tuple-returning transformer
+        and takes ``result[0]``; SANA-WM's returns the noise prediction directly.
+        """
+        return self.transformer(**kwargs)
+
+    def _decode_native_latents(
+        self,
+        latents: torch.Tensor,
+        *,
+        output_type: str,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Any:
+        if output_type == "latent":
+            return latents
+        self._ensure_vae(device=device, dtype=dtype)
+        if self.vae is None:
+            raise RuntimeError("Sana-WM VAE did not initialize.")
+        # NVlabs always passes ``temb=None`` to ``vae.decode`` for this VAE
+        # (``builder.py::vae_decode -> LTX2VAE_diffusers``); match that.
+        denorm = self._vae_denormalize_latent(latents.to(getattr(self.vae, "dtype", dtype)))
+        video = self.vae.decode(denorm, temb=None, return_dict=False)[0]
+        # Post-process to the requested output_type (np/pil/pt). Let failures
+        # surface: silently returning the raw decoder tensor would hand the
+        # caller a wrong-format/range video with no error.
+        from diffusers.video_processor import VideoProcessor
+
+        processor = VideoProcessor(vae_scale_factor=getattr(self.vae, "spatial_compression_ratio", 32))
+        return processor.postprocess_video(video, output_type=output_type)
+
+    def _run_native_backend(
+        self,
+        *,
+        prompt: dict[str, Any],
+        payload: dict[str, Any],
+        sampling_params: Any | None,
+    ) -> DiffusionOutput:
+        params = self._native_params(payload, sampling_params)
+        latent_frames, latent_height, latent_width = resolve_video_latent_shape(
+            params.height,
+            params.width,
+            params.num_frames,
+            vae_spatial_compression_ratio=SANA_WM_VAE_SPATIAL_COMPRESSION,
+            vae_temporal_compression_ratio=SANA_WM_VAE_TEMPORAL_COMPRESSION,
+        )
+        token_count = latent_frames * latent_height * latent_width
+        extra_args = self._extra_args(sampling_params)
+        max_tokens = int(extra_args.get("sana_wm_native_max_tokens", SANA_WM_NATIVE_MAX_TOKENS))
+        if token_count > max_tokens:
+            raise ValueError(
+                "Sana-WM native latent token count exceeds the configured cap. "
+                f"Requested latent tokens={token_count}, max={max_tokens}. Request a smaller "
+                "`height`/`width`/`num_frames`, or raise the cap via "
+                "`sana_wm_native_max_tokens`."
+            )
+
+        device, dtype = self._runtime_device_dtype()
+        generator = torch.Generator(device=device)
+        generator.manual_seed(params.seed)
+        noise = torch.randn(
+            (1, 128, latent_frames, latent_height, latent_width),
+            device=device,
+            dtype=dtype,
+            generator=generator,
+        )
+
+        # Shared flow-matching Euler scheduler. NVlabs ``LTXFlowEuler`` drives
+        # diffusers' ``FlowMatchEulerDiscreteScheduler``; this is vLLM-Omni's
+        # vendored copy of it, including the ``per_token_timesteps`` step branch.
+        # The schedule applies the shift twice (once via the shifted ``sigma_min``
+        # persisted by ``__init__``, once in ``set_timesteps``), which is what
+        # reproduces the NVlabs timestep table (e.g. shift=9.8, N=3 ->
+        # [1000, 909.0, 87.7]).
+        scheduler = FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=SANA_WM_NUM_TRAIN_TIMESTEPS,
+            shift=self.sana_wm_config.inference_flow_shift,
+        )
+        scheduler.set_timesteps(params.num_inference_steps, device=device)
+        timesteps = scheduler.timesteps
+
+        # ``normalize_sana_wm_payload`` guarantees the key is present but says
+        # nothing about its type. Falling back to pure noise on an unusable
+        # image would silently drop the conditioning and produce a
+        # plausible-looking but wrong video, so this raises instead.
+        import numpy as _np
+
+        first_frame_image = (prompt.get("multi_modal_data") or {}).get("image")
+        if not (hasattr(first_frame_image, "convert") or isinstance(first_frame_image, (_np.ndarray, torch.Tensor))):
+            raise TypeError(
+                "Sana-WM first-frame image must be a PIL Image, numpy ndarray, or "
+                f"torch.Tensor; got {type(first_frame_image).__name__}."
+            )
+        # Per-frame timestep contract matching NVlabs ``LTXFlowEuler.sample``:
+        # frame 0 is the clean VAE-encoded conditioning latent and its timestep
+        # is held at 0 while the sampling loop drives the other frames from
+        # noise.
+        first_latent = self._vae_encode_first_frame(
+            first_frame_image,
+            height=params.height,
+            width=params.width,
+            latent_height=latent_height,
+            latent_width=latent_width,
+            device=device,
+            dtype=dtype,
+        )
+        # Place the CLEAN VAE-encoded first frame at frame 0 and rely on
+        # per-frame timesteps + the `condition_mask` torch.where restore (below)
+        # to keep it invariant across denoising steps.
+        latents = torch.cat([first_latent, noise[:, :, 1:]], dim=2)
+
+        prompt_embeds = self._native_prompt_embeds(
+            prompt,
+            device=device,
+            dtype=dtype,
+        )
+        prompt_attention_mask = self._last_prompt_attention_mask
+
+        # Under CFG, encode the negative prompt up front; each denoise step
+        # then runs a two-branch transformer forward. The negative prompt comes
+        # from the request (``prompt["negative_prompt"]``, set by
+        # ``serving_video``); absent, it defaults to the empty string matching
+        # NVlabs ``GenerationParams.negative_prompt = ""``.
+        do_cfg = params.cfg_scale > 1.0
+        if do_cfg:
+            negative_prompt_obj = {"prompt": str(prompt.get("negative_prompt") or "")}
+            negative_prompt_embeds = self._native_prompt_embeds(
+                negative_prompt_obj,
+                device=device,
+                dtype=dtype,
+            )
+            negative_prompt_attention_mask = self._last_prompt_attention_mask
+            # Restore the cond mask so later readers of
+            # `_last_prompt_attention_mask` see the positive branch.
+            self._last_prompt_attention_mask = prompt_attention_mask
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_attention_mask = None
+
+        camera = payload.get("camera") or {}
+        condition = SanaWmCameraCondition(
+            poses=camera.get("poses") if isinstance(camera, dict) else None,
+            intrinsics=payload.get("intrinsics"),
+            action=payload.get("action"),
+            num_frames=params.num_frames,
+            height=params.height,
+            width=params.width,
+            translation_speed=float(payload.get("translation_speed", 0.05)),
+            rotation_speed_deg=float(payload.get("rotation_speed_deg", 1.2)),
+        )
+        camera_tensors = build_plucker_condition(condition)
+        plucker = camera_tensors["chunk_plucker"].to(device=device, dtype=dtype)
+        raymap = camera_tensors["raymap"].to(device=device, dtype=dtype)
+        spatial_raymap = camera_tensors.get("spatial_raymap")
+        if spatial_raymap is not None:
+            spatial_raymap = spatial_raymap.to(device=device, dtype=dtype)
+
+        # Build a per-frame condition mask matching NVlabs ``LTXFlowEuler``:
+        # frame 0 is the conditioning frame at sigma=0 (preserved after
+        # each ``scheduler.step``), all other frames carry the current
+        # sampling sigma. We omit the legacy `add_noise_to_image_conditioning_latents`
+        # motion-continuity term because the public SANA-WM config sets
+        # ``condition_frame_info={0: 0.0}`` (image_cond_noise_scale=0).
+        cam_batch = latents.shape[0]
+        cam_frames = latents.shape[2]
+        condition_mask = torch.zeros(cam_batch, 1, cam_frames, 1, 1, device=latents.device, dtype=latents.dtype)
+        condition_mask[:, :, 0] = 1.0
+
+        for timestep in timesteps:
+            # (B, 1, F) per-frame timestep, frame 0 forced to 0. Stays fp32:
+            # casting through the bf16 latent dtype would quantise it before
+            # ``SanaWmTimestepEmbedder``'s sinusoidal embed.
+            model_timestep = timestep.float().expand(cam_batch, 1, cam_frames).clone()
+            model_timestep[:, :, 0] = 0.0
+            positive_kwargs = {
+                "hidden_states": latents,
+                "timestep": model_timestep,
+                "encoder_hidden_states": prompt_embeds,
+                "encoder_attention_mask": prompt_attention_mask,
+                "plucker": plucker,
+                "raymap": raymap,
+                "spatial_raymap": spatial_raymap,
+            }
+            negative_kwargs = (
+                {
+                    **positive_kwargs,
+                    "encoder_hidden_states": negative_prompt_embeds,
+                    "encoder_attention_mask": negative_prompt_attention_mask,
+                }
+                if do_cfg
+                else None
+            )
+            # ``cfg_normalize=False`` keeps the plain NVlabs combine,
+            # ``uncond + scale * (cond - uncond)``. Under cfg_parallel_size > 1
+            # the mixin runs one branch per rank and all-gathers instead of
+            # evaluating both branches serially.
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_cfg,
+                params.cfg_scale,
+                positive_kwargs,
+                negative_kwargs,
+                cfg_normalize=False,
+            )
+
+            # The per-frame contract pairs with a per-token flow-matching Euler
+            # step. Both step branches end at ``prev = sample + dt *
+            # model_output``, but the scheduler defines dt as
+            # ``sigma - sigma_next`` per-token against ``sigma_next - sigma``
+            # standard, so the per-token call takes the sign-flipped prediction
+            # (matching NVlabs ``LTXFlowEuler.sample``); passing the standard
+            # sign here would integrate towards more noise.
+            #
+            # Broadcast the (B, 1, F) timestep to (B, 1, F, H, W) and flatten
+            # F*H*W; frame-0 conditioning tokens are already 0. The flatten
+            # order (F, H, W) matches ``pack_latents``.
+            pt_t = (
+                model_timestep.unsqueeze(-1)
+                .unsqueeze(-1)
+                .expand(cam_batch, 1, cam_frames, latents.shape[3], latents.shape[4])
+                .reshape(cam_batch, -1)
+            )
+            # The per-token branch consumes/returns the packed (B, N, C) layout
+            # and keeps its result in fp32; cast back so the latent dtype stays
+            # stable across steps.
+            stepped_packed = scheduler.step(
+                -pack_latents(noise_pred),
+                timestep,
+                pack_latents(latents),
+                per_token_timesteps=pt_t,
+                return_dict=False,
+            )[0]
+            stepped = unpack_latents(
+                stepped_packed.to(latents.dtype),
+                cam_frames,
+                latents.shape[3],
+                latents.shape[4],
+            )
+
+            # Match NVlabs LTXFlowEuler exactly. This is stricter than a plain
+            # condition-frame restore: with bf16 latents, the `t=1000`
+            # comparison rounds `1 - 1e-6` to `1`, so the first full-noise step
+            # is discarded for generated tokens as well.
+            tokens_to_denoise_mask = timestep / float(SANA_WM_NUM_TRAIN_TIMESTEPS) - 1e-6 < (1.0 - condition_mask)
+            latents = torch.where(tokens_to_denoise_mask, stepped, latents)
+
+        output_type = getattr(sampling_params, "output_type", None) or "np"
+        output = self._decode_native_latents(latents, output_type=output_type, device=device, dtype=dtype)
+        return DiffusionOutput(
+            output=build_sana_wm_output_envelope(
+                output=output,
+                output_type=output_type,
+                metadata={
+                    "backend": "native_gdn",
+                    "output_space": output_type,
+                    "chi_prompt_applied": bool(self.sana_wm_config.chi_prompt),
+                    "num_frames": params.num_frames,
+                    "height": params.height,
+                    "width": params.width,
+                    "latent_tokens": token_count,
+                    "sampling_steps": params.num_inference_steps,
+                },
+            ),
+        )
+
+    def forward(self, req: DiffusionRequestBatch, *args: Any, **kwargs: Any) -> DiffusionOutput:
+        del args, kwargs
+        if len(req.prompts) != 1:
+            raise ValueError("Sana-WM native backend currently supports exactly one prompt per request.")
+
+        prompt = req.prompts[0]
+        if isinstance(prompt, str):
+            raise ValueError("Sana-WM requires a mapping prompt with first-frame image and camera/action metadata.")
+        # get_sana_wm_pre_process_func is the single canonical entry point: it
+        # normalises the request and republishes the geometry onto
+        # sampling_params. normalize_sana_wm_payload is idempotent, so this only
+        # covers offline callers that bypass the preprocess hook.
+        prompt = normalize_sana_wm_payload(prompt)
+        payload = prompt["additional_information"]["sana_wm"]
+
+        # The checkpoint is resolved once in __init__; this is a cached no-op
+        # kept as a guard for pipelines constructed without od_config.model.
+        if self.od_config is not None and self.od_config.model is not None:
+            self.resolve_checkpoint()
+
+        return self._run_native_backend(
+            prompt=prompt,
+            payload=payload,
+            sampling_params=req.sampling_params,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, Any]]) -> set[str]:
+        # Keys arrive pipeline-prefixed (``transformer.*``, per the component
+        # source). AutoWeightsLoader streams them to the submodule and returns
+        # the prefixed names the loader's strict-coverage check compares.
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/sana_wm/request.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Request normalization for Sana-WM image-to-video.
+
+The released SANA-WM checkpoint is first-frame image-to-video. This module keeps
+request-side normalization lightweight: it validates the image/camera/action
+contract and stores a canonical payload under
+``additional_information["sana_wm"]``. Model-internal raymap / Plucker
+projection lives with the transformer.
+
+This lives beside the model rather than under ``stage_input_processors``
+because nothing loads it as one: the pipeline's ``pre_process_func`` is
+registered from ``pipeline_sana_wm``, the same as every other diffusion model,
+and no deploy config wires a ``custom_process_input_func`` here. Keeping it in
+the stage package forced the pipeline to import *upwards* — the only such
+import in ``diffusion/models/`` — which in turn made the VAE compression
+constants impossible to share and so duplicated.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from vllm_omni.diffusion.models.sana_wm.camera_control import action_rollout_num_frames
+from vllm_omni.diffusion.models.sana_wm.config import (
+    SANA_WM_VAE_SPATIAL_COMPRESSION,
+    SANA_WM_VAE_TEMPORAL_COMPRESSION,
+)
+
+SANA_WM_CANONICAL_KEY = "sana_wm"
+# Request fallbacks used whenever the field is absent, mirroring the deploy YAML
+# (161 frames) and the model's native resolution. Both are genuinely reachable:
+# a served request that omits num_frames never gets one injected, and callers
+# may ask for a different resolution.
+SANA_WM_DEFAULT_NUM_FRAMES = 161
+SANA_WM_DEFAULT_HEIGHT = 704
+SANA_WM_DEFAULT_WIDTH = 1280
+SANA_WM_DEFAULT_CAMERA_FORMAT = "c2w_4x4"
+SANA_WM_DEFAULT_COORDINATE_SYSTEM = "official"
+
+
+def _unwrap_single(value: Any) -> Any:
+    if isinstance(value, list):
+        return value[0] if value else None
+    return value
+
+
+def _as_dict(value: Any, *, name: str) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if not isinstance(value, Mapping):
+        raise ValueError(f"Sana-WM {name} must be a mapping, got {type(value).__name__}.")
+    return dict(value)
+
+
+def _is_sequence(value: Any) -> bool:
+    return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray))
+
+
+def _as_finite_float(value: Any, *, name: str) -> float:
+    """Parse a float and reject NaN/inf.
+
+    Non-finite geometry does not fail loudly downstream: it propagates through
+    the raymap into the Plucker maps and comes back as an all-NaN video with a
+    200 response.
+    """
+    try:
+        parsed = float(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Sana-WM {name} must be a number, got {value!r}.") from exc
+    if not math.isfinite(parsed):
+        raise ValueError(f"Sana-WM {name} must be finite, got {parsed!r}.")
+    return parsed
+
+
+def _validate_numeric_matrix4x4(matrix: Any, *, name: str) -> None:
+    if not _is_sequence(matrix) or len(matrix) != 4:
+        raise ValueError(f"Sana-WM {name} must be a 4x4 matrix.")
+    for row_idx, row in enumerate(matrix):
+        if not _is_sequence(row) or len(row) != 4:
+            raise ValueError(f"Sana-WM {name}[{row_idx}] must contain 4 values.")
+        for col_idx, value in enumerate(row):
+            _as_finite_float(value, name=f"{name}[{row_idx}][{col_idx}]")
+
+
+def _validate_camera_poses(poses: Any, *, num_frames: int | None = None) -> None:
+    if not _is_sequence(poses) or len(poses) == 0:
+        raise ValueError("Sana-WM camera poses must be a non-empty sequence of 4x4 matrices.")
+    if num_frames is not None and len(poses) != num_frames:
+        raise ValueError(f"Sana-WM camera poses length {len(poses)} must equal num_frames {num_frames}.")
+    for idx, pose in enumerate(poses):
+        _validate_numeric_matrix4x4(pose, name=f"camera.poses[{idx}]")
+
+
+def _validate_intrinsics(intrinsics: Any) -> None:
+    # Only the ``{fx, fy, cx, cy}`` mapping is accepted (the one form the recipe
+    # and examples use). Camera intrinsics are optional; when omitted the model
+    # derives them from the output resolution.
+    if intrinsics is None:
+        return
+    if not isinstance(intrinsics, Mapping):
+        raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.")
+    required = {"fx", "fy", "cx", "cy"}
+    missing = sorted(required - set(intrinsics))
+    if missing:
+        raise ValueError(f"Sana-WM intrinsics mapping is missing keys: {missing}.")
+    for key in required:
+        value = _as_finite_float(intrinsics[key], name=f"intrinsics[{key!r}]")
+        # ``compute_raymap`` divides by fx/fy. cx/cy stay unconstrained beyond
+        # finiteness — a principal point outside the frame is unusual, not wrong.
+        if key in ("fx", "fy") and value <= 0:
+            raise ValueError(f"Sana-WM intrinsics[{key!r}] must be positive, got {value}.")
+
+
+def _as_positive_int(value: Any, *, name: str) -> int:
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Sana-WM {name} must be an integer, got {value!r}.") from exc
+    if parsed <= 0:
+        raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.")
+    return parsed
+
+
+def _as_positive_float(value: Any, *, name: str) -> float:
+    # Finiteness first: every comparison against NaN is False, so a bare
+    # ``parsed <= 0`` would accept it.
+    parsed = _as_finite_float(value, name=name)
+    if parsed <= 0:
+        raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.")
+    return parsed
+
+
+def _extract_image(mm_data: Mapping[str, Any]) -> Any:
+    return _unwrap_single(mm_data.get("image"))
+
+
+def _first_present(*values: Any) -> Any:
+    for value in values:
+        if value is not None:
+            return value
+    return None
+
+
+def _validate_latent_geometry(*, num_frames: int, height: int, width: int) -> None:
+    """Reject geometry the VAE would silently floor.
+
+    The SANA-WM VAE compresses 32x spatially and 8x temporally, so a request
+    that is not aligned loses pixels/frames instead of failing.
+    """
+    if height % SANA_WM_VAE_SPATIAL_COMPRESSION or width % SANA_WM_VAE_SPATIAL_COMPRESSION:
+        raise ValueError(
+            f"Sana-WM height and width must be divisible by {SANA_WM_VAE_SPATIAL_COMPRESSION}, "
+            f"got {height} and {width}."
+        )
+    if (num_frames - 1) % SANA_WM_VAE_TEMPORAL_COMPRESSION:
+        raise ValueError(
+            f"Sana-WM num_frames must satisfy (num_frames - 1) % {SANA_WM_VAE_TEMPORAL_COMPRESSION} == 0, "
+            f"got {num_frames}."
+        )
+
+
+def normalize_sana_wm_payload(prompt: Mapping[str, Any]) -> dict[str, Any]:
+    """Return a prompt copy with canonical Sana-WM request metadata.
+
+    The ``sana_wm`` block is read from the top-level ``sana_wm`` key, falling
+    back to ``additional_information["sana_wm"]`` so the function is idempotent
+    when called again on its own output. The first-frame image is read from
+    ``multi_modal_data["image"]``.
+    """
+
+    if not isinstance(prompt, Mapping):
+        raise ValueError(f"Sana-WM prompt must be a mapping, got {type(prompt).__name__}.")
+
+    result = dict(prompt)
+    mm_data = _as_dict(result.get("multi_modal_data"), name="multi_modal_data")
+    image = _extract_image(mm_data)
+    if image is None:
+        raise ValueError("Sana-WM requires a first-frame image in `multi_modal_data['image']`.")
+
+    additional = _as_dict(result.get("additional_information"), name="additional_information")
+    raw = _first_present(
+        result.get(SANA_WM_CANONICAL_KEY),
+        additional.get(SANA_WM_CANONICAL_KEY),
+        {},
+    )
+    raw = _as_dict(raw, name=SANA_WM_CANONICAL_KEY)
+
+    num_frames = _as_positive_int(
+        _first_present(raw.get("num_frames"), result.get("num_frames"), SANA_WM_DEFAULT_NUM_FRAMES),
+        name="num_frames",
+    )
+    height = _as_positive_int(
+        _first_present(raw.get("height"), result.get("height"), SANA_WM_DEFAULT_HEIGHT),
+        name="height",
+    )
+    width = _as_positive_int(_first_present(raw.get("width"), result.get("width"), SANA_WM_DEFAULT_WIDTH), name="width")
+    _validate_latent_geometry(num_frames=num_frames, height=height, width=width)
+
+    camera = raw.get("camera")
+    camera = _as_dict(camera, name="camera") if camera is not None else None
+    action = raw.get("action")
+    if action is not None:
+        if not isinstance(action, str) or not action.strip():
+            raise ValueError("Sana-WM action must be a non-empty string when provided.")
+        # Rejected here rather than in the model: a mismatch used to shorten the
+        # camera tensor silently and then fail deep in the transformer on a
+        # camera/latent frame mismatch, and an unbounded duration used to be
+        # expanded one entry per frame before any of that.
+        rollout_frames = action_rollout_num_frames(action)
+        if rollout_frames != num_frames:
+            raise ValueError(
+                f"Sana-WM action rolls out {rollout_frames} frames but num_frames is {num_frames}. "
+                "Action segment durations must sum to num_frames - 1."
+            )
+    if camera is not None and action is not None:
+        raise ValueError("Sana-WM accepts exactly one of `camera` or `action`, not both.")
+    if camera is None and action is None:
+        raise ValueError("Sana-WM requires either explicit camera poses or an action DSL string.")
+
+    canonical_camera = None
+    if camera is not None:
+        poses = camera.get("poses")
+        if poses is None:
+            raise ValueError("Sana-WM camera payload requires `poses`.")
+        _validate_camera_poses(poses, num_frames=num_frames)
+        # Only the defaults are accepted: nothing downstream reads these two
+        # fields, so honouring another value would mean interpreting w2c poses
+        # as c2w and silently generating the wrong trajectory.
+        camera_format = camera.get("format", SANA_WM_DEFAULT_CAMERA_FORMAT)
+        if camera_format != SANA_WM_DEFAULT_CAMERA_FORMAT:
+            raise ValueError(
+                f"Sana-WM only accepts camera format {SANA_WM_DEFAULT_CAMERA_FORMAT!r}; "
+                f"convert the poses before sending, got {camera_format!r}."
+            )
+        coordinate_system = camera.get("coordinate_system", SANA_WM_DEFAULT_COORDINATE_SYSTEM)
+        if coordinate_system != SANA_WM_DEFAULT_COORDINATE_SYSTEM:
+            raise ValueError(
+                f"Sana-WM only accepts coordinate system {SANA_WM_DEFAULT_COORDINATE_SYSTEM!r}, "
+                f"got {coordinate_system!r}."
+            )
+        canonical_camera = {"poses": poses}
+
+    intrinsics = raw.get("intrinsics")
+    _validate_intrinsics(intrinsics)
+    translation_speed = raw.get("translation_speed")
+    rotation_speed_deg = raw.get("rotation_speed_deg")
+
+    canonical = {
+        "num_frames": num_frames,
+        "height": height,
+        "width": width,
+        "camera": canonical_camera,
+        "action": action,
+        "intrinsics": intrinsics,
+    }
+    if translation_speed is not None:
+        canonical["translation_speed"] = _as_positive_float(translation_speed, name="translation_speed")
+    if rotation_speed_deg is not None:
+        canonical["rotation_speed_deg"] = _as_positive_float(rotation_speed_deg, name="rotation_speed_deg")
+
+    mm_data["image"] = image
+    result["multi_modal_data"] = mm_data
+    additional[SANA_WM_CANONICAL_KEY] = canonical
+    result["additional_information"] = additional
+    result.pop(SANA_WM_CANONICAL_KEY, None)
+    return result

vllm_omni/diffusion/models/sana_wm/ucpe.py

@@ -0,0 +1,536 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SANA-WM UCPE (Unified Camera Pose Embedding) per-block attention transforms.
+
+Ported from NVlabs/Sana ``sana_camctrl_blocks.py`` for the inference-only
+SANA-WM bidirectional 1600M release. Scope is intentionally narrow:
+
+* Pinhole camera only (xi=0); the UCM ``xi`` parameter is fixed.
+* Inference path only — no training-time camera-branch dropout.
+* Online computation only — the precomputed ``cam_pos_embeds`` shortcut
+  used by NVlabs' fused kernels is omitted because vLLM-Omni builds the
+  matrices once per transformer forward and shares them across all blocks.
+* ``apply_vo=True`` only — the SANA-WM checkpoint always uses the
+  inverse-output transform.
+
+Public surface:
+    ``prepare_cam_geometry(...)`` -> :class:`SanaWmCamGeometry`, built once per
+    transformer forward and threaded through every block. It carries both the
+    ``(apply_q, apply_kv, apply_output)`` closures used by the softmax camera
+    branch and the projection/RoPE tensors used by ``cam_prep_func(...)``, the
+    NVlabs fused camera-prep contract on the GDN camera branch.
+
+Each closure transforms a tensor of shape ``(B, num_heads, N, D)`` where
+``N = T * H * W`` (matching the GDN token layout) and ``D == head_dim``,
+returning a tensor of the same shape.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Pinhole ray-grid construction
+# ---------------------------------------------------------------------------
+
+
+def _make_pixel_grid(
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Return ``(H, W, 3)`` grid ``(x, y, 1)`` of pixel coordinates."""
+    xs = torch.arange(width, device=device, dtype=dtype)
+    ys = torch.arange(height, device=device, dtype=dtype)
+    y_grid, x_grid = torch.meshgrid(ys, xs, indexing="ij")
+    return torch.stack((x_grid, y_grid, torch.ones_like(x_grid)), dim=-1)
+
+
+def _unproject_pinhole(
+    fx: torch.Tensor,
+    fy: torch.Tensor,
+    cx: torch.Tensor,
+    cy: torch.Tensor,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Unproject a pixel grid into camera-frame direction vectors.
+
+    Args:
+        fx, fy, cx, cy: shape ``(B, T)`` intrinsics in token-grid units — the
+            caller divides all four by patch size before calling.
+        height, width: token-grid size.
+
+    Returns:
+        ``(B, T, H, W, 3)`` direction tensor (NOT normalised).
+    """
+    B, T = fx.shape
+    device = fx.device
+    dtype = fx.dtype
+    grid = _make_pixel_grid(height, width, device, dtype)  # (H, W, 3)
+    u = grid[..., 0]  # (H, W)
+    v = grid[..., 1]
+    fx_e = fx.view(B, T, 1, 1)
+    fy_e = fy.view(B, T, 1, 1)
+    cx_e = cx.view(B, T, 1, 1)
+    cy_e = cy.view(B, T, 1, 1)
+    x = (u - cx_e) / fx_e
+    y = (v - cy_e) / fy_e
+    z = torch.ones_like(x)
+    # Match NVlabs UCM normalisation at xi=0: alpha = 1, gamma = 1/(1+r^2),
+    # X = gamma*x, Y = gamma*y, Z = gamma. The resulting vector has unit
+    # norm. Computed compactly via direct stack + normalise.
+    d_cam = torch.stack((x, y, z), dim=-1)  # (B, T, H, W, 3)
+    d_cam = d_cam / d_cam.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+    return d_cam
+
+
+def _world_to_ray_mats(
+    d_cam: torch.Tensor,  # (B, T, H, W, 3) unit rays in camera frame
+    c2w: torch.Tensor,  # (B, T, 4, 4) camera-to-world SE(3)
+) -> torch.Tensor:
+    """Build per-pixel ``ray <- world`` SE(3) transforms ``(B, T, H, W, 4, 4)``.
+
+    The ray frame is anchored at the camera origin with the +Z axis along
+    the per-pixel viewing direction and a +Y axis defined by the camera-Y
+    cross product. This matches NVlabs ``world_to_ray_mats``.
+    """
+    B, T, H, W, _ = d_cam.shape
+    device = d_cam.device
+    dtype = d_cam.dtype
+    rot_cam = c2w[..., :3, :3]
+    trans_cam = c2w[..., :3, 3]
+    d_world = torch.einsum("btij,bthwj->bthwi", rot_cam, d_cam)
+    cam_y = rot_cam[..., :, 1].unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3)
+    z_ray = F.normalize(d_world, dim=-1, eps=1e-6)
+    x_ray = F.normalize(torch.cross(cam_y, z_ray, dim=-1), dim=-1, eps=1e-6)
+    y_ray = F.normalize(torch.cross(z_ray, x_ray, dim=-1), dim=-1, eps=1e-6)
+    rot_local_to_world = torch.stack((x_ray, y_ray, z_ray), dim=-1)  # (B,T,H,W,3,3)
+    rot_world_to_local = rot_local_to_world.transpose(-1, -2)
+    t_world = trans_cam.unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3)
+    t_w2l = -torch.einsum("bthwij,bthwj->bthwi", rot_world_to_local, t_world)
+    raymats = torch.zeros(B, T, H, W, 4, 4, device=device, dtype=dtype)
+    raymats[..., :3, :3] = rot_world_to_local
+    raymats[..., :3, 3] = t_w2l
+    raymats[..., 3, 3] = 1.0
+    # Guard NaNs from degenerate pixels (e.g. nan direction at infinity).
+    nan_mask = torch.isnan(d_world).any(-1)
+    if nan_mask.any():
+        eye = torch.eye(4, device=device, dtype=dtype)
+        raymats[nan_mask] = eye
+    return raymats
+
+
+def _process_camera_conditions(
+    camera_conditions: torch.Tensor,  # (B, T, 20)
+    hw: tuple[int, int, int],
+    patch_size: tuple[int, int, int] = (1, 1, 1),
+) -> torch.Tensor:
+    """Convert ``(B, T, 20)`` SANA-WM camera conditions to ``(B, T, H, W, 4, 4)``
+    ``ray <- world`` transforms.
+
+    Layout of the last axis is ``[c2w_flat (16), fx, fy, cx, cy]``.
+    Intrinsics follow the NVlabs UCPE contract: the packed camera payload is
+    expressed on the pre-patch latent image grid, while ``HW`` is the
+    transformer token grid. ``patch_size`` maps those units to the token grid
+    before unprojection. The released SANA-WM P1 config uses ``(1, 1, 1)``.
+
+    Args:
+        camera_conditions: ``(B, T, 20)`` from ``_pack_camera_conditions``.
+        hw: ``(T, H, W)`` token-grid shape. Equal to the latent grid only
+            because the released config patches by ``(1, 1, 1)``.
+    """
+    B, T_cond, last = camera_conditions.shape
+    if last != 20:
+        raise ValueError(f"SANA-WM camera_conditions last dim must be 20, got {last}.")
+    T_latent, H_latent, W_latent = hw
+    if T_cond != T_latent:
+        raise ValueError(f"SANA-WM camera_conditions frames {T_cond} != latent frames {T_latent}.")
+    c2w = camera_conditions[..., :16].reshape(B, T_latent, 4, 4)
+    patch_t, patch_h, patch_w = patch_size
+    del patch_t
+    fx = camera_conditions[..., 16] / float(patch_w)
+    fy = camera_conditions[..., 17] / float(patch_h)
+    cx = camera_conditions[..., 18] / float(patch_w)
+    cy = camera_conditions[..., 19] / float(patch_h)
+    d_cam = _unproject_pinhole(fx, fy, cx, cy, H_latent, W_latent)
+    raymats = _world_to_ray_mats(d_cam, c2w)
+    return raymats
+
+
+def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
+    """Closed-form inverse of a stack of 4x4 SE(3) matrices."""
+    if transforms.shape[-2:] != (4, 4):
+        raise ValueError(f"SE3 inverse expects (..., 4, 4), got {tuple(transforms.shape)}.")
+    rot_inv = transforms[..., :3, :3].transpose(-1, -2)
+    out = torch.zeros_like(transforms)
+    out[..., :3, :3] = rot_inv
+    out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv, transforms[..., :3, 3])
+    out[..., 3, 3] = 1.0
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Per-token transform primitives
+# ---------------------------------------------------------------------------
+
+
+def _apply_ray_projmat(
+    feats: torch.Tensor,  # (B, num_heads, N, D)
+    matrix: torch.Tensor,  # (B, N, 4, 4)
+) -> torch.Tensor:
+    """Apply a per-token 4x4 projection to features grouped by 4 channels.
+
+    ``D`` must be a multiple of 4. Out-of-place; preserves ``feats.shape``.
+    """
+    B, num_heads, N, D = feats.shape
+    if D % 4 != 0:
+        raise ValueError(f"UCPE projmat expects D divisible by 4, got D={D}.")
+    feats_grouped = feats.reshape(B, num_heads, N, D // 4, 4)
+    # matrix: (B, N, 4, 4); feats_grouped: (B, num_heads, N, K, 4)
+    out = torch.einsum("bnij,bhnkj->bhnki", matrix, feats_grouped)
+    return out.reshape(B, num_heads, N, D)
+
+
+def _apply_complex_rope(
+    hidden_states: torch.Tensor,
+    freqs: torch.Tensor,
+    *,
+    inverse: bool = False,
+) -> torch.Tensor:
+    """Apply complex RoPE: treat last-dim pairs as ``(real, imag)`` and multiply
+    by ``freqs``. Computation runs in fp64 to match the NVlabs reference.
+
+    Args:
+        hidden_states: ``(..., D)`` with ``D % 2 == 0``.
+        freqs: complex tensor broadcastable to ``(..., D // 2)``.
+        inverse: when ``True`` apply ``freqs.conj()`` (used for ``apply_o``).
+    """
+    x_real = hidden_states.to(torch.float64).contiguous()
+    x_complex = torch.view_as_complex(x_real.unflatten(-1, (-1, 2)))
+    if inverse:
+        freqs = freqs.conj()
+    x_out = torch.view_as_real(x_complex * freqs).flatten(-2, -1)
+    return x_out.to(hidden_states.dtype)
+
+
+def _apply_block_diagonal(
+    feats: torch.Tensor,
+    func_size_pairs: list[tuple[Callable[[torch.Tensor], torch.Tensor], int]],
+) -> torch.Tensor:
+    """Split ``feats`` along the last axis by ``block_sizes`` and apply one
+    callable per block, then concatenate. Preserves ``feats.shape``.
+    """
+    funcs, block_sizes = zip(*func_size_pairs)
+    if feats.shape[-1] != sum(block_sizes):
+        raise ValueError(f"UCPE block-diagonal block sum {sum(block_sizes)} != feats D {feats.shape[-1]}.")
+    blocks = torch.split(feats, list(block_sizes), dim=-1)
+    return torch.cat([func(block) for func, block in zip(funcs, blocks)], dim=-1)
+
+
+def _slice_rope_for_cam(
+    rotary_emb: torch.Tensor | None,
+    rope_dim: int,
+) -> torch.Tensor | None:
+    """Re-slice WAN-style RoPE frequencies for a smaller ``rope_dim`` using the
+    same ``(T, H, W)`` split as the main branch. Mirrors NVlabs slicing logic.
+
+    The T/H/W band layout is read off the input shape
+    (``rotary_emb.shape[-1] * 2``) rather than a separately supplied head_dim,
+    so the split stays correct without the caller passing it. ``cam_head_dim``
+    still has to equal the main ``head_dim`` — ``SanaWmSelfAttention.__init__``
+    rejects anything else.
+    """
+    if rotary_emb is None:
+        return None
+    # Derive WAN RoPE T/H/W band offsets from the actual input tensor size,
+    # not from a head_dim argument. This is correct regardless of whether
+    # cam_head_dim == main_head_dim.
+    in_head_dim = rotary_emb.shape[-1] * 2
+    orig_t = in_head_dim // 2 - 2 * (in_head_dim // 6)
+    orig_h = in_head_dim // 6
+    new_t = rope_dim // 2 - 2 * (rope_dim // 6)
+    new_h = rope_dim // 6
+    new_w = rope_dim // 6
+    t_part = rotary_emb[..., :new_t]
+    h_part = rotary_emb[..., orig_t : orig_t + new_h]
+    w_part = rotary_emb[..., orig_t + orig_h : orig_t + orig_h + new_w]
+    return torch.cat([t_part, h_part, w_part], dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Apply-fn preparation
+# ---------------------------------------------------------------------------
+
+
+def _identity(x: torch.Tensor) -> torch.Tensor:
+    return x
+
+
+def _prepare_ray_apply_fns(
+    head_dim: int,
+    proj: torch.Tensor,  # (B, N, 4, 4) ray <- world
+    proj_t: torch.Tensor,  # (B, N, 4, 4) ray <- world transpose
+    proj_inv: torch.Tensor,  # (B, N, 4, 4) world <- ray
+    rotary_emb: torch.Tensor | None = None,
+) -> tuple[Callable, Callable, Callable]:
+    """Build ``(apply_q, apply_kv, apply_o)`` block-diagonal callables.
+
+    Each callable transforms a ``(B, num_heads, N, head_dim)`` tensor:
+    the first ``head_dim // 2`` channels go through the per-token 4x4
+    projection, the remaining ``head_dim // 2`` channels go through complex
+    RoPE (or identity when ``rotary_emb`` is None).
+    """
+    if rotary_emb is not None:
+        rope_fwd = partial(_apply_complex_rope, freqs=rotary_emb, inverse=False)
+        rope_inv = partial(_apply_complex_rope, freqs=rotary_emb, inverse=True)
+    else:
+        rope_fwd = _identity
+        rope_inv = _identity
+
+    half = head_dim // 2
+
+    transforms_q = [
+        (partial(_apply_ray_projmat, matrix=proj_t), half),
+        (rope_fwd, half),
+    ]
+    transforms_kv = [
+        (partial(_apply_ray_projmat, matrix=proj_inv), half),
+        (rope_fwd, half),
+    ]
+    transforms_o = [
+        (partial(_apply_ray_projmat, matrix=proj), half),
+        (rope_inv, half),
+    ]
+
+    apply_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q)
+    apply_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv)
+    apply_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o)
+    return apply_q, apply_kv, apply_o
+
+
+@dataclass(frozen=True)
+class SanaWmCamGeometry:
+    """UCPE geometry shared by every block of one transformer forward.
+
+    The ray grid, its SE(3) inverse and the camera RoPE tables depend only on
+    the camera payload, the latent grid, the patch size and the head dim — all
+    fixed for the whole forward — yet each block's camera branch used to
+    rebuild them. At 20 blocks x 60 steps x 2 CFG branches that is ~2400
+    rebuilds of an 18480-token ray grid and its 4x4 inverses per request.
+
+    ``apply_q``/``apply_kv``/``apply_output`` serve the softmax camera branch;
+    ``proj_q``/``proj_kv``/``rope_cos``/``rope_sin`` serve the fused
+    ``cam_prep_func`` contract on the GDN branch. Both are derived from the
+    same projection matrices, so they are built together.
+    """
+
+    proj_q: torch.Tensor
+    proj_kv: torch.Tensor
+    rope_cos: torch.Tensor
+    rope_sin: torch.Tensor
+    apply_q: Callable[[torch.Tensor], torch.Tensor]
+    apply_kv: Callable[[torch.Tensor], torch.Tensor]
+    apply_output: Callable[[torch.Tensor], torch.Tensor]
+
+
+def _prepare_ucpe_rope_tables(
+    rotary_emb_cam: torch.Tensor,
+    token_count: int,
+    half_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert complex RoPE freqs to NVlabs' interleaved cos/sin tables.
+
+    The output convention is ``out[d] = x[d] * cos[d] + x[d ^ 1] * sin[d]``.
+    For every complex pair, ``sin[2i] = -imag(freq[i])`` and
+    ``sin[2i + 1] = imag(freq[i])``.
+    """
+    freqs = rotary_emb_cam
+    while freqs.ndim > 2 and freqs.shape[0] == 1:
+        freqs = freqs.squeeze(0)
+    if freqs.ndim != 2:
+        raise ValueError(f"Sana-WM cam RoPE freqs must be 2D after squeeze, got {tuple(freqs.shape)}.")
+    if freqs.shape[0] != token_count:
+        raise ValueError(f"Sana-WM cam RoPE token count {freqs.shape[0]} != {token_count}.")
+    if freqs.shape[1] * 2 != half_dim:
+        raise ValueError(f"Sana-WM cam RoPE dim {freqs.shape[1] * 2} != half_dim {half_dim}.")
+    cos_half = freqs.real.float()
+    sin_half = freqs.imag.float()
+    rope_cos = cos_half.repeat_interleave(2, dim=-1).contiguous()
+    rope_sin = torch.stack((-sin_half, sin_half), dim=-1).reshape(token_count, half_dim).contiguous()
+    return rope_cos, rope_sin
+
+
+def prepare_cam_geometry(
+    *,
+    camera_conditions: torch.Tensor,
+    spatial_shape: tuple[int, int, int],
+    patch_size: tuple[int, int, int],
+    head_dim: int,
+    rotary_emb: torch.Tensor | None,
+) -> SanaWmCamGeometry:
+    """Build the per-forward UCPE geometry: ray projections, RoPE tables, apply fns.
+
+    Call once per transformer forward and share the result across every block —
+    see :class:`SanaWmCamGeometry` for why.
+    """
+    if head_dim % 4 != 0:
+        raise ValueError(f"UCPE head_dim must be divisible by 4, got {head_dim}.")
+    if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0:
+        raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.")
+
+    batch = camera_conditions.shape[0]
+    frames, height, width = spatial_shape
+    token_count = frames * height * width
+    half_dim = head_dim // 2
+
+    raymats = _process_camera_conditions(camera_conditions, spatial_shape, patch_size=patch_size)
+    proj = raymats.reshape(batch, token_count, 4, 4).contiguous()
+    proj_q = proj.transpose(-1, -2).contiguous()
+    proj_kv = _invert_se3(proj).contiguous()
+
+    rotary_emb_cam = _slice_rope_for_cam(rotary_emb, half_dim)
+    if rotary_emb_cam is None:
+        rope_cos = torch.ones(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32)
+        rope_sin = torch.zeros(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32)
+    else:
+        rope_cos, rope_sin = _prepare_ucpe_rope_tables(rotary_emb_cam, token_count, half_dim)
+
+    apply_q, apply_kv, apply_output = _prepare_ray_apply_fns(head_dim, proj, proj_q, proj_kv, rotary_emb=rotary_emb_cam)
+    return SanaWmCamGeometry(
+        proj_q=proj_q,
+        proj_kv=proj_kv,
+        rope_cos=rope_cos,
+        rope_sin=rope_sin,
+        apply_q=apply_q,
+        apply_kv=apply_kv,
+        apply_output=apply_output,
+    )
+
+
+def _validate_cam_prep_inputs(
+    q_raw: torch.Tensor,
+    k_raw: torch.Tensor,
+    v_raw: torch.Tensor,
+    proj_q: torch.Tensor,
+    proj_kv: torch.Tensor,
+    rope_cos: torch.Tensor,
+    rope_sin: torch.Tensor,
+) -> int:
+    """Validate the fused cam-prep contract and return ``head_dim // 2``."""
+    if q_raw.shape != k_raw.shape or q_raw.shape != v_raw.shape:
+        raise ValueError(f"Sana-WM cam prep q/k/v shapes must match, got {q_raw.shape}, {k_raw.shape}, {v_raw.shape}.")
+    if q_raw.ndim != 4:
+        raise ValueError(f"Sana-WM cam prep q/k/v must be (B,N,H,D), got {tuple(q_raw.shape)}.")
+    batch, token_count, _num_heads, head_dim = q_raw.shape
+    if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0:
+        raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.")
+    half_dim = head_dim // 2
+    if proj_q.shape != (batch, token_count, 4, 4):
+        raise ValueError(f"Sana-WM cam prep proj_q shape {tuple(proj_q.shape)} != {(batch, token_count, 4, 4)}.")
+    if proj_kv.shape != (batch, token_count, 4, 4):
+        raise ValueError(f"Sana-WM cam prep proj_kv shape {tuple(proj_kv.shape)} != {(batch, token_count, 4, 4)}.")
+    if rope_cos.shape != (token_count, half_dim) or rope_sin.shape != (token_count, half_dim):
+        raise ValueError(
+            "Sana-WM cam prep rope table shapes must be "
+            f"{(token_count, half_dim)}, got {tuple(rope_cos.shape)} and {tuple(rope_sin.shape)}."
+        )
+    return half_dim
+
+
+def _apply_ray_projection(feats: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
+    batch, token_count, num_heads, half_dim = feats.shape
+    grouped = feats.reshape(batch, token_count, num_heads, half_dim // 4, 4)
+    out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), grouped)
+    return out.reshape(batch, token_count, num_heads, half_dim)
+
+
+def _apply_interleaved_rope(feats: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor) -> torch.Tensor:
+    half_dim = feats.shape[-1]
+    pair_index = torch.arange(half_dim, device=feats.device) ^ 1
+    paired = feats.index_select(-1, pair_index)
+    return feats * rope_cos.reshape(1, rope_cos.shape[0], 1, half_dim) + paired * rope_sin.reshape(
+        1, rope_sin.shape[0], 1, half_dim
+    )
+
+
+def cam_prep_func(
+    q_normed: torch.Tensor,
+    k_normed: torch.Tensor,
+    v_raw: torch.Tensor,
+    *,
+    proj_q: torch.Tensor,
+    proj_kv: torch.Tensor,
+    rope_cos: torch.Tensor,
+    rope_sin: torch.Tensor,
+    k_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Native equivalent of NVlabs ``cam_prep_func``.
+
+    The caller is responsible for the q/k RMSNorm, which must run through the
+    model's norm modules so a tensor-parallel shard still normalises over the
+    global channel width. This function only applies ReLU, the K scale, the
+    UCPE ray projection and RoPE.
+
+    Args:
+        q_normed: ``(B, N, H, D)`` camera Q, already RMSNorm'd.
+        k_normed: ``(B, N, H, D)`` camera K, already RMSNorm'd. K must also
+            include the temporal short convolution.
+        v_raw: ``(B, N, H, D)`` raw camera V.
+        proj_q: ``(B, N, 4, 4)`` Q ray projection matrices.
+        proj_kv: ``(B, N, 4, 4)`` K/V ray projection matrices.
+        rope_cos: ``(N, D//2)`` interleaved-pair RoPE cosine table.
+        rope_sin: ``(N, D//2)`` interleaved-pair RoPE sine table.
+        k_scale: ``D^-0.5 * spatial_tokens^-0.5``.
+
+    Returns:
+        ``q_trans``, ``k_trans``, ``v_trans`` in ``(B, H, D, N)`` layout, and
+        ``inflation_sq`` in ``(B, H, N)``.
+    """
+    half_dim = _validate_cam_prep_inputs(
+        q_normed,
+        k_normed,
+        v_raw,
+        proj_q,
+        proj_kv,
+        rope_cos,
+        rope_sin,
+    )
+
+    q_norm = torch.relu(q_normed.float())
+    k_norm = torch.relu(k_normed.float()) * float(k_scale)
+    value = v_raw.float()
+
+    q_half = _apply_ray_projection(q_norm[..., :half_dim], proj_q)
+    k_half = _apply_ray_projection(k_norm[..., :half_dim], proj_kv)
+    v_half = _apply_ray_projection(value[..., :half_dim], proj_kv)
+
+    q_rope = _apply_interleaved_rope(q_norm[..., half_dim:], rope_cos.float(), rope_sin.float())
+    k_rope = _apply_interleaved_rope(k_norm[..., half_dim:], rope_cos.float(), rope_sin.float())
+    v_rope = _apply_interleaved_rope(value[..., half_dim:], rope_cos.float(), rope_sin.float())
+
+    k_pre_sq = k_norm.square().sum(dim=-1).permute(0, 2, 1).contiguous()
+    k_post = torch.cat((k_half, k_rope), dim=-1)
+    k_post_sq = k_post.square().sum(dim=-1).permute(0, 2, 1).contiguous()
+    inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12)
+
+    def to_bhdn(tensor: torch.Tensor) -> torch.Tensor:
+        return tensor.permute(0, 2, 3, 1).contiguous().to(v_raw.dtype)
+
+    q_out = to_bhdn(torch.cat((q_half, q_rope), dim=-1))
+    k_out = to_bhdn(k_post)
+    v_out = to_bhdn(torch.cat((v_half, v_rope), dim=-1))
+    return q_out, k_out, v_out, inflation_sq
+
+
+__all__ = [
+    "SanaWmCamGeometry",
+    "cam_prep_func",
+    "prepare_cam_geometry",
+]

vllm_omni/diffusion/offloader/__init__.py

@@ -71,9 +71,8 @@ def get_offload_backend(
     Args:
         od_config: OmniDiffusionConfig with offload settings
         device: Target device (auto-detected if None)
-        host_weight_plan: Exact loader-produced backing plan. Checkpoint plans
-            may have skipped ordinary materialization; host weight cache plans
-            describe already-valid ordinary tensors that can be remapped.
+        host_weight_plan: Exact loader-produced backing plan, if ordinary
+            weight materialization was skipped.
 
     Returns:
         OffloadBackend instance or None if offloading disabled
@@ -86,14 +85,10 @@ def get_offload_backend(
     # Extract and validate configuration
     config = OffloadConfig.from_od_config(od_config)
 
-    plan_requires_backend = host_weight_plan is not None and host_weight_plan.requires_backend
     if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE:
-        if plan_requires_backend:
-            raise RuntimeError(
-                "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
-            )
-        logger.warning("Ignoring optional host weight cache plan because distributed layerwise offload is not selected")
-        host_weight_plan = None
+        raise RuntimeError(
+            "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
+        )
 
     # Return None if no offloading requested
     if config.strategy == OffloadStrategy.NONE:
@@ -101,7 +96,7 @@ def get_offload_backend(
 
     # Validate platform (CUDA required for now)
     if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1:
-        if plan_requires_backend:
+        if host_weight_plan is not None:
             raise RuntimeError(
                 "The loader skipped ordinary weight materialization for DLO, "
                 "but this platform cannot create the required offload backend"
@@ -117,7 +112,7 @@ def get_offload_backend(
         try:
             device = current_omni_platform.get_torch_device()
         except (NotImplementedError, AttributeError) as exc:
-            if plan_requires_backend:
+            if host_weight_plan is not None:
                 raise RuntimeError(
                     "The loader skipped ordinary weight materialization for DLO, "
                     "but the target device could not be resolved"

vllm_omni/diffusion/offloader/base.py

@@ -31,9 +31,6 @@ class OffloadConfig:
     # blocks from the loader-selected host backing with H2D only.
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
-    # Optional per-worker ceiling for registering shared host weight cache mappings.
-    # Zero applies no additional ceiling; pin_cpu_memory controls registration.
-    dlo_host_weight_cache_pin_limit_gib: float = 0.0
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -65,7 +62,7 @@ class OffloadConfig:
         dp_size = 1
         if parallel_config is not None:
             dp_size = getattr(parallel_config, "data_parallel_size", 1)
-            # HSDP's fully_shard_degree also contributes to effective DP
+            # HSDP shard and replica sizes determine the effective group size.
             hsdp_shard_size = getattr(parallel_config, "hsdp_shard_size", -1) if use_hsdp else -1
             hsdp_replicate_size = getattr(parallel_config, "hsdp_replicate_size", 1) if use_hsdp else 1
             if use_hsdp and hsdp_shard_size > 0:
@@ -104,7 +101,6 @@ class OffloadConfig:
         # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
         dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
-        dlo_host_weight_cache_pin_limit_gib = float(getattr(od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0))
         if dlo_resident_layers < 0:
             raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
         if dlo_resident_layers and dlo_use_allgather:
@@ -140,7 +136,6 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
-            dlo_host_weight_cache_pin_limit_gib=dlo_host_weight_cache_pin_limit_gib,
         )
 
 

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -1,280 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CUDA implementation of read-only host-mapping registration."""
-
-from __future__ import annotations
-
-import ctypes
-import mmap
-from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
-from typing import Protocol
-
-import torch
-
-from .host_registration import (
-    HostRegistrationBudgetError,
-    HostRegistrationCleanupError,
-    HostRegistrationError,
-)
-
-_CUDA_HOST_REGISTER_READ_ONLY = 0x08
-_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
-
-CudaHostRegistrationError = HostRegistrationError
-CudaHostRegistrationBudgetError = HostRegistrationBudgetError
-CudaHostRegistrationCleanupError = HostRegistrationCleanupError
-
-
-class _CudaRuntime(Protocol):
-    def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ...
-
-    def cudaHostUnregister(self, address: int) -> int: ...
-
-    def cudaGetErrorString(self, error: int) -> str | bytes: ...
-
-
-@dataclass(frozen=True)
-class _AddressRange:
-    start: int
-    end: int
-
-    @property
-    def size(self) -> int:
-        return self.end - self.start
-
-
-def _coalesce_ranges(
-    ranges: Sequence[tuple[int, int]],
-    page_size: int = mmap.PAGESIZE,
-) -> tuple[_AddressRange, ...]:
-    """Page-align and merge overlapping ranges from one backing mapping."""
-    if page_size <= 0:
-        raise ValueError(f"page_size must be positive, got {page_size}")
-
-    aligned: list[_AddressRange] = []
-    for start, size in ranges:
-        if start <= 0 or size < 0:
-            raise ValueError(f"invalid host range start={start}, size={size}")
-        if size == 0:
-            continue
-        aligned_start = start - start % page_size
-        end = start + size
-        aligned_end = ((end + page_size - 1) // page_size) * page_size
-        aligned.append(_AddressRange(aligned_start, aligned_end))
-
-    merged: list[_AddressRange] = []
-    for region in sorted(aligned, key=lambda item: (item.start, item.end)):
-        if merged and region.start <= merged[-1].end:
-            previous = merged[-1]
-            merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
-        else:
-            merged.append(region)
-    return tuple(merged)
-
-
-def _tensor_regions(
-    sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
-) -> tuple[_AddressRange, ...]:
-    """Resolve storage spans without merging unrelated file mappings."""
-    regions: list[_AddressRange] = []
-    for mapping_name, tensors in sources_by_mapping.items():
-        ranges: list[tuple[int, int]] = []
-        for tensor in tensors:
-            if tensor.device.type != "cpu":
-                raise HostRegistrationError(
-                    f"CUDA host registration requires CPU storage, but {mapping_name!r} contains {tensor.device}"
-                )
-            storage = tensor.untyped_storage()
-            ranges.append((storage.data_ptr(), storage.nbytes()))
-        regions.extend(_coalesce_ranges(ranges))
-    return tuple(regions)
-
-
-def _error_message(runtime: _CudaRuntime, error: int) -> str:
-    try:
-        message = runtime.cudaGetErrorString(error)
-    except Exception:
-        return str(error)
-    if isinstance(message, bytes):
-        return message.decode(errors="replace")
-    return str(message)
-
-
-def _consume_last_cuda_error(expected_error: int) -> None:
-    """Clear a handled CUDA Runtime error before returning to PyTorch.
-
-    PyTorch's ``torch.cuda.cudart()`` binding exposes host registration but not
-    ``cudaGetLastError``. Resolve the already-loaded process symbol so this call
-    clears the same CUDA Runtime instance that produced the registration error.
-    A different pending error is not safe to hide behind the staging fallback.
-    """
-    try:
-        get_last_error = ctypes.CDLL(None).cudaGetLastError
-        get_last_error.argtypes = []
-        get_last_error.restype = ctypes.c_int
-        pending_error = int(get_last_error())
-    except (AttributeError, OSError) as exc:
-        raise HostRegistrationCleanupError("cannot clear CUDA's pending error after host-registration failure") from exc
-
-    if pending_error not in (0, expected_error):
-        raise HostRegistrationCleanupError(
-            f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}"
-        )
-
-
-def _handled_error_message(runtime: _CudaRuntime, error: int) -> str:
-    """Format and consume one nonzero CUDA Runtime return code."""
-    error_code = int(error)
-    message = _error_message(runtime, error)
-    try:
-        _consume_last_cuda_error(error_code)
-    except HostRegistrationCleanupError as exc:
-        raise HostRegistrationCleanupError(f"{message}; {exc}") from exc
-    return message
-
-
-def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool:
-    """Query support required by immutable file-backed cache mappings."""
-    try:
-        get_attribute = getattr(runtime, "cudaDeviceGetAttribute", None)
-        if get_attribute is None:
-            get_attribute = ctypes.CDLL(None).cudaDeviceGetAttribute
-            get_attribute.argtypes = [
-                ctypes.POINTER(ctypes.c_int),
-                ctypes.c_int,
-                ctypes.c_int,
-            ]
-            get_attribute.restype = ctypes.c_int
-    except (AttributeError, OSError) as exc:
-        raise HostRegistrationError("cannot query CUDA read-only host-registration support") from exc
-
-    supported = ctypes.c_int()
-    try:
-        error = get_attribute(
-            ctypes.byref(supported),
-            _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED,
-            torch.accelerator.current_device_index(),
-        )
-    except Exception as exc:
-        raise HostRegistrationError(f"cannot query CUDA read-only host-registration support: {exc}") from exc
-    if int(error) != 0:
-        raise HostRegistrationError(
-            f"cudaDeviceGetAttribute(read-only host registration) failed: {_handled_error_message(runtime, error)}"
-        )
-    return bool(supported.value)
-
-
-class CudaHostRegistration:
-    """Own CUDA registrations for already-existing file-backed CPU tensors."""
-
-    def __init__(
-        self,
-        runtime: _CudaRuntime,
-        regions: tuple[_AddressRange, ...],
-    ) -> None:
-        self._runtime = runtime
-        self._regions = regions
-        self._closed = False
-
-    @classmethod
-    def create(
-        cls,
-        sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
-        *,
-        max_bytes: int | None,
-    ) -> CudaHostRegistration:
-        regions = _tensor_regions(sources_by_mapping)
-        total_bytes = sum(region.size for region in regions)
-        if max_bytes is not None and total_bytes > max_bytes:
-            raise HostRegistrationBudgetError(
-                f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
-            )
-        if not regions:
-            raise HostRegistrationError("no non-empty host ranges were available for registration")
-        if not torch.cuda.is_available():
-            raise HostRegistrationError("CUDA is not available")
-
-        try:
-            runtime = torch.cuda.cudart()
-        except Exception as exc:
-            raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
-        if not _supports_read_only_host_registration(runtime):
-            raise HostRegistrationError(
-                "CUDA device does not support read-only host registration required by immutable "
-                "host weight cache mappings"
-            )
-
-        registered: list[_AddressRange] = []
-        try:
-            for region in regions:
-                error = runtime.cudaHostRegister(
-                    region.start,
-                    region.size,
-                    _CUDA_HOST_REGISTER_READ_ONLY,
-                )
-                if int(error) != 0:
-                    raise HostRegistrationError(
-                        "cudaHostRegister(read-only) failed for "
-                        f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
-                    )
-                registered.append(region)
-
-            unpinned = [
-                mapping_name
-                for mapping_name, tensors in sources_by_mapping.items()
-                if any(tensor.numel() and not tensor.is_pinned() for tensor in tensors)
-            ]
-            if unpinned:
-                raise HostRegistrationError(
-                    f"CUDA registration succeeded but PyTorch did not recognize pinned storage for {unpinned[:3]}"
-                )
-        except Exception as exc:
-            rollback_errors: list[str] = []
-            for region in reversed(registered):
-                try:
-                    error = runtime.cudaHostUnregister(region.start)
-                    if int(error) != 0:
-                        rollback_errors.append(
-                            f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}"
-                        )
-                except Exception as rollback_exc:
-                    rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
-            if rollback_errors:
-                raise HostRegistrationCleanupError(
-                    f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}"
-                ) from exc
-            if isinstance(exc, HostRegistrationError):
-                raise
-            raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc
-
-        return cls(runtime, regions)
-
-    @property
-    def total_bytes(self) -> int:
-        return sum(region.size for region in self._regions)
-
-    @property
-    def region_count(self) -> int:
-        return len(self._regions)
-
-    def close(self) -> list[str]:
-        """Unregister every range, returning errors after best-effort cleanup."""
-        if self._closed:
-            return []
-        errors: list[str] = []
-        failed: list[_AddressRange] = []
-        for region in reversed(self._regions):
-            try:
-                error = self._runtime.cudaHostUnregister(region.start)
-                if int(error) != 0:
-                    errors.append(
-                        f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}"
-                    )
-                    failed.append(region)
-            except Exception as exc:
-                errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
-                failed.append(region)
-        self._regions = tuple(reversed(failed))
-        self._closed = not self._regions
-        return errors

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -18,9 +18,7 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
 from __future__ import annotations
 
-import time
 from itertools import chain
-from pathlib import Path
 from typing import Any
 
 import torch
@@ -29,7 +27,6 @@ from torch import nn
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
-from vllm_omni.diffusion.model_loader.host_weight_cache_utils import file_digest
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
 )
@@ -37,12 +34,6 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
-from .host_registration import (
-    HostRegistration,
-    HostRegistrationCleanupError,
-    HostRegistrationError,
-    register_host_mappings,
-)
 from .module_collector import ModuleDiscovery
 from .offload_plan import (
     OffloadPlan,
@@ -102,7 +93,6 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.rank = rank
         self.pin_memory = pin_memory
         self.rank_local_mmap = rank_local_mmap
-        self.registered_mmap = False
         self.tensor_transforms = tensor_transforms or {}
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -454,26 +444,6 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
-    @staticmethod
-    def _resolve_mmap_source(
-        source_info: dict[str, Any],
-        meta: dict[str, Any],
-        dtype: torch.dtype,
-    ) -> torch.Tensor:
-        source = source_info["tensor"]
-        transform = source_info["transform"]
-        if callable(transform):
-            source = transform(source)
-        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
-            raise ValueError(
-                "mmap weight transform changed tensor layout for "
-                f"{source_info['name']!r}: expected dtype={dtype}, "
-                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
-                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
-                f"stride={source.stride()}"
-            )
-        return source
-
     @staticmethod
     def _pack_mmap_sources(
         cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -486,7 +456,18 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                source = source_info["tensor"]
+                transform = source_info["transform"]
+                if callable(transform):
+                    source = transform(source)
+                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
+                    raise ValueError(
+                        "mmap weight transform changed tensor layout for "
+                        f"{source_info['name']!r}: expected dtype={dtype}, "
+                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
+                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
+                        f"stride={source.stride()}"
+                    )
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -500,32 +481,6 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             staged[dtype] = destination
         return staged
 
-    @staticmethod
-    def _copy_mmap_sources_to_device(
-        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
-        metadata: dict[torch.dtype, list[dict[str, Any]]],
-        gpu_buffers: dict[torch.dtype, torch.Tensor],
-        *,
-        non_blocking: bool,
-    ) -> None:
-        """Copy registered source views directly into flattened device buffers."""
-        for dtype, metas in metadata.items():
-            destination = gpu_buffers[dtype]
-            sources = cpu_sources[dtype]
-            for source_info, meta in zip(sources, metas, strict=True):
-                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
-                start = meta["offset"]
-                physical_storage = destination[start : start + meta["numel"]]
-                async_copy = non_blocking and source.is_pinned()
-                if source.is_contiguous():
-                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
-                else:
-                    torch.as_strided(
-                        physical_storage,
-                        size=source.shape,
-                        stride=source.stride(),
-                    ).copy_(source, non_blocking=async_copy)
-
     def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
         """Pack this block's mmap views into one bounded host staging slot."""
         previous_copy = self.cpu_staging_events[slot]
@@ -559,27 +514,17 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
-            if self.rank_local_mmap and self.registered_mmap:
-                with current_omni_platform.stream(self.copy_stream):
-                    self._copy_mmap_sources_to_device(
-                        self.cpu_sources,
-                        self.metadata,
-                        gpu_weights,
-                        non_blocking=non_blocking,
-                    )
-                    evt.record(self.copy_stream)
-            else:
-                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
-                with current_omni_platform.stream(self.copy_stream):
-                    for dtype, cpu_shard in cpu_weights.items():
-                        gw = gpu_weights[dtype]
-                        async_copy = non_blocking and cpu_shard.is_pinned()
-                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
-                    evt.record(self.copy_stream)
-                if self.rank_local_mmap:
-                    # The CPU slot may be overwritten only after this H2D copy has
-                    # finished.  The shared event protects reuse by another hook.
-                    self.cpu_staging_events[slot] = evt
+            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+            with current_omni_platform.stream(self.copy_stream):
+                for dtype, cpu_shard in cpu_weights.items():
+                    gw = gpu_weights[dtype]
+                    async_copy = non_blocking and cpu_shard.is_pinned()
+                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                evt.record(self.copy_stream)
+            if self.rank_local_mmap:
+                # The CPU slot may be overwritten only after this H2D copy has
+                # finished.  The shared event protects reuse by another hook.
+                self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -789,7 +734,6 @@ class PinnedResidentLayerGroup:
         self.copy_stream = copy_stream
         self.loaded = False
         self.rank_local_mmap = rank_local_mmap
-        self.registered_mmap = False
         self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -860,14 +804,6 @@ class PinnedResidentLayerGroup:
         ready = current_omni_platform.Event()
         with current_omni_platform.stream(self.copy_stream):
             for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
-                if self.rank_local_mmap and self.registered_mmap:
-                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
-                        state["cpu_sources"],
-                        state["metadata"],
-                        block_buffers,
-                        non_blocking=True,
-                    )
-                    continue
                 if self.rank_local_mmap:
                     slot = index % 2
                     previous_copy = self._cpu_staging_events[slot]
@@ -965,11 +901,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
         self._using_mmap = False
         self._using_rank_local_mmap = False
-        self._using_registered_mmap = False
         self.host_weight_plan = host_weight_plan
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
-        self._host_weight_cache_mapped_sources: dict[str, list[torch.Tensor]] = {}
-        self._host_registration: HostRegistration | None = None
 
     def load_resident_layers(self) -> None:
         """Load the model-declared leading blocks for the denoise stage."""
@@ -1090,7 +1023,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     framework="pt",
                     device="cpu",
                 )
-            tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
+            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
             if is_parameter:
                 replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                 parent._parameters[leaf_name] = replacement
@@ -1170,179 +1103,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 f"(first 5: {remaining_meta[:5]})."
             )
 
-    def _load_weights_from_host_weight_cache(
-        self,
-        pipeline: nn.Module,
-        modules,
-        plan: HostWeightPlan,
-    ) -> None:
-        """Replace final ordinary-loader tensors with validated cache views."""
-        from safetensors import safe_open
-
-        if any(binding.transform is not None for binding in plan.bindings.values()):
-            raise RuntimeError("Host weight cache bindings cannot carry deferred checkpoint transforms")
-
-        cache_files = {binding.file_path for binding in plan.bindings.values()}
-        if plan.expected_file_digests is None or set(plan.expected_file_digests) != cache_files:
-            raise RuntimeError("Host weight cache plan does not cover every shard digest")
-        for file_path in sorted(cache_files):
-            try:
-                actual_digest = file_digest(Path(file_path))
-            except OSError as exc:
-                raise RuntimeError(f"Cannot revalidate host weight cache shard {file_path}: {exc}") from exc
-            if actual_digest != plan.expected_file_digests[file_path]:
-                raise RuntimeError(f"Host weight cache shard changed after loader validation: {file_path}")
-
-        required_names: set[str] = set()
-        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
-            required_names.update(f"{dit_name}.{name}" for name, _ in dit_module.named_parameters())
-            for name, _ in dit_module.named_buffers():
-                parent_path, _, leaf_name = name.rpartition(".")
-                owner = dit_module.get_submodule(parent_path)
-                if leaf_name not in owner._non_persistent_buffers_set:
-                    required_names.add(f"{dit_name}.{name}")
-        if required_names != set(plan.bindings):
-            missing = sorted(required_names - set(plan.bindings))
-            unexpected = sorted(set(plan.bindings) - required_names)
-            raise RuntimeError(
-                "Host weight cache plan coverage changed after loader validation: "
-                f"missing={missing[:5]}, unexpected={unexpected[:5]}"
-            )
-
-        file_cache: dict[str, Any] = {}
-        replacements: list[tuple[torch.Tensor, torch.Tensor]] = []
-        mapped_sources: dict[str, list[torch.Tensor]] = {}
-        try:
-            # Validate and map every tensor before mutating the pipeline. A
-            # missing/corrupt shard therefore leaves the ordinary tensors intact.
-            for runtime_name, binding in plan.bindings.items():
-                parent_path, _, leaf_name = runtime_name.rpartition(".")
-                try:
-                    parent = pipeline.get_submodule(parent_path)
-                except AttributeError as exc:
-                    raise RuntimeError(f"Host weight cache target module {parent_path!r} no longer exists") from exc
-
-                target = parent._parameters.get(leaf_name)
-                if target is None:
-                    target = parent._buffers.get(leaf_name)
-                if target is None:
-                    raise RuntimeError(f"Host weight cache target tensor {runtime_name!r} no longer exists")
-                if target.device.type != "cpu" or target.is_meta:
-                    raise RuntimeError(
-                        f"Host weight cache target {runtime_name!r} is not an ordinary CPU tensor: {target.device}"
-                    )
-
-                if binding.file_path not in file_cache:
-                    file_cache[binding.file_path] = safe_open(
-                        binding.file_path,
-                        framework="pt",
-                        device="cpu",
-                    )
-                tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
-                if tensor.shape != target.shape or tensor.dtype != target.dtype:
-                    raise RuntimeError(
-                        f"Host weight cache metadata changed for {runtime_name!r}: "
-                        f"cache={tuple(tensor.shape)}/{tensor.dtype}, "
-                        f"runtime={tuple(target.shape)}/{target.dtype}"
-                    )
-                replacements.append((target, tensor))
-                mapped_sources.setdefault(binding.file_path, []).append(tensor)
-        except Exception:
-            file_cache.clear()
-            raise
-
-        # Preserve Parameter objects and their loader/runtime attributes; only
-        # their storage changes. This also keeps any references held by model
-        # code valid while dropping the private allocator-backed storage.
-        for target, tensor in replacements:
-            set_tensor_storage(target, tensor)
-
-        self._mmap_transforms_by_tensor_id.clear()
-        self._mmap_file_cache = file_cache
-        self._host_weight_cache_mapped_sources = mapped_sources
-        logger.info(
-            "Remapped %d final DiT tensors from host weight cache %s",
-            len(replacements),
-            plan.runtime_layout_key[:12] if plan.runtime_layout_key else "unknown",
-        )
-
-    def _try_register_host_weight_cache_mmap(self) -> bool:
-        """Register complete cache mappings under the pinned-memory policy."""
-        pin_limit_gib = self.config.dlo_host_weight_cache_pin_limit_gib
-        if not self.config.pin_cpu_memory:
-            logger.info("DLO host weight cache registration disabled by pin_cpu_memory=False; using host staging")
-            return False
-        if not self._host_weight_cache_mapped_sources:
-            logger.warning("DLO host weight cache registration found no mapped sources; using bounded host staging")
-            return False
-
-        max_bytes = int(pin_limit_gib * 1024**3) if pin_limit_gib > 0 else None
-        started = time.perf_counter()
-        try:
-            registration = register_host_mappings(
-                self._host_weight_cache_mapped_sources,
-                device=self.device,
-                max_bytes=max_bytes,
-            )
-        except HostRegistrationCleanupError:
-            # Falling back could unmap a range the platform still owns. Abort
-            # the worker so process/context teardown provides final cleanup.
-            logger.exception("DLO host weight cache registration rollback failed")
-            raise
-        except HostRegistrationError as exc:
-            logger.warning(
-                "DLO host weight cache direct H2D unavailable (%s); using bounded host staging",
-                exc,
-            )
-            return False
-
-        self._host_registration = registration
-        logger.info(
-            "Registered %.2f GiB of shared host weight cache mappings in %d range(s) for direct H2D in %.3f s",
-            registration.total_bytes / 1024**3,
-            registration.region_count,
-            time.perf_counter() - started,
-        )
-        return True
-
-    def _configure_host_weight_cache_transfer(
-        self,
-        hooks: list[DistributedLayerwiseOffloadHook],
-    ) -> None:
-        """Select registered mmap or staging once for every cache consumer."""
-        plan = self.host_weight_plan
-        if (
-            (not hooks and self._resident_layer_group is None)
-            or not self._using_rank_local_mmap
-            or plan is None
-            or plan.backing_kind != "host_weight_cache"
-        ):
-            return
-
-        self._using_registered_mmap = self._try_register_host_weight_cache_mmap()
-        self._host_weight_cache_mapped_sources.clear()
-        for hook in hooks:
-            hook.registered_mmap = self._using_registered_mmap
-        if self._resident_layer_group is not None:
-            self._resident_layer_group.registered_mmap = self._using_registered_mmap
-            if self._using_registered_mmap:
-                self._resident_layer_group._cpu_staging_buffers.clear()
-
-    def _release_registered_mmap(self) -> bool:
-        registration = self._host_registration
-        if registration is None:
-            return True
-        errors = registration.close()
-        if errors:
-            logger.error(
-                "DLO host weight cache unregistration failed; retaining mmap handles for retry/process exit: %s",
-                errors[:3],
-            )
-            return False
-        self._host_registration = None
-        logger.info("Unregistered DLO host weight cache mappings")
-        return True
-
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.
 
@@ -1683,41 +1443,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._using_mmap = self.host_weight_plan is not None
         self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
         if self._using_mmap:
-            if self.host_weight_plan.backing_kind == "checkpoint_mmap":
-                self._load_weights_via_mmap(
-                    pipeline,
-                    modules,
-                    self.host_weight_plan,
-                )
-            elif self.host_weight_plan.backing_kind == "host_weight_cache":
-                if self.config.dlo_use_allgather:
-                    raise ValueError("Host weight cache backing is no-AllGather only in the first implementation")
-                self._using_rank_local_mmap = True
-                try:
-                    self._load_weights_from_host_weight_cache(
-                        pipeline,
-                        modules,
-                        self.host_weight_plan,
-                    )
-                except Exception as exc:
-                    logger.warning(
-                        "DLO host weight cache remapping failed; retaining ordinary host tensors: %s",
-                        exc,
-                    )
-                    self.host_weight_plan = None
-                    self._using_mmap = False
-                    self._using_rank_local_mmap = False
-                else:
-                    # Assignments above dropped the model's private tensor
-                    # references. Reclaim those allocations before hook setup
-                    # so steady-state PSS reflects the shared mmap backing.
-                    self._cleanup_after_loading()
-            else:
+            if self.host_weight_plan.backing_kind != "checkpoint_mmap":
                 raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
+            self._load_weights_via_mmap(
+                pipeline,
+                modules,
+                self.host_weight_plan,
+            )
             if self._using_rank_local_mmap:
                 logger.info(
-                    "DLO rank-local mmap storage enabled: file-backed pages are node-shared; "
-                    "transfer setup will select registered direct H2D or bounded host staging"
+                    "DLO rank-local mmap storage enabled: checkpoint pages are "
+                    "node-shared; each worker owns only two bounded host staging slots"
                 )
         else:
             remaining_meta = [
@@ -1887,9 +1623,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             )
             pipeline._dlo_residency_controller = self
 
-        all_hooks = [hook for group in self._all_hook_groups for hook in group]
-        self._configure_host_weight_cache_transfer(all_hooks)
-
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
             if self._using_mmap and not self.enabled:
@@ -1899,14 +1632,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Unified allocation: 2 shared output buffers + 2 shared shard buffers
         # sized to the max block across ALL module groups (gen_layers +
         # language_model).  Groups execute sequentially, so 2 buffers suffice.
+        all_hooks: list[DistributedLayerwiseOffloadHook] = []
+        for group in self._all_hook_groups:
+            all_hooks.extend(group)
+
         unified_buffers = self._allocate_shared_buffers(all_hooks)
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
-
         unified_cpu_staging = None
         cpu_staging_events = None
-        if self._using_rank_local_mmap and not self._using_registered_mmap:
+        if self._using_rank_local_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1988,7 +1724,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     def _release_mmap_handles(self) -> None:
         """Release safetensors mmap file handles."""
         self._mmap_transforms_by_tensor_id.clear()
-        self._host_weight_cache_mapped_sources.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
             del self._mmap_file_cache
@@ -2011,11 +1746,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         if not self.enabled and not hasattr(self, "_mmap_file_cache"):
             return
 
-        # Copy and communication streams are asynchronous in both AllGather
-        # and no-AllGather modes. Drain all device work before hooks, process
-        # groups, registrations, or mapped storage can be torn down.
-        current_omni_platform.synchronize()
-        registration_released = self._release_registered_mmap()
+        if self._using_rank_local_mmap:
+            current_omni_platform.synchronize()
 
         for blocks in self._blocks:
             for block in blocks:
@@ -2031,11 +1763,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
-        if registration_released:
-            self._release_mmap_handles()
+        self._release_mmap_handles()
         self._using_mmap = False
         self._using_rank_local_mmap = False
-        self._using_registered_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/offloader/host_registration.py

@@ -1,62 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Platform-neutral contract for registering existing host mappings."""
-
-from __future__ import annotations
-
-from collections.abc import Mapping, Sequence
-from typing import Protocol
-
-import torch
-
-
-class HostRegistrationError(RuntimeError):
-    """Raised when existing host ranges cannot be registered safely."""
-
-
-class HostRegistrationBudgetError(HostRegistrationError):
-    """Raised before mutation when the complete mapping exceeds its budget."""
-
-
-class HostRegistrationCleanupError(HostRegistrationError):
-    """Raised when partial registration cannot be rolled back safely."""
-
-
-class HostRegistration(Protocol):
-    """Lifecycle owned by DLO for one platform's registered host ranges."""
-
-    @property
-    def total_bytes(self) -> int: ...
-
-    @property
-    def region_count(self) -> int: ...
-
-    def close(self) -> list[str]: ...
-
-
-def register_host_mappings(
-    sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
-    *,
-    device: torch.device,
-    max_bytes: int | None,
-) -> HostRegistration:
-    """Register mappings with the active platform or report unsupported use.
-
-    CUDA is the first implementation. Other platforms retain DLO's bounded
-    host-staging path until they provide an equivalent registration backend.
-    ``max_bytes=None`` applies no additional ceiling to the complete mapping.
-    """
-    if device.type == "cuda":
-        from .cuda_host_registration import CudaHostRegistration
-
-        return CudaHostRegistration.create(sources_by_mapping, max_bytes=max_bytes)
-    raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}")
-
-
-__all__ = [
-    "HostRegistration",
-    "HostRegistrationBudgetError",
-    "HostRegistrationCleanupError",
-    "HostRegistrationError",
-    "register_host_mappings",
-]

vllm_omni/diffusion/registry.py

@@ -171,6 +171,11 @@ _DIFFUSION_MODELS = {
         "pipeline_ming_imagegen",
         "MingImagePipeline",
     ),
+    "SanaWmPipeline": (
+        "sana_wm",
+        "pipeline_sana_wm",
+        "SanaWmPipeline",
+    ),
     "InternVLAA1Pipeline": (
         "internvla_a1",
         "pipeline_internvla_a1",
@@ -633,6 +638,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "LingBotVideoPipeline": "get_lingbot_video_pre_process_func",
     "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func",
     "MagiHumanPipeline": "get_magi_human_pre_process_func",
+    "SanaWmPipeline": "get_sana_wm_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
     "SoulXSingerPipeline": "get_soulxsinger_pre_process_func",

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -279,9 +279,8 @@ class DiffusionWorker:
                 allgather_degree=parallel_config.allgather_degree,
                 tensor_parallel_size=parallel_config.tensor_parallel_size,
                 pipeline_parallel_size=parallel_config.pipeline_parallel_size,
-                fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1,
-                hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1,
                 enable_expert_parallel=parallel_config.enable_expert_parallel,
+                use_hsdp=parallel_config.use_hsdp,
             )
             if (
                 getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
@@ -859,18 +858,11 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        try:
-            if self.model_runner is not 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()
-        finally:
-            destroy_distributed_env()
+        if self.model_runner is not None:
+            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+            if mgr is not None:
+                mgr.shutdown_prefetch()
+        destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -545,10 +545,6 @@ class OrchestratorArgs:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
-    dlo_use_host_weight_cache: bool = False
-    dlo_host_weight_cache_dir: str | None = None
-    dlo_host_weight_cache_lock_timeout: float = 600.0
-    dlo_host_weight_cache_pin_limit_gib: float = 0.0
     dlo_resident_layers: int = 0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -971,7 +971,7 @@ class AsyncOmniEngine:
             ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict"
             sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size")
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
-            data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1
+            data_parallel_size = normalized_kwargs.get("data_parallel_size")
             tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1
             cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
@@ -1004,6 +1004,11 @@ class AsyncOmniEngine:
                 hsdp_replicate_size=hsdp_replicate_size,
             )
 
+        num_gpus = normalized_kwargs.get("num_gpus")
+        if num_gpus is not None:
+            num_gpus = int(num_gpus)
+            parallel_config.resolve_data_parallel_size(num_gpus)
+
         num_devices = max(1, int(parallel_config.world_size))
         devices = ",".join(str(i) for i in range(num_devices))
         model_class_name = kwargs.get("model_class_name", None)
@@ -1040,10 +1045,6 @@ class AsyncOmniEngine:
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
-            "dlo_use_host_weight_cache": kwargs.get("dlo_use_host_weight_cache", False),
-            "dlo_host_weight_cache_dir": kwargs.get("dlo_host_weight_cache_dir", None),
-            "dlo_host_weight_cache_lock_timeout": kwargs.get("dlo_host_weight_cache_lock_timeout", 600.0),
-            "dlo_host_weight_cache_pin_limit_gib": kwargs.get("dlo_host_weight_cache_pin_limit_gib", 0.0),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (
@@ -1096,6 +1097,8 @@ class AsyncOmniEngine:
                 else {}
             ),
         }
+        if num_gpus is not None:
+            stage_engine_args["num_gpus"] = num_gpus
         # Only set dtype if it was already explicitly passed and normalized
         if "dtype" in normalized_kwargs:
             stage_engine_args["dtype"] = normalized_kwargs["dtype"]

vllm_omni/entrypoints/cli/serve.py

@@ -722,25 +722,6 @@ class OmniServeCommand(CLISubcommand):
                 "independently (including existing TP shards)."
             ),
         )
-        omni_config_group.add_argument(
-            "--dlo-use-host-weight-cache",
-            action="store_true",
-            help=(
-                "Use the node-local final-layout host weight cache for no-AllGather DLO. "
-                "With the default pinned-memory policy, supported platforms register the immutable mapping "
-                "for direct H2D; otherwise DLO uses bounded host staging."
-            ),
-        )
-        omni_config_group.add_argument(
-            "--dlo-host-weight-cache-pin-limit-gib",
-            type=float,
-            default=0.0,
-            help=(
-                "Optional maximum GiB of read-only host weight cache mappings each worker may "
-                "register for direct H2D. Zero applies no additional registration ceiling; "
-                "this option requires --dlo-use-host-weight-cache."
-            ),
-        )
         omni_config_group.add_argument(
             "--dlo-resident-layers",
             type=int,

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -10,7 +10,12 @@ from typing import TYPE_CHECKING
 from vllm.inputs import tokens_input
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    ARTTSAdapter,
+    PreparedRequest,
+    apply_max_new_tokens,
+    conditioning_cache_salt,
+)
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -104,6 +109,15 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         prompt["cache_salt"] = conditioning_cache_salt(request, tts_params)
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
 
+    def apply_sampling_overrides(
+        self,
+        sampling_params_list: list,
+        request: "OpenAICreateSpeechRequest",
+        prompt: dict | None = None,
+        request_id: str | None = None,
+    ) -> list:
+        return apply_max_new_tokens(sampling_params_list, request)
+
 
 @register_tts_adapter
 class MossTTSNanoAdapter(_MossTTSAdapterBase):

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -1,3 +1,4 @@
+import os
 import time
 from pathlib import Path
 
@@ -85,3 +86,46 @@ def download_weights_from_hf_specific(
                 time_taken,
             )
     return hf_folder
+
+
+def resolve_model_to_local_path(
+    model: str,
+    *,
+    allow_download: bool = False,
+    allow_patterns: list[str] | None = None,
+    cache_dir: str | None = None,
+) -> str:
+    """Resolve a model reference to a local directory.
+    Failure is always raised: whether a failure is fatal is the caller's decision.
+
+    Args:
+        model: A local directory or an HF Hub repo ID.
+        allow_download: Fetch missing files instead of requiring a warm cache.
+            Downloads go through :func:`download_weights_from_hf_specific`
+        allow_patterns: Restrict the fetch to these glob patterns.
+            Defaults to the whole repository. Ignored when model is already a directory.
+        cache_dir: Download cache location; None uses the HF default.
+
+    Returns:
+        Path to a local directory containing the model files.
+    """
+    if os.path.isdir(model):
+        return model
+
+    if allow_download:
+        return download_weights_from_hf_specific(
+            model_name_or_path=model,
+            cache_dir=cache_dir,
+            allow_patterns=allow_patterns or ["*"],
+            require_all=True,
+        )
+
+    # Cache-only lookups stay on huggingface_hub even under VLLM_USE_MODELSCOPE:
+    # ModelScope's snapshot_download has no local_files_only equivalent,
+    # so a ModelScope deployment that needs resolution here must pass allow_download.
+    return huggingface_hub.snapshot_download(
+        model,
+        cache_dir=cache_dir,
+        allow_patterns=allow_patterns,
+        local_files_only=True,
+    )

vllm_omni/model_executor/models/common/ming/__init__.py

@@ -1,2 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Components shared by the Ming-family model packages.
+
+Naming: The flow-matching pieces all use the cfm_ prefix,
+because "flow matching", "fm" and "cfm" were three names for one concept here:
+
+* dit_blocks    - DiT building blocks (DiTBlock, FinalLayer, CondEmbedder)
+* cfm_dit       - the DiT used as the CFM velocity model
+* cfm_solver    - timestep schedules and the ODE/SDE integrator
+* cfm_head      - inference-only CFM wrapper preserving cfm.model.* names
+* cfm_cudagraph - batch-1 CUDA-graph executor for the sampling loop
+"""

vllm_omni/model_executor/models/common/ming/aggregator.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn as nn
 from x_transformers.x_transformers import RotaryEmbedding
 
-from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.common.ming.dit_blocks import DiTBlock, FinalLayer
 
 
 class Aggregator(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 import torch
 import torch.nn as nn
 
-from vllm_omni.model_executor.models.common.ming.fm import build_timesteps
+from vllm_omni.model_executor.models.common.ming.cfm_solver import build_timesteps
 
 
 class CFMSampler(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_dit.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 from x_transformers.x_transformers import RotaryEmbedding
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
-from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer
 
 
 class DiT(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_head.py

@@ -5,9 +5,8 @@
 import torch
 import torch.nn as nn
 
-from vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps
-
-from .fm.dit import DiT
+from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT
+from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, build_timesteps
 
 
 class FlowLoss(nn.Module):

vllm_omni/model_executor/models/common/ming/cfm_solver.py

@@ -3,7 +3,21 @@
 
 import torch
 
-from .dit import get_epss_timesteps
+
+def get_epss_timesteps(n, device, dtype):
+    dt = 1 / 32
+    predefined_timesteps = {
+        5: [0, 2, 4, 8, 16, 32],
+        6: [0, 2, 4, 6, 8, 16, 32],
+        7: [0, 2, 4, 6, 8, 16, 24, 32],
+        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
+        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+    }
+    t = predefined_timesteps.get(n, [])
+    if not t:
+        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
+    return dt * torch.tensor(t, device=device, dtype=dtype)
 
 
 def apply_sway_sampling(t, sway_sampling_coef):

vllm_omni/model_executor/models/common/ming/dit_blocks.py

@@ -216,19 +216,3 @@ class CondEmbedder(nn.Module):
 
     def forward(self, llm_cond):
         return self.cond_embedder(llm_cond)
-
-
-def get_epss_timesteps(n, device, dtype):
-    dt = 1 / 32
-    predefined_timesteps = {
-        5: [0, 2, 4, 8, 16, 32],
-        6: [0, 2, 4, 6, 8, 16, 32],
-        7: [0, 2, 4, 6, 8, 16, 24, 32],
-        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
-        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-    }
-    t = predefined_timesteps.get(n, [])
-    if not t:
-        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
-    return dt * torch.tensor(t, device=device, dtype=dtype)

vllm_omni/model_executor/models/common/ming/speaker_extractor.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py
+"""CAM++ speaker-embedding extraction, shared by the Ming model packages."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torchaudio
+from vllm.multimodal.media.audio import load_audio
+
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
+
+
+class SpeakerEmbeddingExtractor:
+    """CAM++ ONNX speaker embeddings, resampled to the extractor's rate."""
+
+    def __init__(self, model, target_sr=16000, *, allow_download: bool = False):
+        local_model_path = resolve_model_to_local_path(
+            model,
+            allow_download=allow_download,
+            allow_patterns=["campplus.onnx"],
+        )
+        campplus_path = os.path.join(local_model_path, "campplus.onnx")
+        if not os.path.exists(campplus_path):
+            raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}")
+
+        self.target_sr = int(target_sr)
+        self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr)
+
+    def extract_from_waveform(self, waveform, sample_rate):
+        if not isinstance(waveform, torch.Tensor):
+            waveform = torch.as_tensor(waveform)
+
+        tensor = waveform.detach().to(torch.float32)
+        if tensor.ndim == 1:
+            tensor = tensor.unsqueeze(0)
+        if int(sample_rate) != self.target_sr:
+            tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor)
+
+        embedding = self._core._extract_spk_embedding(tensor)
+        return embedding.squeeze(0).to(dtype=torch.float32)
+
+    def extract_from_file(self, audio_path):
+        audio, sample_rate = load_audio(audio_path, sr=None, mono=False)
+        waveform = torch.from_numpy(audio).contiguous()
+        return self.extract_from_waveform(waveform, sample_rate)
+
+    def extract_many(self, audio_paths):
+        return [self.extract_from_file(path) for path in audio_paths]
+
+
+__all__ = ["SpeakerEmbeddingExtractor"]

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py

@@ -25,14 +25,15 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
-from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig
 from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig
 
 from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT
-from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input
+from .talker_module import CFM, DiT, MingAudioGenerator, build_tts_input
 from .text_processing import segment_and_normalize
 from .voice_presets import VoicePresetRegistry
 
@@ -119,7 +120,11 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin)
             DiT(llm_input_dim=self.hidden_size, **config.flowmodel),
             steps=config.steps,
         )
-        self.aggregator = Aggregator(llm_input_dim=self.hidden_size, **config.aggregator)
+        # config.aggregator still wins `in_channels` if the checkpoint states it.
+        self.aggregator = Aggregator(
+            llm_input_dim=self.hidden_size,
+            **{"in_channels": self.latent_dim, **config.aggregator},
+        )
         self.stop_head = nn.Linear(self.hidden_size, 2, bias=True)
         self.spk_head = nn.Linear(SPEAKER_EMBEDDING_DIM, self.hidden_size, bias=True)
 
@@ -525,27 +530,22 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin)
 
     def _iter_talker_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]:
         """Yield (name, tensor) pairs from talker/model*.safetensors."""
-        model_path = self._model_path
-        # Try local path first
-        for candidate in (os.path.join(model_path, "talker"), model_path):
+        model_root = resolve_model_to_local_path(
+            self._model_path,
+            allow_download=True,
+            allow_patterns=["talker/model*.safetensors"],
+            cache_dir=self.vllm_config.load_config.download_dir,
+        )
+
+        # Nested talker/ layout first, then a flat checkpoint at the root.
+        for candidate in (os.path.join(model_root, "talker"), model_root):
             sf_files = sorted(glob_module.glob(os.path.join(candidate, "model*.safetensors")))
             if sf_files:
                 for sf_path in sf_files:
                     yield from load_file(sf_path, device="cpu").items()
                 return
 
-        # HF hub fallback: download only the talker checkpoint files
-        model_root = download_weights_from_hf_specific(
-            model_path,
-            self.vllm_config.load_config.download_dir,
-            allow_patterns=["talker/model*.safetensors"],
-        )
-        talker_dir = os.path.join(model_root, "talker")
-        sf_files = sorted(glob_module.glob(os.path.join(talker_dir, "model*.safetensors")))
-        if not sf_files:
-            raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.")
-        for sf_path in sf_files:
-            yield from load_file(sf_path, device="cpu").items()
+        raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.")
 
     def _load_vae_weights(self) -> set[str]:
         """Load AudioVAE weights from talker/vae/model.safetensors."""

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -64,6 +64,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config
 from vllm_omni.transformers_utils.processors.ming import (
@@ -715,12 +716,8 @@ class MingFlashOmniThinkerForConditionalGeneration(
 
         p = Path(model_path) / "mlp" / "model.safetensors"
         if not p.exists() and not os.path.exists(model_path):
-            from vllm_omni.model_executor.model_loader.weight_utils import (
-                download_weights_from_hf_specific,
-            )
-
             try:
-                local_path = download_weights_from_hf_specific(model_path, None, ["mlp/*"])
+                local_path = resolve_model_to_local_path(model_path, allow_download=True, allow_patterns=["mlp/*"])
                 p = Path(local_path) / "mlp" / "model.safetensors"
             except Exception:
                 logger.exception("[MingFlashOmniThinker] failed to download mlp/ from %s", model_path)

vllm_omni/model_executor/models/ming_flash_omni/talker_module.py

@@ -33,9 +33,14 @@ from vllm.platforms import current_platform
 from x_transformers.x_transformers import RotaryEmbedding
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
-from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps
-from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps
+from vllm_omni.model_executor.models.common.ming.cfm_solver import (
+    apply_sway_sampling,
+    get_epss_timesteps,
+    integrate_cfm_steps,
+)
+from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer
 
 logger = init_logger(__name__)
 
@@ -458,58 +463,6 @@ def silence_holder(
     return speech_out, sil_cache
 
 
-########################################################################
-# Audio Postprocess
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/aggregator.py
-########################################################################
-
-
-class Aggregator(nn.Module):
-    """Maps generated audio latent patches back to LLM embedding space."""
-
-    def __init__(
-        self,
-        in_channels: int = 64,
-        hidden_size: int = 1152,
-        depth: int = 28,
-        num_heads: int = 16,
-        mlp_ratio: float = 4.0,
-        llm_input_dim: int = 896,
-        **kwargs,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.num_heads = num_heads
-
-        self.word_embedder = nn.Embedding(1, hidden_size)
-        self.x_embedder = nn.Linear(in_channels, hidden_size)
-        self.hidden_size = hidden_size
-
-        self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)
-
-        self.blocks = nn.ModuleList(
-            [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)]
-        )
-        self.final_layer = FinalLayer(hidden_size, llm_input_dim)
-
-    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        x = self.x_embedder(x)
-        cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device))
-        x = torch.cat([cls_embed, x], dim=1)
-
-        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
-        if mask is not None:
-            mask_pad = mask.clone().detach()[:, :1]
-            mask = torch.cat([mask_pad, mask], dim=-1)
-        for block in self.blocks:
-            x = block(x, mask, rope)
-        x = self.final_layer(x)
-        x = x[:, :1, :]
-        return x
-
-
 ########################################################################
 # Prompt Builder
 # Adapted from:

vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py

@@ -16,16 +16,15 @@ import torch
 from transformers.utils.hub import cached_file
 from vllm.logger import init_logger
 
-from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
 
-from .spk_embedding import SpkembExtractor
 from .talker_module import resample
 
 if TYPE_CHECKING:
+    from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
     from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
 
-    from .talker_module import Aggregator
-
 logger = init_logger(__name__)
 
 
@@ -273,18 +272,19 @@ class VoicePresetRegistry:
             if os.path.isfile(path):
                 return path, candidate
 
-        if not os.path.isdir(self._model_path):
-            try:
-                hf_root = download_weights_from_hf_specific(
-                    self._model_path,
-                    self._download_dir,
-                    allow_patterns=["talker/data/**"],
-                    require_all=True,
-                )
-                candidate = os.path.join(hf_root, "talker", "data", "voice_name.json")
-                if os.path.isfile(candidate):
-                    return candidate, os.path.join(hf_root, "talker")
-            except Exception as e:  # pragma: no cover
-                logger.info("Could not download voice presets from HF: %s", e)
+        try:
+            hf_root = resolve_model_to_local_path(
+                self._model_path,
+                allow_download=True,
+                allow_patterns=["talker/data/**"],
+                cache_dir=self._download_dir,
+            )
+        except Exception as e:
+            logger.info("Could not download voice presets from HF: %s", e)
+            return None, None
+
+        candidate = os.path.join(hf_root, "talker", "data", "voice_name.json")
+        if os.path.isfile(candidate):
+            return candidate, os.path.join(hf_root, "talker")
 
         return None, None

vllm_omni/model_executor/models/ming_tts/fm/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -11,15 +11,21 @@ import torch
 import torch.nn as nn
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm.model_executor.models.bailing_moe import BailingMoeModel
 from vllm.model_executor.models.qwen2 import Qwen2Model
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 from vllm.sequence import IntermediateTensors
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
+from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import (
+    CFMGraphExecutor,
+    CFMSampler,
+)
+from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
-from .aggregator import Aggregator
 from .config_ming_tts import (
     KEY_CFG,
     KEY_DECODE_STEP,
@@ -34,7 +40,6 @@ from .config_ming_tts import (
     MingTTSConfig,
 )
 from .constants import SPEAKER_EMBEDDING_DIM
-from .flowloss_head import FlowLoss
 from .patch_emission import (
     MING_STOP_REASON_CODES,
     MING_STOP_REASON_CONTINUE,
@@ -82,8 +87,6 @@ class MingLLMModel(nn.Module):
         self.prefix = prefix
         self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config)
         if self.ming_config.model_variant == "moe":
-            from vllm.model_executor.models.bailing_moe import BailingMoeModel
-
             # BailingMoeModel reads ``vllm_config.model_config.hf_config`` directly
             # (no get_text_config()), so re-wrap with the nested bailing_moe config.
             llm_config = vllm_config.model_config.hf_config.get_text_config()
@@ -335,8 +338,6 @@ class MingLLMModel(nn.Module):
         if self._cfm_graph is not None:
             return self._cfm_graph
         try:
-            from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler
-
             # TODO(perf):
             #   (1) compile per-block (dit_model.blocks) instead of the
             #   whole forward for tighter fusion / fewer graph breaks;

vllm_omni/model_executor/models/ming_tts/prompt_encoder.py

@@ -12,7 +12,7 @@ import torch
 from safetensors import safe_open
 from vllm.logger import init_logger
 
-from vllm_omni.engine.stage_init_utils import _resolve_model_to_local_path
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
 
 from .audio_prep import (
@@ -78,7 +78,7 @@ def _load_prompt_encoder(wrapper: Any) -> AudioVAE:
         loaded_encoder_params = set()
         with torch.no_grad():
             for shard_path in _iter_model_safetensors(
-                _resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model))
+                resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model))
             ):
                 with safe_open(str(shard_path), framework="pt", device="cpu") as handle:
                     for key in handle.keys():