[Diffusion] Simplify DLO host plan setup

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/cuda/test-weekly.yml

@@ -30,14 +30,6 @@ steps:
       - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
     mirror_hardwares: h100_4
 
-  - label: "Reliability Test - voxcpm2"
-    timeout_in_minutes: 180
-    depends_on: upload-weekly-pipeline
-    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
-    commands:
-      - pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
-    mirror_hardwares: l4_1
-
   - group: "Reliability Test - Invalid parameters Test"
     depends_on: upload-weekly-pipeline
     if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"

docs/configuration/stage_configs.md

@@ -38,11 +38,12 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
 |-------|------|----------|---------|-------------|
 | `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
 | `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
-| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
+| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage total memory target; used for automatic KV-cache sizing. |
+| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. |
 | `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
 | `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
-| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
-| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
+| `max_num_batched_tokens` | int | optional | `32768` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
+| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
 | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
 | `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
 | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |

docs/contributing/ci/test_execution_guide.md

@@ -184,9 +184,6 @@ Failed jobs: 1/2
     # Reliability: HunyuanImage DiT (H100 × 4)
     pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
 
-    # Reliability: VoxCPM2 (L4)
-    pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
-
     ```
 
     The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml).

docs/contributing/ci/test_writing_guide.md

@@ -562,7 +562,6 @@ L5 level testing focuses on the performance of model services under ***long-runn
     -   `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
     -   `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
     -   `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
-    -   `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
     -   `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
     -   `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
     -   `README.md` — Short local run commands for this directory.
@@ -662,13 +661,12 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
 | `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
 | `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
 | `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
-| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
 | `README.md` | Minimal run / collect examples. |
 
 ##### Parametrization and markers
 
 - Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
-- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
+- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100).
 - **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
 
 ##### CI trigger
@@ -680,13 +678,11 @@ Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model sui
 | Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
 | Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
 | Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
-| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
 
 ```bash
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
 pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
 ```
 
 ##### Local commands
@@ -696,7 +692,6 @@ pytest --collect-only tests/dfx/reliability
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 ```
 
 ##### Adding a new model suite
@@ -709,4 +704,4 @@ pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 </details>
 
 -   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
-    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
+    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -10,6 +10,30 @@ For user-facing commands, see the
 [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
 and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
 
+## Feature compatibility
+
+Host-storage optimization and runtime compatibility are separate decisions.
+When direct checkpoint mmap is unavailable, DLO can still use tensors produced
+by the ordinary loader. "Compatibility path" below means that fallback is
+implemented but has less end-to-end coverage than the primary path.
+
+Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.
+
+| Feature | DLO + AllGather | DLO without AllGather |
+|---|---|---|
+| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. |
+| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. |
+| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
+| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
+| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
+| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
+| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
+| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
+
+See [Parallelism compatibility](#parallelism-compatibility) and
+[Request and loading constraints](#request-and-loading-constraints) for the
+detailed contracts and validation boundaries.
+
 ## Status
 
 DLO is implemented for multi-device diffusion execution. The default
@@ -19,21 +43,291 @@ 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. 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.
+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.
 
 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, but those configurations do not use checkpoint mmap and must not be
-used to claim shared-mmap host-memory savings.
+tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
+can share each final TP coordinate among equivalent processes.
 
-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.
+The compatibility matrix below defines the current implementation boundary.
+Hardware and performance qualification is tracked separately from this
+long-lived design contract.
 
-## Design
+## 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.
 
 ### DLO consumes the existing parallel topology
 
@@ -41,15 +335,10 @@ 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.
 
-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.
+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 loader owns host-weight planning
 
@@ -81,7 +370,24 @@ 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 path
+## 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.
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
 `1 / group_size` of each streamable block in pinned host memory. The next
@@ -89,6 +395,15 @@ 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]
@@ -114,11 +429,29 @@ 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.
 
-### Rank-local path without DLO AllGather
+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
 
 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
-either a loader-approved checkpoint mapping or ordinary runtime tensors.
+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.
 
 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
@@ -127,10 +460,17 @@ 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.
 
-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.
+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.
 
 This mode means:
 
@@ -141,18 +481,109 @@ 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 ordinary-loader fallback keeps one private runtime
-  copy per rank.
+  mmap is selected. The host weight cache can share equivalent ordinary-loader
+  outputs; its 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 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. |
+| **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. |
 
 ### Combined dimensions
 
@@ -168,7 +599,9 @@ This mode means:
 - **HSDP + DP or TP:** rejected independently by the diffusion parallel
   configuration.
 
-## Request and loading constraints
+## Mode constraints
+
+### AllGather constraints
 
 AllGather DP multi-concurrency requires:
 
@@ -176,104 +609,54 @@ 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.
+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
 
 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. Online quantization remains incompatible with DLO AllGather;
-use `--dlo-no-use-allgather` or disable online quantization.
+ordinary loader.
 
-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).
+## Validation status
 
-## 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;
-- 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.
+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).
 
 ## Recommendations
 
-- Use **DP + DLO AllGather** for the supported throughput and host-memory
-  scaling path.
+### 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 **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; 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.
+  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.

docs/design/feature/quantization.md

@@ -190,12 +190,15 @@ configured parallelism strategy:
 - `WeightsMapper` transformations must be applied before matching ignored-layer
   lists or component prefixes.
 
-Distributed layerwise offload has an additional boundary: its sharded mmap
-AllGather path is not compatible with online quantization because the path
-expects already materialized checkpoint weights. Use the rank-local
-`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
-combining these features. See the [DLO design document](offloader/distributed_layerwise_offload.md)
-for the full compatibility matrix.
+Distributed layerwise offload has an additional boundary: direct checkpoint
+mmap cannot prepare online-quantized tensors because it bypasses the ordinary
+loader that creates their runtime weights and scales. Per-tensor online FP8
+linears can still use DLO AllGather through the ordinary loader: each rank first
+finalizes the FP8 tensors, then DLO shards those runtime tensors for
+reconstruction during prefetch. Other online methods remain limited to
+`--dlo-no-use-allgather` until their packing and scale layouts are validated.
+See the [DLO design document](offloader/distributed_layerwise_offload.md) for
+the full compatibility matrix.
 
 ## Adding a quantization backend
 

docs/design/figures/dlo/dlo_allgather_sharded_host.svg

@@ -0,0 +1,165 @@
+<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

@@ -0,0 +1,202 @@
+<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/design/module/vllm_omni_config.md

@@ -18,8 +18,8 @@ validation_paths:
   - tests/config/**
 upstream_refs:
   - vllm.config
-last_reviewed: 2026-08-07
-last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+last_reviewed: 2026-08-11
+last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f
 ---
 
 # vLLM-Omni configuration
@@ -39,11 +39,66 @@ The PR review proposal to capture environment-derived values in configuration
 objects at initialization is an open design question, not a current
 invariant.
 
+## Current structured-stage reuse behavior
+
+The current RFC #4021 implementation uses upstream vLLM config classes as the
+schema source for AR and generation stages:
+
+- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`.
+- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`.
+- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`.
+- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`.
+
+This inheritance preserves the existing load, cache, scheduler, and parallel
+concern boundaries while reducing duplicate downstream declarations. It does
+not, by itself, make every inherited field an effective vLLM-Omni runtime
+option. Construction and runtime consumption are separate surfaces:
+
+| Surface | Current behavior |
+| --- | --- |
+| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. |
+| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. |
+| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. |
+| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. |
+
+The generated projection maps include known upstream naming differences such
+as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and
+`data_parallel_master_ip` to `data_parallel_address`. A field added to both an
+upstream concern config and `EngineArgs` therefore enters the AR/generation
+projection without requiring a second downstream allowlist. Explicit inherited
+fields that have no projection raise an error rather than being silently
+dropped.
+
+Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`,
+cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns
+`distributed_executor_backend` and `worker_cls`; vLLM's private API-process
+fields are terminal internals. These exclusions prevent one input from being
+constructed or projected through two config concerns. Effective-engine-argument
+tests cover both the dynamic mapping and this exclusion boundary.
+
+`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields:
+`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML
+inputs and complete upstream value objects. Mapping inputs are validated and
+materialized as their concrete upstream types when the structured config is
+constructed, so downstream consumers see one resolved representation.
+Prebuilt upstream objects retain their type across the structured projection
+boundary.
+
+Quantization remains on its existing Omni-owned transport contract. This
+reuse step does not adopt upstream `QuantizationConfigArgs`, because the
+current engine-specific quantization split needs to be resolved separately.
+
+Diffusion keeps its existing engine-field ownership and projection surface.
+Shared Python inheritance must not be treated as evidence that an LLM-only
+parallel input is effective for a diffusion stage.
+
 ## Safe-change guide while deferred
 
 Changes should trace every affected producer to its runtime consumer and test
 the structured, legacy, CLI, and deployment paths that are actually supported.
-Do not infer a stable contract from this placeholder page.
+For inherited vLLM fields, tests should separately cover constructor behavior
+and effective `EngineArgs` projection. Do not infer a stable contract from
+this placeholder page.
 
 ## Promotion gate
 

docs/models/supported_models.md

@@ -78,6 +78,7 @@ th {
 | `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
 | `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
 | `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
+| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -1,153 +1,127 @@
 # Distributed Layerwise Offloading
 
-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.
+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.
 
-See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
-for the implementation contract and compatibility matrix.
+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).
 
-## Execution model
+## Choose a mode
 
-DLO overlaps three operations with a fixed two-block device buffer:
+| 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 |
 
-```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.
+AllGather is normally the best choice for synchronized DP. Host weight cache mode
+targets independently scheduled replicas, especially repeated TP engines on
+one node.
 
 ## Usage
 
 ```bash
-# Four ranks with sharded host weights and AllGather
+# Sharded host weights with DLO AllGather
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --data-parallel-size 4
 
-# Standard-loader rank-local weights, without DLO AllGather
+# Independently streamed rank-local weights
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
-  --data-parallel-size 4 \
   --dlo-no-use-allgather
 
-# Sequence parallel deployment
+# Share final host weights; prefer registered direct H2D under the pinned-memory policy
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
-  --usp 4
+  --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
 ```
 
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="/path/to/model",
-    enable_distributed_layerwise_offload=True,
-    dlo_use_allgather=True,
-)
-```
-
-## Flags
+## Relevant options
 
 | Flag | Meaning | Default |
 | --- | --- | --- |
 | `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
-| `--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` |
+| `--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` |
 
-## Host-weight loading
+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.
 
-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.
+## Operational notes
 
-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.
+- 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 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.
+## Scheduling constraints
 
-With direct checkpoint mmap, the loader:
+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`.
 
-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.
+No-AllGather workers do not have this DLO lockstep requirement. Ordinary TP or
+SP model collectives still synchronize ranks within each engine.
 
 ## Limitations
 
-- 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.
-- Online quantization uses the ordinary loader with no-AllGather. It remains
-  incompatible with DLO AllGather.
-- 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.
+- 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.
+- 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).
 
 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/fault_injection_reliability_matrix.md

@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
 
 ## Source of Conclusions
 
-The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
+The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, and **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
 
 ### Example: `SIGTERM` fault injection log (Qwen3-Omni)
 

docs/user_guide/quantization/fp8.md

@@ -179,6 +179,24 @@ python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers
 vllm serve <your-model> --omni --quantization fp8
 ```
 
+### Distributed layerwise offload
+
+Per-tensor online FP8 linears can be combined with DLO's default AllGather path:
+
+```bash
+vllm serve <your-model> \
+  --omni \
+  --quantization fp8 \
+  --enable-distributed-layerwise-offload
+```
+
+Direct checkpoint mmap is not used for this combination. Each rank runs the
+ordinary loader to create the final FP8 weights and scales before DLO shards
+them, so startup temporarily materializes the complete FP8 model in host memory
+on every rank. Runtime DLO host residency is sharded normally. Use
+`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract
+does not fit the workload.
+
 ## Parameters
 
 | Parameter | Type | Default | Description |

recipes/MiniMaxAI/MiniMax-H3.md

@@ -416,9 +416,10 @@ For example, keep the first main block's attention projections in BF16 with:
 ```
 
 The structured option replaces `--quantization fp8`. Online FP8 can be used
-with H3 layerwise offload and with distributed layerwise offload's full-weight
-per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
-supported for runtime-created FP8 weights.
+with H3 layerwise offload and with both DLO transfer paths. The default
+AllGather path uses the ordinary loader to finalize FP8 weights and scales
+before sharding them across ranks. `--dlo-no-use-allgather` instead retains
+complete rank-local tensors and avoids the synchronized request-wave contract.
 
 ## AMD ROCm (gfx942 / gfx950)
 
@@ -870,8 +871,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
   H3 native `tile` mode only.
 - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
   pure Ulysses.
-- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
-  use `--dlo-no-use-allgather` when combining online FP8 with DLO.
+- Online FP8 with DLO AllGather temporarily materializes the complete FP8 model
+  in host memory on every rank during startup before retaining only each rank's
+  shard. Size startup host memory for that transient peak.
 - TeaCache and Cache-DiT cannot be enabled on the same server.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone

recipes/MiniMaxAI/MiniMax-Music3.md

@@ -0,0 +1,173 @@
+# MiniMax Music 3
+
+> Text-to-music: lyrics plus a style caption in, a 32 kHz stereo song out
+
+## Summary
+
+- Vendor: MiniMax
+- Model: `MiniMaxAI/MiniMax-Music3`
+- Task: Text-to-music generation (vocals and instrumental together)
+- Mode: Online serving through the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: @linyueqian
+
+## When to use this recipe
+
+Use it to serve `MiniMaxAI/MiniMax-Music3` for full-song generation up to six
+minutes. This is not a TTS model with a musical voice: there is no reference
+speaker, no `voice`, and no `temperature`. What you control is the lyrics, the
+caption, the seed, and the length.
+
+## References
+
+- Model card: [MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3)
+- Demo samples: [MiniMax Music 3 demo](https://minimax-ai.github.io/music3-demo/)
+- Deploy config: `vllm_omni/deploy/minimax_music3.yaml` (single GPU),
+  `vllm_omni/deploy/minimax_music3_2gpu.yaml` (split across two)
+
+## Architecture
+
+| Component | Spec |
+|---|---|
+| Backbone | Qwen3 decoder, 36 layers, hidden 4096, GQA 32/8, vocab 200k |
+| Depth decoder | 4 layers, hidden 4096, 16 heads; 7 residual codebooks of 1024 |
+| Frame | 8 codebooks; `c0` in the backbone vocabulary, `c1..c7` in the depth decoder |
+| AR guidance | Classifier-free, scale 1.5, `c0` masked to the conditioned branch's top 50 |
+| Acoustic stage | Flow-matching transformer (36 layers, dim 2048) plus a DAC-style decoder, 512x upsampling |
+| Solver | Euler, 30 steps, acoustic guidance scale 1.7 |
+| Window | 200 frames per acoustic chunk, 100-frame hop |
+| Frame rate | 25 frames per second |
+| Context length | 10,240 tokens |
+| Output | 32 kHz stereo WAV |
+
+Generation is two-stage. The backbone predicts one audio frame per decode step,
+each frame eight RVQ codebooks deep. Every 200 frames are handed to the
+flow-matching transformer, which solves a latent that the decoder turns into
+waveform.
+
+**Guidance doubles the rows.** Every request decodes twice: one row sees the
+real prompt, the other sees the same prompt with the caption-and-lyrics span
+replaced by `<|audio_cfg|>`. Both rows hold their own KV cache for the whole
+song, so a request costs twice the KV its length suggests. Size the pool for
+rows, not requests.
+
+## Hardware Support
+
+## GPU
+
+### 1x H200 141GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- CUDA: 12.9
+- vLLM version: 0.27.0
+- vLLM-Omni version or commit: branch `feat/minimax-music3`
+
+#### Command
+
+```bash
+vllm serve MiniMaxAI/MiniMax-Music3 \
+    --host 0.0.0.0 --port 8000 \
+    --trust-remote-code --omni
+```
+
+The default deploy config `vllm_omni/deploy/minimax_music3.yaml` is discovered
+automatically from the checkpoint's `model_type=minimax_music3`. It colocates
+both stages on device 0.
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "MiniMaxAI/MiniMax-Music3",
+    "input": "[Verse]\nWalking down the empty street at midnight\nStreetlights flicker like a broken dream\n[Chorus]\nAnd I keep on walking\nTill the morning finds me",
+    "instructions": "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. Intimate bedroom production, gentle tape saturation, no bright cymbals.",
+    "seed": 1,
+    "max_new_tokens": 750
+  }' \
+  --output song.wav
+```
+
+Expected: a 30 second 32 kHz stereo WAV. `max_new_tokens` counts audio frames
+at 25 per second, so 750 frames is at most 30 seconds. It is a cap, not a
+target: the model ends the song itself when it emits the audio-end token, and a
+response shorter than the cap is the model finishing rather than a truncation.
+
+#### Notes
+
+- **Memory:** stage 0 takes `gpu_memory_utilization: 0.6` for the bf16 backbone
+  and its KV pool; stage 1 takes 0.2 for the transformer and decoder weights,
+  which sit outside any KV budget.
+- **Key flags:** none required. Guidance, the solver step count and the window
+  geometry are fixed by the checkpoint and are not request parameters.
+- **Known limitations:**
+  - The external API is non-streaming; `stream: true` is rejected.
+  - `max_model_len` is 10,240, so a maximum-length prompt (5,000 tokens) and a
+    maximum-length song (9,000 frames) cannot both fit. Long captions shorten
+    the maximum song. This is a property of the checkpoint.
+  - The acoustic stage runs in float32. bfloat16 is accepted but measurably
+    degrades the solver.
+
+### 2x H200 141GB
+
+Use `minimax_music3_2gpu.yaml` to put the acoustic stage on the second device:
+
+```bash
+vllm serve MiniMaxAI/MiniMax-Music3 \
+    --host 0.0.0.0 --port 8000 \
+    --trust-remote-code --omni \
+    --deploy-config minimax_music3_2gpu.yaml
+```
+
+Only the placement differs; both layouts run the acoustic stage in float32.
+
+## Writing the request
+
+Two fields carry the request, and both are required.
+
+- `input` is the **lyrics**. Structure tags (`[Verse]`, `[Chorus]`, `[Bridge]`,
+  `[Outro]`) steer the arrangement and are part of the prompt contract.
+  **Put a tag on its own line.** Normalization keeps only the tags on a line
+  that starts with one and drops the rest of that line, so lyrics written next
+  to a tag are silently lost:
+
+  ```text
+  "[Verse]\nWalking down the street"   ->   [start] [verse] Walking down the street
+  "[Verse] Walking down the street"    ->   [start] [verse]
+  ```
+
+- `instructions` is the **caption** describing genre, instrumentation, tempo,
+  mood and production. It is the strongest control available: vague captions
+  give generic arrangements. Long structured captions work well, since the
+  model was trained on descriptions covering global attributes, emotional
+  progression, vocal detail and arrangement.
+
+A request is deterministic in its seed: the same lyrics, caption, seed and
+length return the same audio. Changing only the seed gives another take of the
+same song.
+
+## Request parameters
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `input` | string | required | Lyrics, non-empty. Tags on their own lines |
+| `instructions` | string | required | Caption, non-empty |
+| `seed` | int | `0` | Non-negative. Fixes the output for a given request |
+| `max_new_tokens` | int | `9000` | Cap on audio frames, 25 per second. Maximum 9,000 |
+| `response_format` | string | `"wav"` | Output container |
+
+## Parameters the model rejects
+
+The request is refused rather than silently ignored, so a mistake is visible
+immediately.
+
+| Parameter | Reason |
+|---|---|
+| `temperature`, `top_p`, `top_k`, `repetition_penalty` | Sampling is fixed: guidance at 1.5, then a seeded top-k 50 draw |
+| `voice` | There is no speaker to select; the vocal comes from the caption |
+| `ref_audio`, `ref_text`, `language`, `task_type` | No reference-audio conditioning and no language tag |
+| `speed` | Only `1.0`. Tempo belongs in the caption, e.g. "at 92 BPM" |
+| `stream: true` | The external API is non-streaming |

recipes/cosmos3/Cosmos3-DistOffload.md

@@ -189,8 +189,10 @@ wait
 
 - `--dlo-use-allgather` (default) requires all concurrent requests to have the same
   `num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
-- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
-  use `--dlo-no-use-allgather` or disable online quantization
+- Online per-tensor FP8 linears with DLO AllGather use the ordinary loader
+  instead of direct checkpoint mmap, then shard the finalized FP8 weights and
+  scales. Size host memory for a transient complete FP8 model on every rank at
+  startup
 - Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
   `--dlo-no-use-allgather` retains the standard TP loader and should be treated
   as an experimental compatibility path

tests/benchmarks/test_daily_omni_local_source.py

@@ -9,13 +9,17 @@ from pathlib import Path
 import pytest
 
 from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
+    DailyOmniDataset,
     daily_omni_local_qa_json,
     daily_omni_local_videos_dir,
+    ensure_daily_omni_hub_root,
     resolve_daily_omni_local_root,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
 
+_HUB_ID = "liarliar/Daily-Omni"
+
 
 def _write_qa_json(root: Path) -> Path:
     qa = root / "qa.json"
@@ -23,6 +27,28 @@ def _write_qa_json(root: Path) -> Path:
     return qa
 
 
+def _stage_offline_hub_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, with_qa: bool = True) -> Path:
+    """Build a hub cache holding ``liarliar/Daily-Omni`` and cut off all network access.
+
+    Reproduces the air-gapped NPU runner: the dataset is pre-staged with ``hf download``, so it
+    only exists in the *hub* cache -- never in the ``datasets`` cache that ``load_dataset`` reads.
+    """
+    constants = pytest.importorskip("huggingface_hub").constants
+
+    cache = tmp_path / "hub"
+    repo_dir = cache / "datasets--liarliar--Daily-Omni"
+    snapshot = repo_dir / "snapshots" / "deadbeef"
+    snapshot.mkdir(parents=True)
+    if with_qa:
+        _write_qa_json(snapshot)
+    (repo_dir / "refs").mkdir()
+    (repo_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
+
+    monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache))
+    monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
+    return snapshot
+
+
 def test_resolve_local_root_returns_none_for_hub_id() -> None:
     assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
     assert resolve_daily_omni_local_root("") is None
@@ -90,3 +116,37 @@ def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
     root = resolve_daily_omni_local_root(str(tmp_path))
     assert daily_omni_local_qa_json(root) is None
     assert daily_omni_local_videos_dir(root) is None
+
+
+def test_hub_id_resolves_to_offline_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
+
+    root = ensure_daily_omni_hub_root(_HUB_ID)
+
+    assert root == snapshot.resolve()
+    assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
+
+
+def test_dataset_loads_hub_id_without_network(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """The regression this file exists for: a Hub id must not require ``load_dataset``."""
+    snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
+    # Any use of `datasets` here means we went back to the path that cannot read the hub cache.
+    monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
+
+    dataset = DailyOmniDataset(dataset_path=_HUB_ID, random_seed=0)
+
+    assert dataset.qa_json_path == (snapshot / "qa.json").resolve()
+    assert [row["video_id"] for row in dataset.data] == ["vid0"]
+
+
+def test_uncached_hub_id_reports_both_attempts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    _stage_offline_hub_cache(tmp_path, monkeypatch, with_qa=False)
+    monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        DailyOmniDataset(dataset_path="someone-else/Daily-Omni", random_seed=0)
+
+    message = str(excinfo.value)
+    assert "someone-else/Daily-Omni" in message
+    assert "snapshot_download failed" in message
+    assert "--daily-omni-qa-json" in message

tests/config/test_config_import_cycle.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Guard the NPU pytest circular import through DiffusionOutput."""
+
+from __future__ import annotations
+
+import ast
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def _top_level_import_modules(path: Path) -> set[str]:
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    modules: set[str] = set()
+    for node in tree.body:
+        if isinstance(node, ast.ImportFrom) and node.module:
+            modules.add(node.module)
+        elif isinstance(node, ast.Import):
+            modules.update(alias.name for alias in node.names)
+    return modules
+
+
+def _assert_isolated_import_succeeds(script: str) -> None:
+    env = os.environ.copy()
+    env.pop("CUDA_VISIBLE_DEVICES", None)
+    env.pop("HIP_VISIBLE_DEVICES", None)
+    result = subprocess.run(
+        [sys.executable, "-c", script],
+        cwd=REPO_ROOT,
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=120,
+        check=False,
+    )
+    assert result.returncode == 0, result.stderr
+
+
+def test_config_package_does_not_eagerly_import_pipeline_registry() -> None:
+    imported = _top_level_import_modules(REPO_ROOT / "vllm_omni" / "config" / "__init__.py")
+    assert "vllm_omni.config.config_factory" not in imported
+    assert "vllm_omni.config.pipeline_registry" not in imported
+
+
+def test_config_package_lazy_exports_still_resolve() -> None:
+    from vllm_omni.config import StageConfigFactory, register_pipeline
+    from vllm_omni.config.config_factory import StageConfigFactory as DirectStageConfigFactory
+    from vllm_omni.config.pipeline_registry import register_pipeline as direct_register_pipeline
+
+    assert StageConfigFactory is DirectStageConfigFactory
+    assert register_pipeline is direct_register_pipeline
+
+
+def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
+    src = (REPO_ROOT / "vllm_omni" / "platforms" / "npu" / "platform.py").read_text(encoding="utf-8")
+    init_start = src.index("def __init__(self) -> None:")
+    runtime_start = src.index("def init_diffusion_model_runner_runtime")
+    init_src = src[init_start:runtime_start]
+    runtime_src = src[runtime_start:]
+    assert "apply_minimax_h3_qwen3vl_patch" not in init_src
+    assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
+    assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
+
+
+def test_lora_config_import_does_not_load_pipeline_registry() -> None:
+    """Import LoRAConfig and diffusion.data in a fresh interpreter.
+
+    The NPU failure was ``data.py`` importing LoRAConfig while still loading,
+    then ``vllm_omni.config`` re-entering ``pipeline_registry`` → PI0 →
+    DiffusionOutput. This child process must not preload those exports.
+    """
+    _assert_isolated_import_succeeds(
+        """
+import sys
+
+from vllm_omni.config.lora import LoRAConfig
+
+if "vllm_omni.config.pipeline_registry" in sys.modules:
+    raise SystemExit("LoRAConfig import loaded pipeline_registry")
+if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
+    raise SystemExit("LoRAConfig import loaded PI0_PIPELINE")
+
+from vllm_omni.diffusion.data import DiffusionOutput
+
+if DiffusionOutput is None or LoRAConfig is None:
+    raise SystemExit("expected LoRAConfig and DiffusionOutput")
+"""
+    )
+
+
+def test_pipeline_registry_does_not_import_pi0_runtime() -> None:
+    """Resolving the pi0 topology must not import its runtime pipeline.
+
+    ``pipeline_pi0`` imports ``diffusion.data``. Importing it while
+    ``diffusion.data`` is only partially initialized closes the Ascend plugin's
+    circular-import loop, so the registry must depend only on a lightweight
+    pipeline-config module.
+    """
+    _assert_isolated_import_succeeds(
+        """
+import sys
+
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+
+if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
+    raise SystemExit("pipeline_registry imported the pi0 runtime pipeline")
+
+pipeline = resolve_pipeline_config("pi0")
+if pipeline is None or pipeline.model_type != "pi0":
+    raise SystemExit("expected the registered pi0 pipeline config")
+"""
+    )

tests/config/test_omni_config.py

@@ -5,11 +5,20 @@
 from __future__ import annotations
 
 from dataclasses import fields
+from inspect import Parameter, signature
 from pathlib import Path
 
+import msgspec
 import pytest
 from pydantic import ValidationError
+from pydantic.fields import FieldInfo
 from transformers import Qwen3OmniMoeConfig
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
 
 from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.config import omni_config as omni_config_module
@@ -40,6 +49,7 @@ from vllm_omni.config.stage_config import (
     merge_pipeline_deploy,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
 from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -243,9 +253,29 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
     ],
     ids=["global", "ar-stage", "generation-stage"],
 )
-def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
-    with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
-        _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
+def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id):
+    stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id)
+
+    assert stage.cache_config.cache_dtype == "fp8"
+    assert "cache_dtype" in stage.cache_config._omni_explicit_fields
+
+
+@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"])
+def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    deploy = DeployConfig(
+        stages=[
+            StageDeployConfig(
+                stage_id=0,
+                engine_extras={"parallel_config": {field_name: "10.0.0.1"}},
+            )
+        ]
+    )
+
+    stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
+
+    assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1"
+    assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
 
 
 def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
@@ -437,6 +467,8 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
         "connector_config",
         "runtime_config",
         "parallel_config",
+        "compilation_config",
+        "profiler_config",
         "quantization_config",
     }
     assert "diffusion_config" not in public_fields
@@ -455,10 +487,47 @@ def test_runtime_config_fields_match_structured_runtime_scope():
         "num_gpus",
         "log_level",
         "log_stats",
-        "profiler_config",
     }
 
 
+def test_upstream_config_compatible_fields_materialize_mapping_inputs():
+    compilation_config = {"backend": "eager"}
+    profiler_config = {"profiler": "cuda"}
+
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+        compilation_config=compilation_config,
+        profiler_config=profiler_config,
+    )
+
+    assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
+    assert stage_config.compilation_config.backend == "eager"
+    assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
+    assert stage_config.profiler_config.profiler == "cuda"
+
+
+def test_upstream_config_compatible_fields_validate_mapping_inputs():
+    with pytest.raises(ValidationError):
+        VllmOmniARStageConfig(
+            stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+            profiler_config={"delay_iterations": -1},
+        )
+
+
+def test_upstream_config_compatible_fields_keep_prebuilt_config_types():
+    compilation_config = VllmCompilationConfig(backend="eager")
+    profiler_config = VllmProfilerConfig(profiler="cuda")
+
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+        compilation_config=compilation_config,
+        profiler_config=profiler_config,
+    )
+
+    assert stage_config.compilation_config is compilation_config
+    assert stage_config.profiler_config is profiler_config
+
+
 def test_sub_config_fields_match_structured_scopes():
     assert {f.name for f in fields(OmniStageModelConfig)} == {
         "model",
@@ -489,7 +558,6 @@ def test_sub_config_fields_match_structured_scopes():
         "enforce_eager",
         "max_cudagraph_capture_size",
         "enable_flashinfer_autotune",
-        "compilation_config",
         "enable_multithread_weight_load",
         "num_weight_load_threads",
         "disable_autocast",
@@ -499,29 +567,25 @@ def test_sub_config_fields_match_structured_scopes():
         "model_subdir",
         "tokenizer_subdir",
     }
-    assert {f.name for f in fields(OmniStageLoadConfig)} == {
+    vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
+    assert issubclass(OmniStageLoadConfig, VllmLoadConfig)
+    assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | {
         "tokenizer",
-        "download_dir",
         "skip_tokenizer_init",
-        "load_format",
         "tokenizer_mode",
         "config_format",
         "skip_mm_profiling",
     }
-    assert {f.name for f in fields(OmniStageCacheConfig)} == {
-        "kv_cache_memory_bytes",
-        "gpu_memory_utilization",
-        "enable_prefix_caching",
+    assert OmniStageLoadConfig(load_format="PT").load_format == "pt"
+    assert issubclass(OmniStageCacheConfig, VllmCacheConfig)
+    assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | {
         "disable_hybrid_kv_cache_manager",
         "mm_processor_cache_gb",
         "mamba_ssm_cache_dtype",
     }
-    assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
-        "max_num_seqs",
-        "max_num_batched_tokens",
+    assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig)
+    assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | {
         "max_model_len",
-        "enable_chunked_prefill",
-        "async_scheduling",
     }
     assert {f.name for f in fields(OmniStageConnectorConfig)} == {
         "async_chunk",
@@ -530,17 +594,10 @@ def test_sub_config_fields_match_structured_scopes():
         "output_connectors",
         "input_connectors",
     }
-    assert {f.name for f in fields(OmniStageParallelConfig)} == {
-        "pipeline_parallel_size",
-        "data_parallel_size",
-        "tensor_parallel_size",
-        "enable_expert_parallel",
-        "world_size",
-    }
-    assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
-        "pipeline_parallel_size",
-        "data_parallel_size",
-        "tensor_parallel_size",
+    vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
+    assert issubclass(OmniStageParallelConfig, VllmParallelConfig)
+    assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields
+    assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | {
         "sequence_parallel_size",
         "ulysses_degree",
         "ring_degree",
@@ -554,20 +611,86 @@ def test_sub_config_fields_match_structured_scopes():
         "mask_sp_padding",
         "hsdp_shard_size",
         "hsdp_replicate_size",
-        "enable_expert_parallel",
-        "world_size",
     }
 
 
+def test_inherited_sub_configs_initialize_transport_safe_derived_fields():
+    scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096)
+    assert scheduler_config.max_num_encoder_input_tokens == 4096
+    assert scheduler_config.encoder_cache_size == 4096
+
+    unresolved_scheduler_config = OmniStageSchedulerConfig()
+    assert unresolved_scheduler_config.max_num_encoder_input_tokens is None
+    assert unresolved_scheduler_config.encoder_cache_size is None
+
+    parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2)
+    assert parallel_config.data_parallel_index == 2
+
+    diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2)
+    assert diffusion_parallel_config.data_parallel_index == 2
+
+
+@pytest.mark.parametrize(
+    "config_cls",
+    [
+        OmniStageLoadConfig,
+        OmniStageCacheConfig,
+        OmniStageSchedulerConfig,
+        OmniStageParallelConfig,
+        OmniStageDiffusionParallelConfig,
+    ],
+)
+def test_inherited_sub_configs_are_keyword_only(config_cls):
+    assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values())
+
+    with pytest.raises(TypeError):
+        config_cls(1)
+
+
+@pytest.mark.parametrize(
+    "config_cls",
+    [
+        OmniStageLoadConfig,
+        OmniStageCacheConfig,
+        OmniStageSchedulerConfig,
+        OmniStageParallelConfig,
+        OmniStageDiffusionParallelConfig,
+    ],
+)
+def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls):
+    config = config_cls()
+
+    uninitialized_fields = [
+        config_field.name
+        for config_field in fields(config)
+        if isinstance(getattr(config, config_field.name), FieldInfo)
+    ]
+    assert uninitialized_fields == []
+
+    msgspec.msgpack.encode(_serialize_stage_config(config))
+
+
+def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe():
+    omni_config = _from_pipeline_key("qwen3_tts")
+
+    for stage_config in omni_config.stage_configs:
+        payload = {
+            "stage_id": stage_config.stage_id,
+            "stage_config": _serialize_stage_config(stage_config),
+        }
+        msgspec.msgpack.encode(payload)
+
+
 def test_diffusion_parallel_config_fields_cover_legacy_surface():
     from vllm_omni.diffusion.data import DiffusionParallelConfig
 
     legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
     structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
     expected_upstream_fields = {"mask_sp_padding"}
+    vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
 
     assert legacy_fields | expected_upstream_fields <= structured_fields
-    assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
+    assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set()
 
 
 def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
@@ -785,6 +908,52 @@ 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]
@@ -1078,3 +1247,21 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa
     cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
 
     assert cfg.quantization_config == quantization_config
+
+
+def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    quantization_config = {"method": "int8", "activation_scheme": "dynamic"}
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+        quantization_config=quantization_config,
+    )
+
+    # Exercise the terminal construction performed by the future typed startup
+    # path without probing ports or loading remote model metadata.
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+    monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
+    cfg.enrich_config()
+
+    assert cfg.quantization_config is not None
+    assert cfg.quantization_config.get_name() == "int8"

tests/dfx/reliability/README.md

@@ -7,12 +7,10 @@ The detailed RFC document is maintained outside this repository (local/internal
 - Qwen3 tests: `test_reliability_qwen3_omni.py`
 - Wan2.2 tests: `test_reliability_wan22.py`
 - HunyuanImage tests: `test_reliability_hunyuan_image.py`
-- VoxCPM2 tests: `test_reliability_voxcpm2.py`
 
 ```bash
 pytest --collect-only tests/dfx/reliability
 pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
 pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
 ```

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -976,3 +976,68 @@ def test_voices_delete_invalid_requests(
             "err_message": err_message,
         }
     )
+
+
+# ─── POST /v1/audio/speech · MiniMax Music 3 ───
+
+# Text-to-music on the speech endpoint. Most of the speech contract does not
+# apply: there is no speaker, no reference audio and no temperature, and the
+# lyrics alone are not a request. Each case below pins a rejection so a caller
+# who assumes the usual TTS shape finds out immediately.
+_MINIMAX_MUSIC3_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="MiniMaxAI/MiniMax-Music3",
+            stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="minimax_music3",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight"
+_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor."
+# 25 frames per second, capped at six minutes of song.
+_MINIMAX_MUSIC3_MAX_FRAMES = 9000
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
+        pytest.param({"instructions": "   "}, ("instructions",), id="caption_blank"),
+        pytest.param({"input": "   "}, ("input",), id="lyrics_blank"),
+        pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
+        pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
+        pytest.param(
+            {"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1},
+            ("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)),
+            id="max_new_tokens_over_cap",
+        ),
+        pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True)
+def test_minimax_music3_speech_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body: dict[str, Any] = {
+        "model": omni_server.model,
+        "input": _MINIMAX_MUSIC3_LYRICS,
+        "instructions": _MINIMAX_MUSIC3_CAPTION,
+        "seed": 1,
+        "max_new_tokens": 250,
+        "response_format": "wav",
+    }
+    for key, value in overrides.items():
+        if value is None:
+            body.pop(key, None)
+        else:
+            body[key] = value
+    openai_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -71,7 +71,7 @@ def _default_oom_device_spec() -> str:
 
 OOM_INJECTION_CONFIG = {
     "device": _default_oom_device_spec(),
-    "target_mem_ratio": 0.92,
+    "target_mem_ratio": 0.95,
     "hold_seconds": 0,
     "startup_timeout_sec": 20,
     "strict": False,

tests/dfx/reliability/test_reliability_voxcpm2.py

@@ -1,508 +0,0 @@
-"""
-VoxCPM2 reliability integration tests.
-
-Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``.
-"""
-
-from __future__ import annotations
-
-import concurrent.futures
-import os
-import time
-from pathlib import Path
-from typing import Any
-
-import pytest
-
-from tests.dfx.conftest import (
-    assert_fault_exception,
-    create_reliability_omni_server_params,
-    resolve_oom_device_spec,
-)
-from tests.dfx.reliability.helpers import (
-    FaultInjector,
-    assert_no_server_tree_process_residual_and_gpu_release,
-    extract_openai_error_contract_from_bytes,
-    get_health_raw,
-    inject_gpu_oom,
-    make_process_kill_fault_injector,
-    make_server_root_kill_fault_injector,
-    make_server_tree_kill_fault_injector,
-    post_json_raw,
-    run_fault_injection_with_rate_load,
-    stop_gpu_oom_hogs,
-    worker_residual_timeout_after_kill_signal,
-)
-from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
-from tests.helpers.runtime import OpenAIClientHandler
-
-RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
-    {
-        "test_name": "voxcpm2_reliability_default",
-        "server_params": {
-            "model": "openbmb/VoxCPM2",
-            "stage_config_name": "voxcpm2.yaml",
-            # VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False
-            # would otherwise override deploy YAML settings during server startup.
-            "server_args": ["--trust-remote-code", "--disable-log-stats"],
-        },
-    },
-]
-
-DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
-OOM_INJECTION_CONFIG = {
-    "target_mem_ratio": 0.96,
-    "hold_seconds": 0,
-    "startup_timeout_sec": 20,
-    "strict": False,
-}
-FAULT_ERROR_KEYWORDS = (
-    "the request failed",
-    "oom",
-    "out of memory",
-    "cuda",
-    "orchestrator",
-    "timeout",
-    "connection",
-    "500",
-    "503",
-)
-PROCESS_KILL_ERROR_KEYWORDS = (
-    "timeout",
-    "did not complete within",
-    "connection",
-    "engine",
-    "orchestrator",
-    "dead",
-    "internal",
-    "500",
-    "503",
-)
-RUNTIME_WORKER_PATTERN = "VLLM::"
-SERVE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGINT", id="sigint"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-SERVE_WITH_LOAD_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGINT", id="sigint"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_WITH_LOAD_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-WORKER_SIGNAL_FAULT_PARAMS = [
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns=RUNTIME_WORKER_PATTERN,
-            signal_name="SIGTERM",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigterm",
-    ),
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns=RUNTIME_WORKER_PATTERN,
-            signal_name="SIGKILL",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigkill",
-    ),
-]
-
-VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
-INFLIGHT_INJECTION_REQUEST_RATE = 0.3
-INFLIGHT_INJECTION_REQUEST_COUNT = 10
-_SPEECH_FF_HTTP_TIMEOUT_SEC = 30
-_SPEECH_FF_MAX_ELAPSED_SEC = 30
-
-
-def _speech_request_config(omni_server: Any) -> dict[str, Any]:
-    return {
-        "model": omni_server.model,
-        "input": "The weather is nice today, perfect for a walk in the park.",
-        "stream": False,
-        "timeout": 120.0,
-        "response_format": "wav",
-        "voice": "default",
-    }
-
-
-def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
-    openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1)
-
-
-def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]:
-    base = "The weather is nice today, perfect for a walk in the park. "
-    return {
-        "model": omni_server.model,
-        "input": base * 80,
-        "max_new_tokens": 4096,
-        "ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"),
-        "stream": False,
-        "timeout": 300.0,
-        "response_format": "wav",
-        "voice": "default",
-    }
-
-
-def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
-    openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1)
-
-
-def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
-    payload = {
-        "model": model,
-        "input": "post-fault fast-fail check",
-        "voice": "default",
-    }
-    start = time.monotonic()
-    try:
-        status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20)
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s"
-        assert status >= 500, (
-            f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
-        )
-    except Exception:  # noqa: BLE001
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
-
-
-def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
-    deadline = time.monotonic() + 20.0
-    last_observation = ""
-    while time.monotonic() < deadline:
-        try:
-            status, body = get_health_raw(host, port, timeout_sec=5)
-            last_observation = f"http={status}, body={body[:200]!r}"
-            if status == 503:
-                return
-        except Exception as exc:  # noqa: BLE001
-            last_observation = f"exception={exc!r}"
-            return
-        time.sleep(0.5)
-    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skip(reason="issue#4285")
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_speech_request_failure(
-    omni_server_function,
-    openai_client_function,
-) -> None:
-    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
-    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
-    handle = inject_gpu_oom(
-        device=device_spec,
-        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
-        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
-        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
-        strict=OOM_INJECTION_CONFIG["strict"],
-    )
-    try:
-        try:
-            _submit_oom_speech_request(openai_client_function, omni_server_function)
-        except Exception as exc:
-            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
-        else:
-            pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection")
-    finally:
-        stop_gpu_oom_hogs(handle)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skip(reason="issue#4285")
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_speech_error_contract(
-    omni_server_function,
-) -> None:
-    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
-    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
-    handle = inject_gpu_oom(
-        device=device_spec,
-        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
-        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
-        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
-        strict=OOM_INJECTION_CONFIG["strict"],
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    oom_config = _oom_speech_request_config(omni_server_function)
-    payload = {
-        "model": oom_config["model"],
-        "input": oom_config["input"],
-        "voice": oom_config["voice"],
-        "max_new_tokens": oom_config["max_new_tokens"],
-        "ref_audio": oom_config["ref_audio"],
-    }
-    try:
-        status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300)
-    finally:
-        stop_gpu_oom_hogs(handle)
-
-    assert status >= 500, f"expected speech runtime error under OOM, got status={status}"
-    speech_error = extract_openai_error_contract_from_bytes(body)
-    assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}"
-    assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_request_failure(
-    omni_server_after_fault_function,
-    openai_client_function,
-) -> None:
-    try:
-        _submit_speech_request(openai_client_function, omni_server_after_fault_function)
-    except Exception as exc:
-        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
-    else:
-        pytest.fail("expected /v1/audio/speech request failure after process-kill injection")
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_fast_fail_and_concurrent(
-    omni_server_after_fault_function,
-    openai_client_function,
-) -> None:
-    host = omni_server_after_fault_function.host
-    port = omni_server_after_fault_function.port
-    model = omni_server_after_fault_function.model
-    payload = {
-        "model": model,
-        "input": "Say hello in one short sentence.",
-        "voice": "default",
-    }
-
-    start = time.monotonic()
-    try:
-        status, body = post_json_raw(
-            host,
-            port,
-            "/v1/audio/speech",
-            payload,
-            timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC,
-        )
-        elapsed = time.monotonic() - start
-        assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
-            f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s"
-        )
-        assert status >= 500, (
-            "[process_kill fast_fail] expected server-side failure after fault, "
-            f"got status={status}, body={body[:200]!r}"
-        )
-    except Exception:
-        elapsed = time.monotonic() - start
-        assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
-            f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s"
-        )
-
-    start = time.monotonic()
-    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
-        futures = [
-            executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function)
-            for _ in range(3)
-        ]
-        done, pending = concurrent.futures.wait(
-            futures,
-            timeout=30,
-            return_when=concurrent.futures.ALL_COMPLETED,
-        )
-
-    elapsed = time.monotonic() - start
-    assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}"
-    assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s"
-
-    fault_observed = False
-    for future in done:
-        try:
-            future.result()
-        except Exception as exc:
-            fault_observed = True
-            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
-    assert fault_observed, (
-        "[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault"
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_worker_with_load_request_failure(
-    omni_server_function,
-    openai_client_function,
-    fault_injector: FaultInjector,
-) -> None:
-    scenario = "kill_worker_with_load"
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
-        inject_fault=lambda: fault_injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
-    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
-    omni_server_function,
-    openai_client_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
-    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
-        inject_fault=lambda: injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=300.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
-    omni_server_function,
-    openai_client_function,
-    signal_name: str,
-) -> None:
-    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
-        inject_fault=lambda: injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )

tests/dfx/stability/scripts/test_stability_voxcpm2.py

@@ -1,69 +0,0 @@
-"""
-VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration.
-
-Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from tests.dfx.conftest import (
-    create_benchmark_indices,
-    create_test_parameter_mapping,
-    create_unique_server_params,
-    load_configs,
-)
-from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
-
-STABILITY_DIR = Path(__file__).resolve().parent.parent
-DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
-CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json")
-DEFAULT_NUM_PROMPTS_PER_BATCH = 20
-STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"]
-
-try:
-    BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-except FileNotFoundError:
-    BENCHMARK_CONFIGS = []
-
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
-server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
-benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
-def test_stability_voxcpm2(omni_server, stability_benchmark_params):
-    test_name = stability_benchmark_params["test_name"]
-    params = stability_benchmark_params["params"]
-    duration_sec = params.get("duration_sec", 300)
-    num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
-    request_rate = params.get("request_rate")
-    max_concurrency = params.get("max_concurrency")
-
-    bench_params = {
-        k: v
-        for k, v in params.items()
-        if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
-    }
-
-    result = run_stability_benchmark_loop(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=omni_server.model,
-        duration_sec=duration_sec,
-        params=bench_params,
-        request_rate=request_rate,
-        max_concurrency=max_concurrency,
-        result_dir=str(STABILITY_DIR),
-        num_prompts_per_batch=num_prompts_per_batch,
-        run_one_batch=_run_one_vllm_bench_batch,
-    )
-
-    assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
-    assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"

tests/diffusion/diffusion_kv/test_initialization.py

@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+import vllm_omni.diffusion.diffusion_kv.initialization as initialization
+import vllm_omni.diffusion.vllm_config as diffusion_vllm_config
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Executor:
+    def __init__(self) -> None:
+        self.specs = [{"layer0": object()}, {"layer0": object()}]
+        self.memory = [1024, 1024]
+        self.configs = None
+        self.profile_requests = None
+
+    def get_kv_cache_specs(self):
+        return self.specs
+
+    def determine_available_kv_memory(self, profile_requests):
+        self.profile_requests = profile_requests
+        return self.memory
+
+    def set_kv_cache_configs(self, configs) -> None:
+        self.configs = configs
+
+
+def _od_config(**overrides):
+    values = dict(
+        model="test-diffusion",
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        dtype=torch.bfloat16,
+        max_model_len=64,
+        quantization_config=None,
+        tf_model_config=None,
+        enforce_eager=True,
+        is_moe=False,
+        additional_config={},
+        profiler_config=None,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            enable_expert_parallel=False,
+        ),
+        gpu_memory_utilization=0.9,
+        kv_cache_memory_bytes=None,
+        max_num_batched_tokens=64,
+        num_gpus=1,
+    )
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
+    executor = _Executor()
+    od_config = _od_config(num_gpus=2)
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=4096),
+        max_in_flight_tokens=128,
+    )
+    worker_configs = [object(), object()]
+    scheduler_kv_cache_config = object()
+    profile_requests = [object()]
+    monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config)
+    build = Mock(return_value=(worker_configs, scheduler_kv_cache_config))
+    monkeypatch.setattr(
+        initialization,
+        "build_native_kv_cache_configs",
+        build,
+    )
+    monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16))
+
+    result = initialization.initialize_diffusion_kv_control_plane(
+        executor,
+        od_config,
+        profile_requests=profile_requests,
+        device=torch.device("cpu"),
+    )
+
+    build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
+    assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
+    assert executor.configs == worker_configs
+    assert executor.profile_requests is profile_requests
+
+
+def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None:
+    executor = _Executor()
+    od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+    get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device"))
+    monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device)
+
+    result = initialization.initialize_diffusion_kv_control_plane(executor, od_config)
+
+    assert result is None
+    assert executor.configs is None
+    get_device.assert_not_called()
+
+
+def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None:
+    executor = _Executor()
+    executor.memory = [1024]
+    od_config = _od_config(num_gpus=2)
+    monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object())
+
+    with pytest.raises(ValueError, match="rank count mismatch"):
+        initialization.initialize_diffusion_kv_control_plane(
+            executor,
+            od_config,
+            profile_requests=[object()],
+            device=torch.device("cpu"),
+        )
+
+
+def test_paged_control_plane_requires_prepared_profile_request() -> None:
+    with pytest.raises(ValueError, match="prepared memory profile requests"):
+        initialization.initialize_diffusion_kv_control_plane(
+            _Executor(),
+            _od_config(),
+            device=torch.device("cpu"),
+        )
+
+
+def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None:
+    od_config = _od_config()
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+    num_blocks = 32
+
+    worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * num_blocks],
+    )
+
+    assert len(worker_configs) == 1
+    assert worker_configs[0].num_blocks == num_blocks
+    assert scheduler_kv_cache_config.num_blocks == num_blocks
+    assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups
+
+
+def test_native_control_plane_config_initializes_scheduler() -> None:
+    od_config = _od_config()
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+    _, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * 32],
+    )
+
+    scheduler = RequestScheduler()
+    scheduler.initialize(
+        od_config,
+        kv_cache_config=scheduler_kv_cache_config,
+        scheduler_block_size=spec.block_size,
+        hash_block_size=spec.block_size,
+        kv_vllm_config=vllm_config,
+    )
+
+    assert scheduler._diffusion_kv_manager is not None
+    scheduler.close()
+
+
+def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None:
+    od_config = _od_config(
+        diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=4,
+            data_parallel_size=1,
+            enable_expert_parallel=False,
+        ),
+        gpu_memory_utilization=0.5,
+        kv_cache_memory_bytes=4096,
+        max_num_batched_tokens=64,
+    )
+    vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config)
+    original_pp_size = vllm_config.parallel_config.pipeline_parallel_size
+    original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization
+    original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes
+    original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens
+
+    diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config)
+
+    assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size
+    assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization
+    assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes
+    assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget
+
+
+def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None:
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+        torch.device("cpu"),
+        _od_config(gpu_memory_utilization=0.42),
+    )
+
+    assert vllm_config.cache_config.gpu_memory_utilization == 0.42
+
+
+def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
+    quantization_config = build_quant_config(
+        {
+            "quant_method": "modelopt",
+            "quant_algo": "FP8",
+            "ignore": [],
+        }
+    )
+    model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+            quantization_config=quantization_config,
+        )
+    )
+
+    assert model_config.dtype is torch.bfloat16
+    assert model_config.quantization == "modelopt"
+    assert model_config.quantization_config is quantization_config
+    assert model_config.is_quantized
+    assert model_config.is_moe is False
+    assert model_config.is_nvfp4_quantized() is False
+    assert model_config.original_max_model_len == model_config.max_model_len
+
+    unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+            quantization_config=None,
+        )
+    )
+    assert not unquantized_model_config.is_quantized
+
+
+def test_minus_one_max_model_len_uses_model_limit() -> None:
+    od_config = SimpleNamespace(
+        max_model_len=-1,
+        tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+    )
+
+    assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096
+
+    model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+        _od_config(
+            max_model_len=-1,
+            tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        )
+    )
+    assert model_config.max_model_len == 4096
+    assert model_config.original_max_model_len == -1
+
+
+def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None:
+    vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+        torch.device("cpu"),
+        _od_config(
+            max_model_len=-1,
+            tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+        ),
+    )
+    spec = FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        non_causal=True,
+    )
+
+    initialization.build_native_kv_cache_configs(
+        vllm_config,
+        [{"layer0": spec}],
+        [spec.page_size_bytes * 16],
+    )
+
+    assert vllm_config.model_config.original_max_model_len == -1
+    assert vllm_config.model_config.max_model_len == 256

tests/diffusion/diffusion_kv/test_manager.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+BLOCK_SIZE = 4
+
+
+def _config(num_blocks: int) -> KVCacheConfig:
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+
+
+def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
+    return DiffusionKVRequest(
+        f"{public_id}/diffusion-kv/{sequence_id}",
+        sequence_id=sequence_id,
+        prefix_len=4,
+        target_len=4,
+        seq_len=seq_len,
+        kv_contexts=kv_contexts,
+    )
+
+
+def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
+    return DiffusionKVCacheManager(
+        _config(num_blocks),
+        max_model_len=max_model_len,
+        scheduler_block_size=BLOCK_SIZE,
+        hash_block_size=BLOCK_SIZE,
+    )
+
+
+def test_reserve_and_free_multi_cfg_request() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    requests = (_request("public", 0), _request("public", 1))
+
+    metadata = manager.reserve_request("public", requests)
+
+    assert metadata is not None
+    assert metadata.request_id == "public"
+    assert metadata.allocation_generation == 1
+    assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1]
+    assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4]
+    assert [sequence.target_len for sequence in metadata.sequences] == [4, 4]
+    assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8]
+    assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+    assert all(sequence.context_ids == () for sequence in metadata.sequences)
+    assert metadata.contexts == ()
+    assert metadata == manager.get_metadata("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4
+
+    manager.free_request("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None:
+    manager = _manager(3)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+    with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0), _request("public", 1)),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None:
+    manager = _manager(5)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+    with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"):
+        manager.reserve_request(
+            "impossible",
+            tuple(_request("impossible", sequence_id) for sequence_id in range(3)),
+        )
+
+    assert manager.has_request("impossible") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_temporary_capacity_pressure_returns_none() -> None:
+    manager = _manager(3)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+
+    allocation = manager.reserve_request("waiting", (_request("waiting", 0),))
+
+    assert allocation is None
+    assert manager.has_request("waiting") is False
+
+
+def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None:
+    manager = _manager(5)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+    waiting = (_request("waiting", 0), _request("waiting", 1))
+
+    assert manager.reserve_request("waiting", waiting) is None
+    assert manager.has_request("waiting") is False
+
+    manager.free_request("running")
+    metadata = manager.reserve_request("waiting", waiting)
+
+    assert metadata is not None
+    assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+
+
+def test_reallocated_request_gets_a_new_allocation_generation() -> None:
+    manager = _manager(3)
+    request = (_request("public", 0),)
+
+    first = manager.reserve_request("public", request)
+    assert first is not None
+    manager.free_request("public")
+    second = manager.reserve_request("public", request)
+
+    assert second is not None
+    assert second.allocation_generation == first.allocation_generation + 1
+
+
+def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    native_allocate = manager.native_manager.allocate_slots
+    num_calls = 0
+
+    def fail_second_allocation(*args, **kwargs):
+        nonlocal num_calls
+        num_calls += 1
+        if num_calls == 2:
+            raise RuntimeError("injected allocation failure")
+        return native_allocate(*args, **kwargs)
+
+    monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation)
+
+    with pytest.raises(RuntimeError, match="injected allocation failure"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0), _request("public", 1)),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_rejects_independent_context_until_role_routing_is_implemented() -> None:
+    manager = _manager(8)
+    context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0, kv_contexts=(context,)),),
+        )
+
+
+def test_rejects_sequence_longer_than_admission_bound() -> None:
+    manager = _manager(8, max_model_len=8)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"):
+        manager.reserve_request("public", (_request("public", 0, seq_len=9),))
+
+
+@pytest.mark.parametrize(
+    "requests",
+    [
+        (_request("public", 1),),
+        (_request("public", 1), _request("public", 0)),
+        (_request("public", 0), _request("public", 2)),
+    ],
+)
+def test_rejects_ambiguous_sequence_order(requests) -> None:
+    manager = _manager(8)
+
+    with pytest.raises(ValueError, match="ordered by contiguous sequence_id"):
+        manager.reserve_request("public", requests)
+
+
+def test_close_releases_every_public_request() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    assert manager.reserve_request("first", (_request("first", 0),)) is not None
+    assert manager.reserve_request("second", (_request("second", 0),)) is not None
+
+    manager.close()
+
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+    assert manager.has_request("first") is False
+    assert manager.has_request("second") is False

tests/diffusion/diffusion_kv/test_request.py

@@ -107,11 +107,8 @@ def test_request_rejects_non_context_values() -> None:
 def _manager(*, enable_caching: bool) -> KVCacheManager:
     # vLLM 0.26 lazily registers its built-in KV specs. Another test may have
     # already registered an out-of-tree spec, making the registry non-empty
-    # before FullAttentionSpec is registered. Initialize the native mappings
-    # explicitly; vLLM 0.25 has no registration hook and needs no setup.
-    register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
-    if register_all is not None:
-        register_all(None)
+    # before FullAttentionSpec is registered. Initialize the native mappings.
+    native_kv_managers.register_all_kvcache_specs(None)
 
     spec = FullAttentionSpec(
         block_size=BLOCK_SIZE,

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Backend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return True
+
+
+def _attention(*, enabled: bool) -> Attention:
+    attention = Attention.__new__(Attention)
+    nn.Module.__init__(attention)
+    attention.paged_kv_cache_role = "primary" if enabled else None
+    attention.paged_kv_cache_dtype = torch.bfloat16
+    attention.num_kv_heads = 2
+    attention.head_size = 8
+    attention.causal = False
+    attention.attn_backend = _Backend
+    return attention
+
+
+def _runner(attention: Attention) -> DiffusionModelRunner:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+    runner.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=4),
+        model_config=SimpleNamespace(dtype=torch.float16),
+    )
+    pipeline = nn.Module()
+    pipeline.image_attention = attention
+    runner.pipeline = pipeline
+    return runner
+
+
+def test_runner_discovers_native_spec_from_loaded_attention() -> None:
+    specs = _runner(_attention(enabled=True)).get_kv_cache_spec()
+
+    assert set(specs) == {"image_attention"}
+    spec = specs["image_attention"]
+    assert isinstance(spec, FullAttentionSpec)
+    assert spec.block_size == 4
+    assert spec.num_kv_heads == 2
+    assert spec.head_size == 8
+    assert spec.dtype is torch.bfloat16
+    assert spec.indexes_kv_by_block_stride is True
+    assert spec.non_causal is True
+
+
+def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
+    runner = _runner(_attention(enabled=False))
+
+    with pytest.raises(RuntimeError, match="no cache-enabled Attention"):
+        runner.get_kv_cache_spec()
+
+
+def test_runner_retains_matching_rank_local_config() -> None:
+    runner = _runner(_attention(enabled=True))
+    spec = runner.get_kv_cache_spec()["image_attention"]
+    config = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
+    )
+
+    runner.set_kv_cache_config(config)
+
+    assert runner.kv_cache_config is config
+
+
+def test_runner_rejects_rank_local_config_for_different_layers() -> None:
+    runner = _runner(_attention(enabled=True))
+    spec = runner.get_kv_cache_spec()["image_attention"]
+    config = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
+    )
+
+    with pytest.raises(ValueError, match="layer mismatch"):
+        runner.set_kv_cache_config(config)
+
+
+def test_worker_selects_its_rank_local_config() -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 1
+    worker.od_config = SimpleNamespace(num_gpus=2)
+    worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
+    configs = [object(), object()]
+
+    worker.set_kv_cache_configs(configs)
+
+    assert worker.installed is configs[1]
+    with pytest.raises(ValueError, match="rank count mismatch"):
+        worker.set_kv_cache_configs(configs[:1])
+
+
+def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096))
+    worker.init_snapshot = SimpleNamespace(free_memory=16384)
+    worker.requested_memory = 8192
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [4096]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=900)
+    worker.requested_memory = 750
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=0,
+            gpu_memory_utilization=0.75,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=400,
+        after_profile=SimpleNamespace(free_memory=500),
+    )
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        lambda snapshot, weights_memory: nullcontext(profile_result),
+    )
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [350]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+@pytest.mark.parametrize("non_kv_cache_memory", [750, 800])
+def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=900)
+    worker.requested_memory = 750
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=None,
+            gpu_memory_utilization=0.75,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=non_kv_cache_memory,
+        after_profile=SimpleNamespace(free_memory=100),
+    )
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        lambda snapshot, weights_memory: nullcontext(profile_result),
+    )
+    profile_request = object()
+
+    with pytest.raises(
+        RuntimeError,
+        match=(
+            "rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: "
+            f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes"
+        ),
+    ):
+        worker.determine_available_kv_memory(profile_request)
+
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = 0
+    worker.init_snapshot = SimpleNamespace(free_memory=1200)
+    worker.requested_memory = 1000
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            kv_cache_memory_bytes=None,
+            gpu_memory_utilization=0.8,
+        )
+    )
+    worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250)
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+    profile_result = SimpleNamespace(
+        non_kv_cache_memory=650,
+        after_profile=SimpleNamespace(free_memory=550),
+    )
+    profiling_calls = []
+
+    def profile(snapshot, *, weights_memory):
+        profiling_calls.append((snapshot, weights_memory))
+        return nullcontext(profile_result)
+
+    monkeypatch.setattr(
+        diffusion_worker_module,
+        "memory_profiling",
+        profile,
+    )
+    profile_request = object()
+
+    assert worker.determine_available_kv_memory(profile_request) == [350]
+    assert profiling_calls == [(worker.init_snapshot, 250)]
+    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None:
+    worker = object.__new__(DiffusionWorker)
+    worker.init_snapshot = None
+    worker.requested_memory = None
+    worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None))
+    worker.model_runner = SimpleNamespace()
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+
+    with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"):
+        worker.determine_available_kv_memory(object())
+
+
+def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
+    gathered = []
+
+    def gather(local_result):
+        gathered.append(local_result)
+        return [local_result, (True, "peer-result")]
+
+    def fail_local_probe():
+        raise ValueError("local failure")
+
+    monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather)
+
+    with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"):
+        diffusion_worker_module._run_and_gather_rank_values(
+            "test probe",
+            fail_local_probe,
+        )
+
+    assert gathered == [(False, "ValueError: local failure")]

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -85,19 +85,15 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
     assert new_req.diffusion_kv_metadata is metadata
 
 
-def test_paged_request_without_metadata_reaches_request_forward() -> None:
+def test_paged_request_without_metadata_fails_before_request_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
-    expected_output = DiffusionOutput(output=None)
-    runner._execute_request_list = Mock(
-        return_value=SimpleNamespace(
-            runner_outputs=[SimpleNamespace(result=expected_output)],
-        )
-    )
+    runner._execute_request_list = Mock()
     req = SimpleNamespace(request_id="req-0")
 
-    assert runner.execute_model(req) is expected_output
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+        runner.execute_model(req)
 
-    runner._execute_request_list.assert_called_once()
+    runner._execute_request_list.assert_not_called()
 
 
 def test_dense_request_rejects_metadata_before_request_forward() -> None:
@@ -141,18 +137,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None:
     assert calls == []
 
 
-def test_batch_path_allows_missing_metadata_before_forward() -> None:
+def test_batch_path_rejects_missing_metadata_before_forward() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
-    expected_output = object()
-    runner._execute_request_list = Mock(return_value=expected_output)
+    runner._execute_request_list = Mock()
     new_req = NewRequestData(
         request_id="req-0",
         req=SimpleNamespace(request_id="req-0"),
     )
 
-    assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+        runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config)
 
-    runner._execute_request_list.assert_called_once()
+    runner._execute_request_list.assert_not_called()
 
 
 def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -170,7 +166,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
     runner._execute_request_list.assert_not_called()
 
 
-def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
+def test_step_path_rejects_missing_metadata_before_step_support_check() -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner.pipeline = object()
     runner._supports_step_mode = Mock(return_value=False)
@@ -179,10 +175,10 @@ def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
         req=SimpleNamespace(request_id="req-0"),
     )
 
-    with pytest.raises(ValueError, match="does not support step execution"):
+    with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
         runner.execute_stepwise(make_scheduler_output(new_req))
 
-    runner._supports_step_mode.assert_called_once()
+    runner._supports_step_mode.assert_not_called()
 
 
 def test_step_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -288,3 +284,42 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None:
     worker.execute_model(req, SimpleNamespace())
 
     assert calls == [(req, None)]
+
+
+def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None:
+    calls: list[tuple] = []
+
+    class ModelRunner:
+        def execute_model(self, req, **kwargs):
+            calls.append((req, kwargs))
+            return DiffusionOutput(output=None)
+
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = ModelRunner()
+    worker.lora_manager = None
+    worker._get_profiler = lambda: None
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank",
+        lambda: 1,
+    )
+    req_0 = SimpleNamespace(request_id="req-0")
+    req_1 = SimpleNamespace(request_id="req-1")
+    metadata_0 = make_metadata("req-0")
+    metadata_1 = make_metadata("req-1")
+    envelopes = [
+        NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0),
+        NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1),
+    ]
+
+    result = worker.execute_model(envelopes, SimpleNamespace())
+
+    assert calls == [
+        (
+            req_1,
+            {
+                "kv_prefetch_job": None,
+                "diffusion_kv_metadata": metadata_1,
+            },
+        )
+    ]
+    assert result["dp_rank"] == 1

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -12,9 +12,10 @@ 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 OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
@@ -91,6 +92,17 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade
     return loader
 
 
+def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model="",
+        dtype=torch.float32,
+        quantization_config="fp8",
+        parallel_config=DiffusionParallelConfig(data_parallel_size=2),
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=True,
+    )
+
+
 def test_strict_check_only_validates_source_prefix_parameters():
     model = _DummyPipelineModel(source_prefix="transformer.")
     loader = _make_loader_with_weights(["transformer.weight"])
@@ -298,6 +310,7 @@ 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
 
@@ -335,7 +348,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
         backing_kind="checkpoint_mmap",
         bindings={
             "transformer.weight": TensorBinding(
-                checkpoint_key="weight",
+                storage_key="weight",
                 file_path="unused",
             )
         },
@@ -363,8 +376,9 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
     assert loader.take_host_weight_plan() is plan
 
 
-def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
+def test_dlo_plan_fallback_retains_ordinary_loader_weights_by_default(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,
@@ -389,11 +403,274 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
         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,
+    )
+
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
+    model.transformer.quant_method.uses_meta_device = True
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # 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,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == [model]
+    assert calls == [("load", True), "process"]
+    assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    class UnsupportedOnlineMethod:
+        uses_meta_device = True
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnsupportedOnlineMethod()
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # 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,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    with pytest.raises(ValueError, match="per-tensor FP8 linears"):
+        loader.load_model(load_device="cpu")
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     from vllm_omni.diffusion.offloader.module_collector import PipelineModules

tests/diffusion/model_loader/test_host_weight_cache.py

@@ -0,0 +1,260 @@
+# 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/hunyuan_image3/test_image_kv_cache_manager.py

@@ -118,6 +118,14 @@ def _call_mgr(
     )
 
 
+def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
+    mgr = _make_cache_mgr()
+
+    assert isinstance(mgr, nn.Module)
+    assert dict(mgr.named_modules())["attn"] is mgr.attn
+    assert mgr.state_dict() == {}
+
+
 # ============================================================
 # Test 1: No AR KV — basic cache → reuse
 # ============================================================

tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression test: SenseNovaU1ForCausalLM instantiation under the diffusion
+config shim must not crash on missing ``head_dtype`` and the resulting
+LogitsProcessor must have ``head_dtype is None`` (= use model dtype)."""
+
+import os
+
+import pytest
+import torch
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+from vllm.distributed.parallel_state import (
+    cleanup_dist_env_and_memory,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+    SenseNovaU1ForCausalLM,
+)
+from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
+from vllm_omni.transformers_utils.configs.sensenova_u1 import (
+    SenseNovaU1Config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _build_tiny_sensenova_u1(quant_config=None, prefix="model"):
+    """Build a minimal SenseNovaU1ForCausalLM for unit tests."""
+    config = SenseNovaU1Config(
+        llm_config={
+            "hidden_size": 64,
+            "num_attention_heads": 2,
+            "num_hidden_layers": 1,
+            "intermediate_size": 128,
+            "vocab_size": 32,
+            "max_position_embeddings": 128,
+            "max_position_embeddings_hw": 128,
+        },
+    )
+    return SenseNovaU1ForCausalLM(config.llm_config, quant_config=quant_config, prefix=prefix)
+
+
+def test_logits_processor_head_dtype_under_diffusion_shim():
+    # Initialize the distributed environment
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29543")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+    initialize_model_parallel()
+
+    # Create a vLLM config with a wrapped diffusion config
+    fake_diff_config = _DiffusionVllmModelConfig(
+        model="sensenova-u1-test",
+        dtype=torch.bfloat16,
+        max_model_len=8192,
+        original_max_model_len=8192,
+    )
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig(device="cpu"),
+    )
+    vllm_config.modelconfig = fake_diff_config  # type: ignore[assignment]
+
+    # Build a tiny version of the Causal LM component
+    with set_current_vllm_config(vllm_config):
+        model = _build_tiny_sensenova_u1()
+
+    # Ensure that we have a set head_dtype attribute. Currently, the head_dtype
+    # is set to None
+    head_dtype = model.logits_processor.head_dtype
+    assert head_dtype is None
+    cleanup_dist_env_and_memory()

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -0,0 +1,346 @@
+# 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,6 +18,7 @@ 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,
@@ -185,6 +186,79 @@ class TestDistributedLayerwiseOffloadHook:
         assert next_block.weight.stride() == expected_stride
         assert torch.equal(next_block.weight, expected)
 
+    def test_allgather_reconstructs_online_fp8_weight_and_scale(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        """DLO must reconstruct finalized online-FP8 tensors byte-for-byte."""
+
+        class OnlineFp8Block(nn.Module):
+            def __init__(self):
+                super().__init__()
+                logical_weight = torch.tensor(
+                    [
+                        [1.0, 2.0, 3.0, 4.0],
+                        [5.0, 6.0, 7.0, 8.0],
+                        [9.0, 10.0, 11.0, 12.0],
+                    ],
+                    dtype=torch.float32,
+                ).to(torch.float8_e4m3fn)
+                self.weight = nn.Parameter(logical_weight.t(), requires_grad=False)
+                self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False)
+
+        current_block = OnlineFp8Block()
+        rank0_block = OnlineFp8Block()
+        rank1_block = OnlineFp8Block()
+        expected_weight = rank0_block.weight.detach().clone()
+        expected_scale = rank0_block.weight_scale.detach().clone()
+        expected_stride = rank0_block.weight.stride()
+
+        rank0_hook = DistributedLayerwiseOffloadHook(
+            next_block=rank0_block,
+            device=torch.device("cpu"),
+            dp_group=object(),
+            dp_size=2,
+            rank=0,
+            copy_stream=DummyStream(),
+            comm_stream=DummyStream(),
+            pin_memory=False,
+        )
+        rank1_hook = DistributedLayerwiseOffloadHook(
+            next_block=rank1_block,
+            device=torch.device("cpu"),
+            dp_group=object(),
+            dp_size=2,
+            rank=1,
+            copy_stream=DummyStream(),
+            comm_stream=DummyStream(),
+            pin_memory=False,
+        )
+        rank0_hook.initialize_hook(current_block)
+        rank1_hook.initialize_hook(OnlineFp8Block())
+        rank0_hook.gpu_shard_buffers = [
+            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+        ]
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+            shard_size = local_shard.numel()
+            output[:shard_size].copy_(local_shard)
+            output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+        rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+        assert rank0_block.weight.stride() == expected_stride
+        torch.testing.assert_close(
+            rank0_block.weight.float(),
+            expected_weight.float(),
+            rtol=0,
+            atol=0,
+        )
+        torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
+
     def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
         """DTensor wrapper should be preserved through prefetch/offload cycle."""
         current_block = TinyBlock(_make_values(1.0))
@@ -473,6 +547,31 @@ 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):
@@ -563,6 +662,24 @@ 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
     ):
@@ -692,6 +809,140 @@ 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()}
@@ -715,7 +966,7 @@ class TestMmapWeightLoading:
             backing_kind="checkpoint_mmap",
             bindings={
                 name: TensorBinding(
-                    checkpoint_key=name,
+                    storage_key=name,
                     file_path=str(weight_file),
                 )
                 for name in weights
@@ -786,6 +1037,292 @@ 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):
@@ -1091,7 +1628,7 @@ class TestMmapValidation:
         assert result.fallback_reason is None
         assert result.plan is not None
         assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
-            checkpoint_key="blocks.0.weight",
+            storage_key="blocks.0.weight",
             file_path=str(checkpoint_file),
         )
         assert result.plan.planned_source_prefixes == frozenset({"transformer."})
@@ -1381,7 +1918,7 @@ class TestMmapValidation:
 
         assert result.plan is not None
         assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
-            checkpoint_key="proj_in.weight",
+            storage_key="proj_in.weight",
             file_path=str(checkpoint_file),
         )
 

tests/diffusion/test_diffusion_config_propagation.py

@@ -130,6 +130,21 @@ class TestCreateDefaultDiffusion:
 
         assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
 
+    def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
+        monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)
+        od = _roundtrip_diffusion_config(
+            model="x",
+            kv_cache_memory_bytes=4096,
+            gpu_memory_utilization=0.75,
+            max_num_batched_tokens=2048,
+            max_model_len=4096,
+        )
+
+        assert od.kv_cache_memory_bytes == 4096
+        assert od.gpu_memory_utilization == 0.75
+        assert od.max_num_batched_tokens == 2048
+        assert od.max_model_len == 4096
+
 
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -161,9 +161,83 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa
     assert exc_info.value is initialization_error
     custom_scheduler.initialize.assert_called_once_with(od_config)
     custom_scheduler.close.assert_called_once_with()
+
+
+def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fake_executor = SimpleNamespace(shutdown=Mock())
+    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_resolve_execution_mode",
+        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+    )
+    monkeypatch.setattr(
+        DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor)
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+        Mock(side_effect=RuntimeError("KV initialization failed")),
+    )
+
+    with pytest.raises(RuntimeError, match="KV initialization failed"):
+        DiffusionEngine(SimpleNamespace())
+
     fake_executor.shutdown.assert_called_once_with()
 
 
+def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None:
+    events = []
+    expected_profile_requests = [object()]
+    fake_executor = SimpleNamespace(shutdown=Mock())
+    kv_cache_config = object()
+    kv_vllm_config = object()
+    od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler")
+
+    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_resolve_execution_mode",
+        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+    )
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_init_executor",
+        lambda self, config: setattr(self, "executor", fake_executor),
+    )
+    monkeypatch.setattr(
+        DiffusionEngine,
+        "_prepare_diffusion_kv_profile_requests",
+        lambda self: events.append("prepare-profile") or expected_profile_requests,
+    )
+
+    def initialize(executor, config, *, profile_requests):
+        events.append("profile-workers")
+        assert executor is fake_executor
+        assert config is od_config
+        assert profile_requests is expected_profile_requests
+        return kv_cache_config, 16, 16, kv_vllm_config
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+        initialize,
+    )
+
+    def init_scheduler(self, config, scheduler, *args, **kwargs):
+        events.append("initialize-scheduler")
+        self.scheduler = SimpleNamespace(close=Mock())
+
+    monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler)
+    monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None)
+    monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None)
+    monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None)
+
+    DiffusionEngine(od_config)
+
+    assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"]
+
+
 @pytest.mark.asyncio
 async def test_step_compatibility_wrapper_returns_final_batch() -> None:
     engine = _make_engine()

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -1,9 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 
 from vllm_omni.diffusion import io_support
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -55,3 +63,151 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk
     )
 
     assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
+
+
+def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
+    assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3
+    assert io_support.get_dummy_run_num_image_inputs("unknown") == 1
+
+
+def test_dense_mode_does_not_build_kv_profile_request() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+    engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile"))
+
+    assert engine._prepare_diffusion_kv_profile_requests() is None
+    engine._make_dummy_request.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
+    [
+        (DiffusionExecutionMode.STEP_BATCH, False, 2, 2),
+        (DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1),
+        (DiffusionExecutionMode.STEP_BATCH, True, 3, 3),
+    ],
+)
+def test_paged_kv_profile_requests_match_per_rank_batch(
+    execution_mode: DiffusionExecutionMode,
+    uses_dlo_dp: bool,
+    max_num_seqs: int,
+    expected_profile_requests: int,
+) -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine.execution_mode = execution_mode
+    engine.od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        model_class_name="HunyuanImage3ForCausalMM",
+        max_num_seqs=max_num_seqs,
+        parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1),
+        enable_distributed_layerwise_offload=uses_dlo_dp,
+        dlo_use_allgather=True,
+    )
+    request = OmniDiffusionRequest(
+        prompt="profile",
+        request_id="profile-request",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )
+    engine._make_dummy_request = Mock(return_value=request)
+    prepared_layout = object()
+    scheduler_kv_state = (
+        DiffusionKVRequest(
+            "profile-request/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=3,
+            target_len=8,
+            seq_len=12,
+        ),
+        DiffusionKVRequest(
+            "profile-request/diffusion-kv/1",
+            sequence_id=1,
+            prefix_len=4,
+            target_len=8,
+            seq_len=13,
+        ),
+    )
+
+    def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        req.prepared_layout = prepared_layout
+        req.diffusion_kv_requests = scheduler_kv_state
+        return req
+
+    engine._prepare_request_for_admission = Mock(side_effect=preprocess)
+
+    result = engine._prepare_diffusion_kv_profile_requests()
+
+    assert result is not None
+    assert len(result) == expected_profile_requests
+    assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests
+    assert all(profile_request.is_dummy_run() for profile_request in result)
+    assert all(profile_request.prepared_layout is prepared_layout for profile_request in result)
+    assert all(profile_request.diffusion_kv_requests is None for profile_request in result)
+    assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests
+    assert engine._diffusion_kv_profile_limits == (2, 13, 8)
+    assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2
+    engine._make_dummy_request.assert_called_once_with(
+        height=1024,
+        width=1024,
+        guidance_scale=5.0,
+        num_image_inputs=3,
+    )
+    engine._prepare_request_for_admission.assert_called_once_with(request)
+
+
+@pytest.mark.parametrize(
+    ("num_sequences", "seq_len", "target_len"),
+    [
+        (3, 16, 8),
+        (1, 17, 8),
+        (1, 16, 9),
+    ],
+)
+def test_paged_kv_admission_rejects_shape_beyond_profile_envelope(
+    num_sequences: int,
+    seq_len: int,
+    target_len: int,
+) -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine._diffusion_kv_profile_limits = (2, 16, 8)
+    engine.pre_process_func = lambda request: request
+    request = OmniDiffusionRequest(
+        prompt="too large",
+        request_id="too-large",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        diffusion_kv_requests=tuple(
+            DiffusionKVRequest(
+                f"too-large/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=0,
+                target_len=target_len,
+                seq_len=seq_len,
+            )
+            for sequence_id in range(num_sequences)
+        ),
+    )
+
+    with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
+        engine._prepare_request_for_admission(request)
+
+
+def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None:
+    engine = object.__new__(DiffusionEngine)
+    engine._diffusion_kv_profile_limits = (2, 16, 8)
+    engine.pre_process_func = lambda request: request
+    request = OmniDiffusionRequest(
+        prompt="fits",
+        request_id="fits",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        diffusion_kv_requests=tuple(
+            DiffusionKVRequest(
+                f"fits/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=8,
+                target_len=8,
+                seq_len=16,
+            )
+            for sequence_id in range(2)
+        ),
+    )
+
+    assert engine._prepare_request_for_admission(request) is request

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -473,6 +473,38 @@ async def test_cancelled_future_does_not_kill_busy_loop():
         _stop_engine(engine)
 
 
+@pytest.mark.asyncio
+async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+
+    def fail_reservation(*args, **kwargs):
+        raise RuntimeError("native allocation bug")
+
+    engine.scheduler._diffusion_kv_manager = SimpleNamespace(
+        has_request=lambda request_id: False,
+        reserve_request=fail_reservation,
+        free_request=lambda request_id: None,
+    )
+    request = _make_request("kv-error")
+    request.diffusion_kv_requests = (object(),)
+    try:
+        response_stream = engine.async_add_req_and_stream_response(request)
+        output_queue = engine._out_streams[request.request_id]
+        for _ in range(300):
+            if not output_queue.empty():
+                break
+            await asyncio.sleep(0.01)
+        assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream"
+
+        output = await anext(response_stream)
+        assert output.error == "native allocation bug"
+        assert output.finished
+        assert engine.worker_thread.is_alive()
+    finally:
+        _stop_engine(engine)
+
+
 # ─────────────────────────── _RpcTask basics ───────────────────────────────
 
 

tests/diffusion/test_diffusion_model_runner.py

@@ -3,6 +3,7 @@
 
 from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -512,6 +513,84 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
     assert runner.pipeline.last_req.num_reqs == 1
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    request = _make_request()
+    runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+    record_names = []
+    original_execute = runner._execute_request_list
+
+    def execute(*args, **kwargs):
+        record_names.append(kwargs["record_name"])
+        return original_execute(*args, **kwargs)
+
+    runner._execute_request_list = execute
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    reset_peak_memory_stats = Mock()
+    monkeypatch.setattr(
+        model_runner_module.current_omni_platform,
+        "reset_peak_memory_stats",
+        reset_peak_memory_stats,
+    )
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    synchronize = Mock()
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+    runner.profile_run([request])
+
+    assert runner.pipeline.forward_calls == 1
+    assert record_names == ["pipeline_memory_profile"]
+    runner._validate_diffusion_kv_metadata.assert_not_called()
+    reset_peak_memory_stats.assert_not_called()
+    synchronize.assert_called_once_with()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.pipeline = _FinalOnlyStepPipeline()
+    runner.od_config.step_execution = True
+    requests = [_make_request(), _make_request()]
+    requests[0].request_id = "profile-0"
+    requests[1].request_id = "profile-1"
+    for request in requests:
+        request.sampling_params.num_inference_steps = 1
+
+    observed_batch_rows = []
+    original_denoise_step = runner.pipeline.denoise_step
+
+    def denoise_step(input_batch, states):
+        observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
+        return original_denoise_step(input_batch, states)
+
+    runner.pipeline.denoise_step = denoise_step
+    runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    reset_peak_memory_stats = Mock()
+    monkeypatch.setattr(
+        model_runner_module.current_omni_platform,
+        "reset_peak_memory_stats",
+        reset_peak_memory_stats,
+    )
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    synchronize = Mock()
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+    runner.profile_run(requests)
+
+    assert observed_batch_rows == [(2, 2)]
+    assert runner.state_cache == {}
+    assert runner.input_batch is None
+    runner._validate_diffusion_kv_metadata.assert_not_called()
+    reset_peak_memory_stats.assert_not_called()
+    synchronize.assert_called_once_with()
+
+
 class _BatchPipeline:
     """Pipeline returning a configurable list of outputs from forward()."""
 

tests/diffusion/test_diffusion_request.py

@@ -52,6 +52,7 @@ def test_dummy_run_request_is_identified_by_reserved_request_id():
 
     assert req.is_dummy_run()
     assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
+    assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1")
 
 
 def test_non_dummy_request_is_not_identified_as_dummy_run():

tests/diffusion/test_diffusion_scheduler.py

@@ -8,12 +8,14 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from pytest_mock import MockerFixture
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
-from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
     BaseScheduler,
@@ -87,6 +89,52 @@ def _cached_ids(sched_output) -> list[str]:
     return list(sched_output.scheduled_cached_reqs.request_ids)
 
 
+def _initialize_paged_scheduler(
+    scheduler: BaseScheduler,
+    *,
+    num_blocks: int = 64,
+    max_num_seqs: int = 1,
+) -> None:
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+    )
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+    scheduler.initialize(
+        SimpleNamespace(
+            diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+            max_model_len=64,
+            max_num_seqs=max_num_seqs,
+        ),
+        kv_cache_config=config,
+        scheduler_block_size=4,
+        hash_block_size=4,
+        kv_vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(max_model_len=64),
+            max_in_flight_tokens=64,
+        ),
+    )
+
+
+def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
+    request.diffusion_kv_requests = (
+        DiffusionKVRequest(
+            f"{request.request_id}/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=4,
+            target_len=4,
+            seq_len=seq_len,
+        ),
+    )
+
+
 class _StubScheduler:
     def __init__(self, request: OmniDiffusionRequest, output) -> None:
         self._request = request
@@ -400,18 +448,16 @@ class TestRequestScheduler:
         assert self.scheduler.has_requests() is False
 
     def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
         prepared_layout = object()
         request.prepared_layout = prepared_layout
-        text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
         kv_request = DiffusionKVRequest(
             "diffusion-kv/diffusion-kv/0",
             sequence_id=0,
             prefix_len=4,
             target_len=8,
             seq_len=16,
-            kv_contexts=(text_context,),
         )
         request.diffusion_kv_requests = (kv_request,)
 
@@ -421,12 +467,179 @@ class TestRequestScheduler:
 
         assert state is not None
         assert state.diffusion_kv_requests == (kv_request,)
-        assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
         assert request.diffusion_kv_requests is None
         assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
+        metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.request_id == request.request_id
+        assert len(metadata.sequences[0].block_ids[0]) == 4
+
+    def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
+        first_request = _make_request("first")
+        second_request = _make_request("second")
+        _attach_diffusion_kv(first_request)
+        _attach_diffusion_kv(second_request)
+        self.scheduler.add_request(first_request)
+        self.scheduler.add_request(second_request)
+
+        first_output = self.scheduler.schedule()
+
+        assert _new_ids(first_output) == ["first"]
+        assert first_output.num_waiting_reqs == 1
+
+        self.scheduler.update_from_output(first_output, _make_request_output("first"))
+        second_output = self.scheduler.schedule()
+
+        assert _new_ids(second_output) == ["second"]
+        assert second_output.num_waiting_reqs == 0
+        # The same output carries both the release and the newly admitted block
+        # table; the future Worker data plane must process finished ids first.
+        assert second_output.finished_req_ids == {"first"}
+        assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None
+
+    def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2)
+        impossible = _make_request("impossible")
+        schedulable = _make_request("schedulable")
+        _attach_diffusion_kv(impossible, seq_len=8)
+        schedulable.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "schedulable/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            ),
+        )
+        self.scheduler.add_request(impossible)
+        self.scheduler.add_request(schedulable)
+
+        sched_output = self.scheduler.schedule()
+
+        failed_state = self.scheduler.get_request_state("impossible")
+        assert failed_state is not None
+        assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert failed_state.error is not None
+        assert "cannot fit even when the block pool is empty" in failed_state.error
+        assert sched_output.finished_req_ids == {"impossible"}
+        assert _new_ids(sched_output) == ["schedulable"]
+
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_request_output("schedulable"),
+        )
+
+        assert finished == {"impossible", "schedulable"}
+
+    def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2)
+        running = _make_request("running")
+        _attach_diffusion_kv(running)
+        self.scheduler.add_request(running)
+        assert _new_ids(self.scheduler.schedule()) == ["running"]
+
+        impossible = _make_request("impossible")
+        impossible.diffusion_kv_requests = tuple(
+            DiffusionKVRequest(
+                f"impossible/diffusion-kv/{sequence_id}",
+                sequence_id=sequence_id,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            )
+            for sequence_id in range(4)
+        )
+        schedulable = _make_request("schedulable")
+        schedulable.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "schedulable/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=0,
+                target_len=4,
+                seq_len=4,
+            ),
+        )
+        self.scheduler.add_request(impossible)
+        self.scheduler.add_request(schedulable)
+
+        sched_output = self.scheduler.schedule()
+
+        failed_state = self.scheduler.get_request_state("impossible")
+        assert failed_state is not None
+        assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert failed_state.error is not None
+        assert "required_blocks=4, available_blocks=3" in failed_state.error
+        assert sched_output.finished_req_ids == {"impossible"}
+        assert _new_ids(sched_output) == ["schedulable"]
+
+    def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None:
+        _initialize_paged_scheduler(self.scheduler)
+        request = _make_request("native-error")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+
+        def raise_native_error(*args, **kwargs):
+            raise ValueError("native allocation bug")
+
+        monkeypatch.setattr(
+            manager,
+            "reserve_request",
+            raise_native_error,
+        )
+
+        sched_output = self.scheduler.schedule()
+
+        state = self.scheduler.get_request_state("native-error")
+        assert state is not None
+        assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert state.error == "native allocation bug"
+        assert sched_output.finished_req_ids == {"native-error"}
+        assert sched_output.scheduled_request_ids == []
+
+    def test_diffusion_kv_preemption_retains_allocation(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+        request = _make_request("preempted")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        self.scheduler.schedule()
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        metadata = manager.get_metadata("preempted")
+
+        assert self.scheduler.preempt_request("preempted") is True
+        resumed_output = self.scheduler.schedule()
+
+        assert resumed_output.scheduled_new_reqs == []
+        assert _cached_ids(resumed_output) == ["preempted"]
+        assert manager.get_metadata("preempted") == metadata
+
+    @pytest.mark.parametrize(
+        "status",
+        [
+            DiffusionRequestStatus.FINISHED_COMPLETED,
+            DiffusionRequestStatus.FINISHED_ABORTED,
+            DiffusionRequestStatus.FINISHED_ERROR,
+        ],
+    )
+    def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None:
+        _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+        request = _make_request("terminal")
+        _attach_diffusion_kv(request)
+        self.scheduler.add_request(request)
+        self.scheduler.schedule()
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        assert manager.has_request("terminal") is True
+
+        self.scheduler.finish_requests("terminal", status)
+
+        assert manager.has_request("terminal") is False
 
     def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
 
         with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
@@ -443,7 +656,7 @@ class TestRequestScheduler:
         ],
     )
     def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
-        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        _initialize_paged_scheduler(self.scheduler)
         request = _make_request("diffusion-kv")
         owner = request if owner_name == "request" else request.sampling_params
         setattr(owner, field_name, object())
@@ -921,6 +1134,45 @@ class TestDiffusionEngine:
         assert engine.scheduler.max_num_running_reqs == 1
         fake_executor_cls.assert_called_once_with(od_config)
 
+    def test_initializes_paged_scheduler_from_native_control_plane(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
+        fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
+        scheduler = mocker.Mock()
+        kv_cache_config = object()
+        kv_vllm_config = object()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+            lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config),
+        )
+
+        engine = DiffusionEngine(od_config, scheduler=scheduler)
+
+        assert engine.scheduler is scheduler
+        scheduler.initialize.assert_called_once_with(
+            od_config,
+            kv_cache_config=kv_cache_config,
+            scheduler_block_size=16,
+            hash_block_size=16,
+            kv_vllm_config=kv_vllm_config,
+        )
+
     def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
         scheduler = Scheduler()
         scheduler.initialize(SimpleNamespace())
@@ -987,6 +1239,21 @@ class TestDiffusionEngine:
         assert output.aborted is True
         assert output.abort_message == "Request req-finalize aborted."
 
+    def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = StepScheduler()
+        engine.scheduler.initialize(SimpleNamespace())
+
+        req_id = engine.scheduler.add_request(_make_request("req-error"))
+        engine.scheduler._finish_requests(
+            {req_id: DiffusionRequestStatus.FINISHED_ERROR},
+            {req_id: "KV request cannot fit"},
+        )
+
+        output = engine._finalize_finished_request(req_id)
+
+        assert output.error == "KV request cannot fit"
+
     @pytest.mark.asyncio
     async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)

tests/diffusion/test_diffusion_worker.py

@@ -13,11 +13,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
-from vllm_omni.diffusion.worker.diffusion_worker import (
-    DiffusionWorker,
-    _create_diffusion_worker_vllm_config,
-    _make_diffusion_vllm_model_config,
-)
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
@@ -58,34 +54,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
     return worker
 
 
-def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods():
-    from types import SimpleNamespace
-
-    from vllm_omni.quantization import build_quant_config
-
-    od_config = SimpleNamespace(
-        model="dummy",
-        dtype=torch.bfloat16,
-        quantization_config=build_quant_config(
-            {
-                "quant_method": "modelopt",
-                "quant_algo": "FP8",
-                "ignore": [],
-            }
-        ),
-        tf_model_config=SimpleNamespace(),
-        enforce_eager=True,
-        is_moe=False,
-    )
-
-    model_config = _make_diffusion_vllm_model_config(od_config)
-
-    assert model_config.dtype is torch.bfloat16
-    assert model_config.quantization == "modelopt"
-    assert model_config.quantization_config is od_config.quantization_config
-    assert model_config.is_quantized()
-
-
 class TestDiffusionWorkerSleep:
     """Test DiffusionWorker.sleep method."""
 
@@ -329,49 +297,3 @@ class TestDiffusionWorkerWakeUp:
         mock_buffer2.data.copy_.assert_not_called()
 
         assert result is True
-
-
-class TestWorkerVllmConfigAdditionalConfig:
-    """Test worker-side VllmConfig construction with additional_config."""
-
-    def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture):
-        mock_vllm_config = mocker.Mock()
-        mock_vllm_cls = mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            return_value=mock_vllm_config,
-        )
-        od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}})
-
-        result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
-        assert result is mock_vllm_config
-        assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config
-
-    def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config(
-        self, mocker: MockerFixture
-    ):
-        mock_vllm_config = mocker.Mock()
-        mock_vllm_cls = mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            side_effect=[
-                TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"),
-                mock_vllm_config,
-            ],
-        )
-        od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
-        result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
-        assert result is mock_vllm_config
-        assert mock_vllm_cls.call_count == 2
-        assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config
-
-    def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture):
-        mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
-            side_effect=TypeError("additional_config values must be hashable"),
-        )
-        od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
-        with pytest.raises(TypeError, match="hashable"):
-            _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -413,6 +413,26 @@ class TestRequestModeDispatch:
         assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
         executor.collective_rpc.assert_called_once()
 
+    def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id)
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+        forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0]
+        assert forwarded_envelopes is scheduler_output.scheduled_new_reqs
+        assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"]
+
     @pytest.mark.parametrize(
         ("field", "value"),
         [

tests/diffusion/test_worker_wrapper_base.py

@@ -12,11 +12,13 @@ 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,
@@ -204,6 +206,29 @@ 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/online_serving/test_minimax_music3.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online tests for MiniMax Music 3 text-to-music.
+
+This is a music model on the speech endpoint, so the request shape differs
+from every other TTS model here: ``input`` carries the lyrics, ``instructions``
+carries the caption that decides genre and arrangement, and both are required.
+
+Weekly rather than per-PR: the checkpoint is 27 GB across five components and a
+single request occupies two decode rows, because classifier-free guidance is
+mandatory for this model. Rejection cases live in
+``tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py``.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "MiniMaxAI/MiniMax-Music3"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 600.0
+
+# 25 frames per second, so 250 frames is ten seconds and 750 is thirty. The
+# byte floors sit below the true payload (32 kHz stereo, 16-bit) to leave room
+# for the model ending a song before the cap.
+_FRAMES_10S = 250
+_FRAMES_30S = 750
+_MIN_BYTES_10S = 500_000
+_MIN_BYTES_30S = 1_500_000
+
+LYRICS = (
+    "[Verse]\nWalking down the empty street at midnight\n"
+    "Streetlights flicker like a broken dream\n"
+    "[Chorus]\nAnd I keep on walking\nTill the morning finds me"
+)
+
+CAPTION = (
+    "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, "
+    "soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. "
+    "Intimate bedroom production, gentle tape saturation, no bright cymbals."
+)
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="minimax_music3",
+    )
+]
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_music_001(omni_server, openai_client) -> None:
+    """Baseline: lyrics plus caption in, one ten-second song out.
+
+    Deploy Setting: minimax_music3.yaml
+    Input Modal: text (lyrics) + text (caption)
+    Output Modal: audio, 32 kHz stereo
+    Input Setting: stream=False, seed pinned, 250 frames
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": LYRICS,
+        "instructions": CAPTION,
+        "seed": 1,
+        "max_new_tokens": _FRAMES_10S,
+        "stream": False,
+        "response_format": "wav",
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "min_audio_bytes": _MIN_BYTES_10S,
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_music_multi_window_002(omni_server, openai_client) -> None:
+    """A song long enough to span several acoustic windows.
+
+    The acoustic stage decodes overlapping 200-frame windows on a 100-frame
+    hop and crops the overlap back off, so a request that fits in one window
+    never exercises the stitching. 750 frames covers seven of them.
+
+    Deploy Setting: minimax_music3.yaml
+    Input Setting: stream=False, seed pinned, 750 frames
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": LYRICS,
+        "instructions": CAPTION,
+        "seed": 1,
+        "max_new_tokens": _FRAMES_30S,
+        "stream": False,
+        "response_format": "wav",
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "min_audio_bytes": _MIN_BYTES_30S,
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/engine/test_async_omni_engine_input.py

@@ -371,3 +371,106 @@ def test_build_add_request_message_releases_preselected_replica_on_preprocess_er
         )
 
     assert stage_pool.get_bound_replica_id("req-error") is None
+
+
+def test_cfg_companion_build_failure_admits_nothing(mocker: MockerFixture):
+    """A guided request is all-or-nothing.
+
+    Admitting the parent before its companion is built leaves an orphan: a
+    model whose guidance is mandatory never completes the pair, so the request
+    holds scheduler and KV capacity for the scheduler's whole hold budget and
+    then produces no audio. The build has to raise before anything is enqueued.
+    """
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = mocker.Mock(side_effect=ValueError("cannot build the null twin"))
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.request_queue = mocker.Mock()
+
+    with pytest.raises(ValueError, match="null twin"):
+        engine._build_cfg_companions(
+            parent_id="req",
+            original_prompt={"prompt": "positive"},
+            stage0_params=params,
+            sampling_params_list=[params],
+        )
+
+    engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_companion_processing_failure_admits_nothing(mocker: MockerFixture):
+    """The same holds when expansion succeeds but input processing does not."""
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: [
+        ExpandedPrompt(
+            prompt={"prompt": "negative"},
+            role="cfg_text",
+            request_id_suffix="__cfg_text",
+        )
+    ]
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.input_processor.process_inputs.side_effect = RuntimeError("tokenizer rejected the prompt")
+    engine.request_queue = mocker.Mock()
+
+    with pytest.raises(RuntimeError, match="tokenizer rejected"):
+        engine._build_cfg_companions(
+            parent_id="req",
+            original_prompt={"prompt": "positive"},
+            stage0_params=params,
+            sampling_params_list=[params],
+        )
+
+    engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_companion_build_returns_messages_without_enqueueing(mocker: MockerFixture):
+    """Construction and admission are separate steps."""
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: [
+        ExpandedPrompt(
+            prompt={"prompt": "negative"},
+            role="cfg_text",
+            request_id_suffix="__cfg_text",
+        )
+    ]
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
+    engine.request_queue = mocker.Mock()
+
+    companions = engine._build_cfg_companions(
+        parent_id="req",
+        original_prompt={"prompt": "positive"},
+        stage0_params=params,
+        sampling_params_list=[params],
+    )
+
+    assert len(companions) == 1
+    assert companions[0].companion_id == "req__cfg_text"
+    assert companions[0].parent_id == "req"
+    engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_expansion_returning_nothing_is_not_an_error(mocker: MockerFixture):
+    """An unguided request expands to no companions and admits normally."""
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: []
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.request_queue = mocker.Mock()
+
+    assert (
+        engine._build_cfg_companions(
+            parent_id="req",
+            original_prompt={"prompt": "positive"},
+            stage0_params=params,
+            sampling_params_list=[params],
+        )
+        == []
+    )
+    engine.request_queue.sync_q.put.assert_not_called()

tests/engine/test_stage_engine_args.py

@@ -8,8 +8,23 @@ from pathlib import Path
 
 import pytest
 from pydantic.fields import FieldInfo
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
+from vllm.engine.arg_utils import EngineArgs
 
-from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.omni_config import (
+    _LLM_STAGE_ENGINE_FIELDS,
+    OmniStageCacheConfig,
+    OmniStageLoadConfig,
+    OmniStageParallelConfig,
+    OmniStageSchedulerConfig,
+    VllmOmniARStageConfig,
+    VllmOmniConfig,
+)
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
 from vllm_omni.config.stage_config import (
     DeployConfig,
@@ -53,6 +68,30 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str,
 
 _LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
 _DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
+_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
+    {
+        "active_stream_window",
+        "codec_frame_rate_hz",
+        "custom_voice_dir",
+        "devices",
+        "disable_autocast",
+        "enable_multithread_weight_load",
+        "env",
+        "has_sampling_extra_args",
+        "log_level",
+        "log_stats",
+        "model_arch",
+        "model_subdir",
+        "num_gpus",
+        "num_replicas",
+        "num_weight_load_threads",
+        "omni_kv_config",
+        "parallel_config",
+        "subtalker_sampling_params",
+        "task_type",
+        "tokenizer_subdir",
+    }
+)
 
 
 @pytest.fixture(autouse=True)
@@ -197,6 +236,100 @@ def _legacy_and_typed_stages(
     return legacy_stages, omni_config
 
 
+def test_llm_stage_engine_field_schema_tracks_upstream_engine_args():
+    upstream_engine_fields = frozenset(field.name for field in fields(EngineArgs))
+
+    assert _LLM_STAGE_ENGINE_FIELDS - upstream_engine_fields == _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS
+
+
+def test_typed_llm_projection_discovers_explicit_upstream_config_fields():
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+        load_config=OmniStageLoadConfig(ignore_patterns=["*.bin"]),
+        cache_config=OmniStageCacheConfig(block_size=32, cache_dtype="fp8"),
+        scheduler_config=OmniStageSchedulerConfig(
+            prefill_schedule_interval=2,
+            policy="priority",
+        ),
+        parallel_config=OmniStageParallelConfig(
+            data_parallel_master_ip="10.0.0.1",
+            enable_dbo=True,
+        ),
+    )
+
+    engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+    assert engine_args["ignore_patterns"] == ["*.bin"]
+    assert engine_args["block_size"] == 32
+    assert engine_args["kv_cache_dtype"] == "fp8"
+    assert engine_args["prefill_schedule_interval"] == 2
+    assert engine_args["scheduling_policy"] == "priority"
+    assert engine_args["data_parallel_address"] == "10.0.0.1"
+    assert engine_args["enable_dbo"] is True
+
+
+def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+    )
+
+    engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+    inherited_defaults = {
+        "ignore_patterns",
+        "block_size",
+        "kv_cache_dtype",
+        "prefill_schedule_interval",
+        "scheduling_policy",
+        "data_parallel_address",
+        "data_parallel_rank",
+        "data_parallel_size_local",
+        "data_parallel_rpc_port",
+        "enable_dbo",
+    }
+    assert inherited_defaults.isdisjoint(engine_args)
+
+
+def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary():
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+        scheduler_config=OmniStageSchedulerConfig(scheduler_cls="test.CustomScheduler"),
+    )
+
+    with pytest.raises(ValueError, match=r"no EngineArgs projection: scheduler_cls"):
+        stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+
+@pytest.mark.parametrize("stage_id", [0, 1], ids=["ar", "generation"])
+def test_typed_llm_engine_args_preserve_upstream_config_objects(tmp_path, stage_id):
+    pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+    compilation_config = VllmCompilationConfig(backend="eager")
+    profiler_config = VllmProfilerConfig(profiler="cuda")
+    stage_prefix = f"stage_{stage_id}_"
+    omni_config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=copy.deepcopy(deploy),
+        cli_overrides={
+            "model": model,
+            f"{stage_prefix}compilation_config": compilation_config,
+            f"{stage_prefix}profiler_config": profiler_config,
+        },
+    )
+
+    stage_config = omni_config.stage_by_id(stage_id)
+    typed_args = build_engine_args_dict_from_omni_stage_config(stage_config, model)
+
+    assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
+    assert isinstance(typed_args["compilation_config"], VllmCompilationConfig)
+    assert typed_args["compilation_config"] is not stage_config.compilation_config
+    assert typed_args["compilation_config"].backend == "eager"
+
+    assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
+    assert isinstance(typed_args["profiler_config"], VllmProfilerConfig)
+    assert typed_args["profiler_config"] is not stage_config.profiler_config
+    assert typed_args["profiler_config"].profiler == "cuda"
+
+
 def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
     pipeline, deploy, model = _engine_arg_inputs(tmp_path)
     legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
@@ -222,6 +355,13 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
         assert {name: typed_args[name] for name in legacy_args} == legacy_args
 
     thinker_args = typed_args_by_stage[0]
+    inherited_vllm_fields = {
+        field.name
+        for config_cls in (VllmLoadConfig, VllmCacheConfig, VllmSchedulerConfig, VllmParallelConfig)
+        for field in fields(config_cls)
+    }
+    topology_projected_fields = {"scheduler_cls"}
+    assert (inherited_vllm_fields - _LLM_STAGE_ENGINE_FIELDS - topology_projected_fields).isdisjoint(thinker_args)
     assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model")
     assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer")
     assert thinker_args["stage_id"] == 0

tests/entrypoints/openai_api/test_serving_speech.py

@@ -4059,9 +4059,7 @@ class TestTTSAsyncOffloading:
 
         assert adapter_model_type != legacy_tts_model_type
         assert any(
-            call.args
-            and call.args[0] == "TTS speech request %s: text=%r, model=%s"
-            and call.args[3] == adapter_model_type
+            call.args and call.args[0] == "TTS speech request %s: model=%s" and call.args[2] == adapter_model_type
             for call in log_info.call_args_list
         )
 

tests/entrypoints/openai_api/test_tts_detection.py

@@ -153,6 +153,7 @@ _PIPELINE_STAGES = [
     "indextts2_talker",
     "latent_generator",
     "llm",
+    "minimax_music3_ar",
     "ming_tts",
     "moss_tts",
     "moss_tts_codec",

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -328,6 +328,9 @@ 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",
         ]
@@ -339,9 +342,15 @@ 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/models/test_minimax_music3_batch_payload.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 batched payload partitioning.
+
+The stage-0 payload is split per request by indexing list values: a list is
+indexed by batch row, anything else is handed to every request unchanged. So
+per-request metadata has to be carried as lists. Merging it into one dict made
+every request in a batch receive the last request's seed, window offset and
+terminal flag, which crosses seeds between songs, misplaces acoustic windows,
+and lets one request's terminal flag end another request's stream early.
+
+These drive ``make_omni_output`` directly with a hand-built batch, so no GPU
+and no engine are needed.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+    MiniMaxMusic3TalkerForConditionalGeneration as Talker,
+)
+from vllm_omni.utils.mm_outputs import to_payload_element
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_WIDTH = 4
+
+
+def _info(*, seed: int, start: int, frames: int, last: bool) -> dict:
+    """One batch row's intermediate buffer, as ``postprocess`` leaves it."""
+    return {
+        "latent": torch.full((frames, _WIDTH), float(seed)),
+        "codes": {"audio": torch.full((frames, 8), seed)},
+        "meta": {
+            "audio_seed": seed,
+            "num_processed_tokens": start,
+            "codec_chunk_frames": frames,
+            "last_chunk": last,
+        },
+    }
+
+
+def _payload_for(output, row: int) -> dict:
+    """Partition the batch payload down to one request, as the runner does."""
+    flat = {}
+    for key, value in output.multimodal_outputs.items():
+        if isinstance(value, dict):
+            for qual, val in value.items():
+                flat[f"{key}.{qual}"] = val
+        else:
+            flat[key] = value
+    return {key: to_payload_element(value, row, 1, 1) for key, value in flat.items()}
+
+
+class TestBatchedMetadataPartitioning:
+    """Two requests in one batch must keep their own metadata."""
+
+    @staticmethod
+    def _two_request_batch():
+        infos = [
+            _info(seed=111, start=100, frames=3, last=False),
+            _info(seed=222, start=200, frames=5, last=True),
+        ]
+        talker = Talker.__new__(Talker)
+        return talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+    def test_seeds_are_not_crossed(self) -> None:
+        """A song decoded with another request's seed is a different song."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.audio_seed"] == 111
+        assert _payload_for(output, 1)["meta.audio_seed"] == 222
+
+    def test_window_offsets_stay_with_their_request(self) -> None:
+        """The start frame decides where a window lands in the song."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.num_processed_tokens"] == 100
+        assert _payload_for(output, 1)["meta.num_processed_tokens"] == 200
+
+    def test_terminal_flag_does_not_leak_across_requests(self) -> None:
+        """One request finishing must not end another request's stream."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.last_chunk"] is False
+        assert _payload_for(output, 1)["meta.last_chunk"] is True
+
+    def test_span_length_matches_its_own_latent(self) -> None:
+        output = self._two_request_batch()
+        for row, frames in ((0, 3), (1, 5)):
+            entry = _payload_for(output, row)
+            assert entry["meta.codec_chunk_frames"] == frames
+            assert entry["latent"].shape[0] == frames
+
+    def test_latent_and_codes_follow_the_same_row(self) -> None:
+        """The tensors were already per-request; pin that they stay aligned."""
+        output = self._two_request_batch()
+        for row, seed in ((0, 111), (1, 222)):
+            entry = _payload_for(output, row)
+            assert entry["latent"][0, 0].item() == pytest.approx(float(seed))
+            assert int(entry["codes.audio"][0, 0]) == seed
+
+
+class TestPartialBatches:
+    """Only some rows carry a span on any given step."""
+
+    def test_row_without_a_span_does_not_take_another_row_metadata(self) -> None:
+        infos = [{}, _info(seed=777, start=300, frames=2, last=True)]
+        talker = Talker.__new__(Talker)
+        output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+        idle = _payload_for(output, 0)
+        assert idle["latent"].numel() == 0
+        assert idle["meta.audio_seed"] is None
+        assert idle["meta.last_chunk"] is None
+
+        active = _payload_for(output, 1)
+        assert active["meta.audio_seed"] == 777
+        assert active["meta.last_chunk"] is True
+
+    def test_batch_with_no_spans_emits_no_payload(self) -> None:
+        talker = Talker.__new__(Talker)
+        output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=[{}, {}])
+        assert output.multimodal_outputs is None
+
+    def test_consumed_entries_are_cleared_from_the_buffer(self) -> None:
+        """The runner never clears this buffer, so a span left behind would be
+        re-emitted on every later step."""
+        infos = [_info(seed=1, start=0, frames=2, last=False)]
+        talker = Talker.__new__(Talker)
+        talker.make_omni_output(torch.zeros(1, _WIDTH), model_intermediate_buffer=infos)
+        assert "latent" not in infos[0]
+        assert "meta" not in infos[0]

tests/model_executor/models/test_minimax_music3_frame_state.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 per-request AR frame bookkeeping.
+
+The talker emits overlapping windows while dropping frames it can no longer
+need, and flushes whatever is left as one span when the song ends. Getting the
+eviction boundary wrong by one frame silently corrupts the conditioning handed
+to the acoustic stage from the second window onward, which is audible but hard
+to attribute, so the lifecycle is simulated here frame by frame.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import chunk_windows
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    AR_CHUNK_FRAMES,
+    AR_CHUNK_HOP_FRAMES,
+    NUM_CODEBOOKS,
+)
+from vllm_omni.model_executor.models.minimax_music3.frame_state import (
+    FrameBuffer,
+    RequestARState,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# One scalar per frame instead of 32768 dims: the bookkeeping is what is under
+# test, and a frame's value doubles as its identity.
+_WIDTH = 4
+
+
+def _frame(index: int) -> torch.Tensor:
+    return torch.full((1, _WIDTH), float(index))
+
+
+def _identities(chunk: torch.Tensor) -> list[int]:
+    return [int(row[0].item()) for row in chunk]
+
+
+class TestFrameBuffer:
+    def test_slice_returns_the_requested_frames(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        assert _identities(buffer.slice(3, 7)) == [3, 4, 5, 6]
+
+    def test_slice_is_relative_to_the_live_base(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(4)
+        assert buffer.base_frame == 4
+        assert _identities(buffer.slice(4, 8)) == [4, 5, 6, 7]
+
+    def test_slicing_a_discarded_range_raises(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(5)
+        with pytest.raises(RuntimeError, match="unavailable"):
+            buffer.slice(3, 6)
+
+    def test_slicing_past_the_end_raises(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(4):
+            buffer.append(_frame(index))
+        with pytest.raises(RuntimeError, match="unavailable"):
+            buffer.slice(2, 9)
+
+    def test_discard_is_clamped_at_both_ends(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(5):
+            buffer.append(_frame(index))
+        buffer.discard_before(-3)
+        assert buffer.base_frame == 0
+        buffer.discard_before(99)
+        assert buffer.base_frame == 5
+        assert buffer.end_frame == 5
+
+    def test_peak_tracks_the_high_water_mark(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(8)
+        assert buffer.peak_frames == 10
+
+
+class TestWindowLifecycle:
+    """Drive the state exactly as the talker does, one frame per step."""
+
+    @staticmethod
+    def _run(total_frames: int) -> tuple[list[list[int]], int]:
+        """Generate ``total_frames`` frames, returning the emitted spans.
+
+        Returns:
+            The frame identities of every emitted span, and the peak buffer
+            occupancy reached along the way.
+        """
+        state = RequestARState(seed=1, max_frames=total_frames)
+        spans: list[list[int]] = []
+        for index in range(total_frames):
+            state.frames.append(_frame(index))
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+            state.generated_frames += 1
+            window = state.ready_window()
+            if window is not None:
+                spans.append(_identities(state.take_window(window, is_final=False)))
+
+        # Terminal flush: everything not yet emitted, as one contiguous span.
+        start = state.emitted_windows * AR_CHUNK_HOP_FRAMES
+        end = state.frames.end_frame
+        if end > start:
+            spans.append(_identities(state.frames.slice(start, end)))
+        return spans, state.frames.peak_frames
+
+    def test_short_song_is_one_span(self) -> None:
+        spans, _ = self._run(50)
+        assert spans == [list(range(50))]
+
+    def test_every_frame_is_delivered_exactly_once_in_order(self) -> None:
+        """Spans overlap by design, but their union must be the whole song.
+
+        Reconstruct by taking each span's non-overlapping remainder, which is
+        what the acoustic stage does after cropping.
+        """
+        for total in (250, 437, 600, 901):
+            spans, _ = self._run(total)
+            rebuilt: list[int] = []
+            for span in spans:
+                rebuilt.extend(frame for frame in span if frame >= len(rebuilt))
+            assert rebuilt == list(range(total)), f"gap or duplicate at {total} frames"
+
+    def test_windows_advance_by_the_hop(self) -> None:
+        spans, _ = self._run(600)
+        for index, span in enumerate(spans[:-1]):
+            assert span[0] == index * AR_CHUNK_HOP_FRAMES
+            assert len(span) == AR_CHUNK_FRAMES
+
+    def test_mid_stream_windows_are_never_the_final_frame(self) -> None:
+        """A window is held back until a further frame proves it is not last."""
+        total = 600
+        spans, _ = self._run(total)
+        for span in spans[:-1]:
+            assert span[-1] < total - 1
+
+    def test_terminal_span_covers_the_tail(self) -> None:
+        total = 437
+        spans, _ = self._run(total)
+        assert spans[-1][-1] == total - 1
+
+    def test_terminal_span_is_bounded(self) -> None:
+        """The tail is at most one full window plus one hop."""
+        for total in (250, 301, 437, 600, 901, 1000):
+            spans, _ = self._run(total)
+            assert len(spans[-1]) <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+    def test_buffer_stays_bounded_over_a_long_song(self) -> None:
+        """Peak occupancy must not grow with song length; that is the point
+        of evicting retired frames."""
+        _, peak_short = self._run(600)
+        _, peak_long = self._run(3000)
+        assert peak_long == peak_short
+        assert peak_long <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+    def test_emitted_window_count_matches_the_window_plan(self) -> None:
+        """What the stream emitted must reconcile with chunk_windows()."""
+        for total in (250, 437, 600):
+            spans, _ = self._run(total)
+            planned = chunk_windows(total)
+            # Every window is accounted for: the pre-terminal spans are one
+            # window each, and the terminal span carries whatever remains.
+            assert len(spans) <= len(planned)
+            assert spans[0][0] == planned[0].start
+
+
+class TestSamplingPositions:
+    def test_each_frame_reserves_one_slot_per_codebook(self) -> None:
+        """Codebook streams must never collide, or codes repeat across frames."""
+        state = RequestARState(seed=1, max_frames=10)
+        positions = [state.advance_position() for _ in range(5)]
+        assert positions == [0, 8, 16, 24, 32]
+
+    def test_draw_slots_are_globally_unique(self) -> None:
+        state = RequestARState(seed=1, max_frames=64)
+        slots: set[int] = set()
+        for _ in range(64):
+            base = state.advance_position()
+            for codebook in range(NUM_CODEBOOKS):
+                slot = base + codebook
+                assert slot not in slots
+                slots.add(slot)
+        assert len(slots) == 64 * NUM_CODEBOOKS
+
+
+class TestCodeRetention:
+    def test_codes_span_matches_the_frame_span(self) -> None:
+        state = RequestARState(seed=1, max_frames=10)
+        for index in range(10):
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+        codes = state.take_codes(2, 6)
+        assert codes is not None
+        assert codes.shape == (4, NUM_CODEBOOKS)
+        assert [int(row[0]) for row in codes] == [2, 3, 4, 5]
+
+    def test_out_of_range_span_is_dropped_not_raised(self) -> None:
+        """Codes are diagnostic; the audio path must not fail without them."""
+        state = RequestARState(seed=1, max_frames=10)
+        for index in range(4):
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+        assert state.take_codes(2, 99) is None
+        assert state.take_codes(0, 0) is None

tests/model_executor/models/test_minimax_music3_prompt.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic.
+
+The prompt transformations are not cosmetic: they decide token positions, and
+token positions decide the audio. These tests pin the documented behaviours
+that a well-meaning cleanup would otherwise "fix".
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import (
+    chunk_windows,
+    crop_sample_bounds,
+    overlap_mel_length,
+)
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    AR_CHUNK_FRAMES,
+    AR_CHUNK_HOP_FRAMES,
+    BLEND_MEL_FRAMES,
+    DAV_HOP_SAMPLES,
+    HOP_MEL_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.prompt import (
+    SPECIAL_TOKEN_IDS,
+    build_cfg_null_token_ids,
+    build_prompt,
+    clean_caption,
+    normalize_lyrics,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestNormalizeLyrics:
+    def test_prepends_start_tag(self) -> None:
+        assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n")
+
+    def test_lowercases_structure_tags(self) -> None:
+        assert "[verse]" in normalize_lyrics("[Verse]\nhello")
+        assert "[Verse]" not in normalize_lyrics("[Verse]\nhello")
+
+    def test_text_after_a_leading_tag_is_dropped(self) -> None:
+        """Documented contract: a tag must sit on its own line.
+
+        Lyrics written next to a tag are silently lost, which is why the
+        serving docs tell callers to put tags on their own line.
+        """
+        out = normalize_lyrics("[Verse] Walking down the street")
+        assert "Walking" not in out
+        assert "[verse]" in out
+
+    def test_text_on_its_own_line_survives(self) -> None:
+        out = normalize_lyrics("[Verse]\nWalking down the street")
+        assert "Walking down the street" in out
+
+    def test_bracket_space_becomes_newline(self) -> None:
+        assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]"
+
+
+class TestCleanCaption:
+    def test_special_token_form_is_rewritten_not_passed_through(self) -> None:
+        """A caption must not be able to inject a real special token."""
+        assert clean_caption("<|bpm 92|>") == "bpm is 92"
+        assert "<|" not in clean_caption("<|bpm 92|>")
+
+    def test_lone_tag_keeps_its_inner_text(self) -> None:
+        assert clean_caption("<|instrumental|>") == "instrumental"
+
+    def test_markdown_is_stripped(self) -> None:
+        assert clean_caption("## Heading") == "Heading"
+        assert clean_caption("**bold**") == "bold"
+        assert clean_caption("- item") == "item"
+
+    def test_blank_lines_collapse_at_two_not_three(self) -> None:
+        """The source collapses runs of two or more, not three or more.
+
+        This shifts blank-line token positions and therefore the audio, so it
+        is pinned deliberately.
+        """
+        assert clean_caption("a\n\nb") == "a\nb"
+        assert clean_caption("a\n\n\n\nb") == "a\nb"
+
+
+class TestBuildPrompt:
+    def test_framing_tokens_are_present_and_ordered(self) -> None:
+        prompt = build_prompt("a bright pop song", "[Verse]\nhello")
+        assert prompt.startswith("<|im_start|><|caption_start|>")
+        assert prompt.endswith("<|im_end|><|audio_start|>")
+        assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>")
+
+
+class TestBuildCfgNullTokenIds:
+    @staticmethod
+    def _prompt_ids(body_len: int) -> list[int]:
+        return (
+            [SPECIAL_TOKEN_IDS["<|im_start|>"]]
+            + list(range(1000, 1000 + body_len))
+            + [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]]
+        )
+
+    def test_length_is_preserved(self) -> None:
+        """Equal length is what lets the two rows decode in lockstep."""
+        ids = self._prompt_ids(12)
+        assert len(build_cfg_null_token_ids(ids)) == len(ids)
+
+    def test_body_is_replaced_and_framing_survives(self) -> None:
+        ids = self._prompt_ids(12)
+        null = build_cfg_null_token_ids(ids)
+        assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"]
+        assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"]
+        assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"]
+        assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]}
+
+    def test_input_is_not_mutated(self) -> None:
+        ids = self._prompt_ids(6)
+        before = list(ids)
+        build_cfg_null_token_ids(ids)
+        assert ids == before
+
+    def test_minimal_prompt_with_empty_body(self) -> None:
+        ids = [
+            SPECIAL_TOKEN_IDS["<|im_start|>"],
+            SPECIAL_TOKEN_IDS["<|audio_cfg|>"],
+            SPECIAL_TOKEN_IDS["<|im_end|>"],
+            SPECIAL_TOKEN_IDS["<|audio_start|>"],
+        ]
+        assert build_cfg_null_token_ids(ids) == ids
+
+    def test_too_short_is_rejected(self) -> None:
+        with pytest.raises(ValueError, match="too short"):
+            build_cfg_null_token_ids([1, 2, 3])
+
+    def test_wrong_framing_is_rejected(self) -> None:
+        with pytest.raises(ValueError, match="framing"):
+            build_cfg_null_token_ids([9, 9, 9, 9, 9])
+
+
+class TestChunkWindows:
+    def test_no_frames_yields_nothing(self) -> None:
+        assert chunk_windows(0) == []
+
+    def test_short_song_is_one_final_window(self) -> None:
+        windows = chunk_windows(50)
+        assert len(windows) == 1
+        assert (windows[0].start, windows[0].end) == (0, 50)
+        assert windows[0].is_first and windows[0].is_last
+
+    def test_exactly_one_window(self) -> None:
+        windows = chunk_windows(AR_CHUNK_FRAMES)
+        assert len(windows) == 1
+        assert windows[0].is_last
+
+    def test_windows_advance_by_the_hop(self) -> None:
+        windows = chunk_windows(450)
+        starts = [w.start for w in windows]
+        assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES))
+        assert windows[-1].is_last
+        assert not any(w.is_last for w in windows[:-1])
+
+    def test_windows_cover_every_frame(self) -> None:
+        frames = 437
+        windows = chunk_windows(frames)
+        covered: set[int] = set()
+        for window in windows:
+            covered.update(range(window.start, window.end))
+        assert covered == set(range(frames))
+
+    def test_final_window_is_clamped_to_the_song(self) -> None:
+        windows = chunk_windows(437)
+        assert windows[-1].end == 437
+
+    def test_only_the_first_window_is_first(self) -> None:
+        windows = chunk_windows(600)
+        assert windows[0].is_first
+        assert not any(w.is_first for w in windows[1:])
+
+
+class TestCropBounds:
+    def test_single_window_song_is_not_cropped(self) -> None:
+        window = chunk_windows(50)[0]
+        assert crop_sample_bounds(window) == (0, 0)
+
+    def test_first_window_keeps_its_head(self) -> None:
+        left, right = crop_sample_bounds(chunk_windows(600)[0])
+        assert left == 0
+        assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES
+
+    def test_last_window_keeps_its_tail(self) -> None:
+        left, right = crop_sample_bounds(chunk_windows(600)[-1])
+        assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES
+        assert right == 0
+
+    def test_middle_window_is_cropped_on_both_sides(self) -> None:
+        windows = chunk_windows(900)
+        middle = windows[len(windows) // 2]
+        left, right = crop_sample_bounds(middle)
+        assert left > 0 and right > 0
+
+    def test_overlap_is_half_a_window(self) -> None:
+        assert overlap_mel_length() == HOP_MEL_FRAMES // 2

tests/model_executor/models/test_minimax_music3_row_binding.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 row-to-request binding.
+
+vLLM compacts and reorders the persistent batch whenever a request finishes or
+is preempted, so a batch row number is only meaningful within the step that
+produced it. Every failure pinned here was a real defect: state keyed by row
+survived into a step where the row belonged to someone else, or state keyed by
+one spelling of the request id was released under another and leaked.
+
+These exercise the pure helpers and the binding arithmetic directly, so they
+need no GPU and no engine.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    DEFAULT_MAX_AUDIO_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+    _CFG_UNCOND_SUFFIX,
+    _first_int,
+    _pair_id_from_request,
+    _request_key,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestRequestKey:
+    def test_engine_list_form_is_unwrapped(self) -> None:
+        """The engine injects global_request_id as a single-element list.
+
+        Keying state on the list's ``str()`` would produce "['req-1']", which
+        never matches the bare id that ``on_requests_finished`` delivers, so
+        every finished song would leak its frame buffer.
+        """
+        assert _request_key(["req-1"]) == "req-1"
+
+    def test_bare_string_is_unchanged(self) -> None:
+        assert _request_key("req-1") == "req-1"
+
+    def test_list_and_bare_forms_agree(self) -> None:
+        """Allocation and release must land on the same key."""
+        assert _request_key(["req-1"]) == _request_key("req-1")
+
+    def test_empty_and_missing_are_stable(self) -> None:
+        assert _request_key(None) == ""
+        assert _request_key([]) == ""
+
+    def test_tuple_form(self) -> None:
+        assert _request_key(("req-9",)) == "req-9"
+
+
+class TestPairIdDerivation:
+    def test_conditioned_request_is_its_own_pair(self) -> None:
+        assert _pair_id_from_request("req-1") == "req-1"
+
+    def test_companion_maps_to_its_partner(self) -> None:
+        assert _pair_id_from_request(f"req-1{_CFG_UNCOND_SUFFIX}") == "req-1"
+
+    def test_pair_members_agree_on_the_pair_id(self) -> None:
+        cond = "req-42"
+        uncond = f"req-42{_CFG_UNCOND_SUFFIX}"
+        assert _pair_id_from_request(cond) == _pair_id_from_request(uncond)
+
+    def test_a_request_ending_in_a_similar_string_is_not_confused(self) -> None:
+        assert _pair_id_from_request("req-cfg_uncond") == "req-cfg_uncond"
+
+
+class TestFirstInt:
+    def test_serving_layer_list_wrapping(self) -> None:
+        """tts_params values arrive wrapped in a single-element list."""
+        assert _first_int([750], DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+    def test_bare_value(self) -> None:
+        assert _first_int(750, DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+    def test_missing_falls_back(self) -> None:
+        assert _first_int(None, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int([], DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+    def test_nonsense_falls_back_rather_than_raising(self) -> None:
+        assert _first_int("abc", DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int(0, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int(-5, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+
+def _feedback_plan(
+    request_ids: list[str],
+    computed: list[int],
+    scheduled: list[int],
+    prompt_tokens: dict[str, int],
+    has_codes: set[str],
+) -> list[int]:
+    """Mirror of the talker's feedback row selection.
+
+    Returns the flat token offsets whose embedding gets replaced by an audio
+    frame. Kept as a standalone function so the arithmetic can be checked
+    without constructing a model.
+    """
+    offsets: list[int] = []
+    running = 0
+    for count in scheduled:
+        offsets.append(running)
+        running += count
+
+    rows: list[int] = []
+    for row, request_id in enumerate(request_ids):
+        if request_id not in has_codes:
+            continue
+        if scheduled[row] != 1 or computed[row] < prompt_tokens.get(request_id, 0):
+            continue
+        rows.append(offsets[row])
+    return rows
+
+
+class TestFeedbackRowSelection:
+    """The guard that decides which token embeddings become audio frames."""
+
+    def test_decode_only_step_replaces_every_row(self) -> None:
+        ids = ["a", "a__cfg_uncond", "b", "b__cfg_uncond"]
+        rows = _feedback_plan(
+            ids,
+            computed=[120, 120, 90, 90],
+            scheduled=[1, 1, 1, 1],
+            prompt_tokens={i: 50 for i in ids},
+            has_codes=set(ids),
+        )
+        assert rows == [0, 1, 2, 3]
+
+    def test_prefill_row_is_never_overwritten(self) -> None:
+        """A prompt token must keep its real embedding.
+
+        The earlier guard compared the flat token count against a row budget,
+        so a short prompt slipped through and had its first embeddings
+        replaced with a previous request's audio frame.
+        """
+        rows = _feedback_plan(
+            ["a", "b"],
+            computed=[0, 300],
+            scheduled=[8, 1],
+            prompt_tokens={"a": 8, "b": 40},
+            has_codes={"b"},
+        )
+        # Row 0 occupies flat tokens 0..7 and must be untouched; row 1's single
+        # decode token sits at flat offset 8.
+        assert rows == [8]
+
+    def test_short_prefill_below_the_row_budget_is_still_safe(self) -> None:
+        """The exact shape that defeated the old token-count guard."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[0],
+            scheduled=[10],
+            prompt_tokens={"a": 10},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_chunked_prefill_midway_is_not_fed_back(self) -> None:
+        rows = _feedback_plan(
+            ["a"],
+            computed=[2048],
+            scheduled=[2048],
+            prompt_tokens={"a": 5000},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_final_prefill_chunk_is_not_fed_back(self) -> None:
+        """The step that completes a prompt still carries prompt tokens."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[4096],
+            scheduled=[904],
+            prompt_tokens={"a": 5000},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_mixed_batch_offsets_follow_the_flat_layout(self) -> None:
+        """Tokens are a flat concatenation, so offsets are a running sum."""
+        ids = ["p", "d1", "d2"]
+        rows = _feedback_plan(
+            ids,
+            computed=[0, 500, 700],
+            scheduled=[64, 1, 1],
+            prompt_tokens={"p": 64, "d1": 40, "d2": 40},
+            has_codes={"d1", "d2"},
+        )
+        assert rows == [64, 65]
+
+    def test_request_without_codes_yet_is_skipped(self) -> None:
+        """The first decode step after prefill has no previous frame."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[50],
+            scheduled=[1],
+            prompt_tokens={"a": 50},
+            has_codes=set(),
+        )
+        assert rows == []
+
+    def test_reordered_batch_binds_by_request_not_position(self) -> None:
+        """The same two requests in swapped rows get swapped offsets.
+
+        This is the compaction case: holding codes in a row-indexed buffer
+        would have handed each request the other's frame.
+        """
+        prompt_tokens = {"a": 50, "b": 50}
+        first = _feedback_plan(["a", "b"], [100, 200], [1, 1], prompt_tokens, {"a"})
+        swapped = _feedback_plan(["b", "a"], [200, 100], [1, 1], prompt_tokens, {"a"})
+        assert first == [0]
+        assert swapped == [1]

tests/model_executor/test_cfg_pairing.py

@@ -0,0 +1,389 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the model-neutral CFG request-pairing scheduler patches.
+
+The patches keep a cond/uncond request pair step-locked inside one engine:
+admitted together, kept at equal ``num_computed_tokens``, finished together.
+They are installed onto a lightweight fake scheduler here; vLLM's real
+``Scheduler`` needs a full engine (and a GPU) to construct.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vllm_omni.model_executor.models.common import cfg_pairing
+from vllm_omni.model_executor.models.common.cfg_pairing import (
+    _equalize_cfg_pair_progress,
+    _patch_scheduler,
+    cfg_uncond_suffixes,
+    register_cfg_uncond_suffix,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ───────── fakes ─────────
+
+
+class _FakeRequest:
+    """Minimal stand-in for ``vllm.v1.request.Request`` (identity equality)."""
+
+    def __init__(self, request_id: str, extra_args: dict[str, Any] | None = None, step_tokens: int = 1) -> None:
+        self.request_id = request_id
+        self.sampling_params = SimpleNamespace(extra_args=extra_args)
+        self.num_computed_tokens = 0
+        self.step_tokens = step_tokens
+        self.skip_reading_prefix_cache = False
+
+
+class _FakeWaitingQueue(deque):
+    """FCFS waiting queue: vLLM's is also a ``deque`` subclass."""
+
+    def remove_requests(self, requests: list[_FakeRequest]) -> None:
+        keep = [req for req in self if not any(req is dropped for dropped in requests)]
+        self.clear()
+        self.extend(keep)
+
+    def prepend_request(self, request: _FakeRequest) -> None:
+        self.appendleft(request)
+
+
+@dataclass
+class _FakeSchedulerOutput:
+    num_scheduled_tokens: dict[str, int] = field(default_factory=dict)
+    total_num_scheduled_tokens: int = 0
+
+
+def _make_scheduler_cls() -> type:
+    """Fresh class per test so the in-place patches never stack."""
+
+    class FakeScheduler:
+        def __init__(self) -> None:
+            self.waiting = _FakeWaitingQueue()
+            self.running: list[_FakeRequest] = []
+            self.requests: dict[str, _FakeRequest] = {}
+            self.finished: list[set[str]] = []
+            self.observed_waiting: list[list[str]] = []
+            self.max_num_scheduled_tokens = 512
+            self.scheduler_config = SimpleNamespace(long_prefill_token_threshold=4096)
+
+        def add_request(self, request: _FakeRequest) -> None:
+            self.requests[request.request_id] = request
+            self.waiting.append(request)
+
+        def finish_requests(self, request_ids: Any, finished_status: Any) -> None:
+            if request_ids is None:
+                ids = set(self.requests)
+            elif isinstance(request_ids, str):
+                ids = {request_ids}
+            else:
+                ids = set(request_ids)
+            self.finished.append(ids)
+            self.running = [req for req in self.running if req.request_id not in ids]
+            for req_id in ids:
+                self.requests.pop(req_id, None)
+
+        def schedule(self) -> _FakeSchedulerOutput:
+            self.observed_waiting.append([req.request_id for req in self.waiting])
+            while self.waiting:
+                self.running.append(self.waiting.popleft())
+            output = _FakeSchedulerOutput()
+            for req in self.running:
+                req.num_computed_tokens += req.step_tokens
+                output.num_scheduled_tokens[req.request_id] = req.step_tokens
+                output.total_num_scheduled_tokens += req.step_tokens
+            return output
+
+    return FakeScheduler
+
+
+@pytest.fixture
+def scheduler() -> Any:
+    cls = _make_scheduler_cls()
+    _patch_scheduler(cls)
+    return cls()
+
+
+@pytest.fixture(autouse=True)
+def _restore_suffix_registry():
+    original = set(cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES)
+    yield
+    cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.clear()
+    cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.update(original)
+
+
+def _add_pair(
+    scheduler: Any,
+    pair_id: str | None = None,
+    base_id: str = "req-1",
+    suffix: str = "__cfg_uncond",
+) -> tuple[_FakeRequest, _FakeRequest]:
+    cond_extra: dict[str, Any] = {"cfg_role": "cond"}
+    uncond_extra: dict[str, Any] = {"cfg_role": "uncond"}
+    if pair_id is not None:
+        cond_extra["cfg_pair_id"] = pair_id
+        uncond_extra["cfg_pair_id"] = pair_id
+    cond = _FakeRequest(base_id, cond_extra)
+    uncond = _FakeRequest(f"{base_id}{suffix}", uncond_extra)
+    scheduler.add_request(cond)
+    scheduler.add_request(uncond)
+    return cond, uncond
+
+
+# ───────── pair-id derivation ─────────
+
+
+def test_pair_id_is_derived_from_request_id_when_absent(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+
+    assert scheduler._cfg_pairs == {"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}}
+    assert scheduler._cfg_req_to_pair == {cond.request_id: "req-1", uncond.request_id: "req-1"}
+
+
+def test_registered_companion_suffix_is_stripped(scheduler: Any) -> None:
+    register_cfg_uncond_suffix("::negative")
+    assert "::negative" in cfg_uncond_suffixes()
+
+    cond, uncond = _add_pair(scheduler, base_id="music-7", suffix="::negative")
+
+    assert scheduler._cfg_pairs == {"music-7": {"cond": cond.request_id, "uncond": uncond.request_id}}
+
+
+def test_register_rejects_empty_suffix() -> None:
+    with pytest.raises(ValueError):
+        register_cfg_uncond_suffix("")
+
+
+def test_uncond_without_known_suffix_keeps_its_own_request_id(scheduler: Any) -> None:
+    # Degenerate input: the companion never joins the cond pair, so the pair
+    # stays incomplete and the hold budget eventually releases it.
+    scheduler.add_request(_FakeRequest("solo-uncond", {"cfg_role": "uncond"}))
+
+    assert scheduler._cfg_pairs == {"solo-uncond": {"uncond": "solo-uncond"}}
+
+
+def test_explicit_pair_id_takes_precedence(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler, pair_id="pair-A")
+
+    assert scheduler._cfg_pairs == {"pair-A": {"cond": cond.request_id, "uncond": uncond.request_id}}
+    assert cond.request_id not in scheduler._cfg_pairs
+    assert scheduler._cfg_req_to_pair[uncond.request_id] == "pair-A"
+
+
+def test_cfg_requests_skip_prefix_cache_reads(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+
+    assert cond.skip_reading_prefix_cache is True
+    assert uncond.skip_reading_prefix_cache is True
+
+
+# ───────── admission holding ─────────
+
+
+def test_lone_member_is_held_then_released(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+
+    scheduler.schedule()
+
+    assert scheduler.observed_waiting == [[]]  # the base scheduler saw nothing
+    assert [req.request_id for req in scheduler.waiting] == ["req-1"]
+    assert scheduler.running == []
+
+
+def test_complete_pair_is_admitted_together(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+    scheduler.schedule()
+    scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
+
+    scheduler.schedule()
+
+    assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond"]
+    assert [req.request_id for req in scheduler.running] == ["req-1", "req-1__cfg_uncond"]
+    assert list(scheduler.waiting) == []
+
+
+def test_hold_budget_expiry_drops_the_stale_pair(scheduler: Any, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(cfg_pairing, "_MAX_PAIR_HOLD_STEPS", 2)
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+
+    scheduler.schedule()
+    scheduler.schedule()
+    assert scheduler.running == []
+    assert scheduler._cfg_pairs  # still hoping for the partner
+
+    scheduler.schedule()
+
+    assert scheduler._cfg_pairs == {}
+    assert scheduler._cfg_req_to_pair == {}
+    assert [req.request_id for req in scheduler.running] == ["req-1"]
+
+
+def test_waiting_queue_reorder_puts_partners_adjacent(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+    scheduler.add_request(_FakeRequest("other", None))
+    scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
+
+    scheduler.schedule()
+
+    assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond", "other"]
+
+
+# ───────── progress equalization ─────────
+
+
+def _equalize_fixture(
+    cond_computed: int,
+    uncond_computed: int,
+    cond_sched: int,
+    uncond_sched: int,
+) -> tuple[Any, _FakeSchedulerOutput, _FakeRequest, _FakeRequest]:
+    cond = _FakeRequest("req-1", {"cfg_role": "cond"})
+    uncond = _FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})
+    cond.num_computed_tokens = cond_computed
+    uncond.num_computed_tokens = uncond_computed
+    fake = SimpleNamespace(
+        _cfg_pairs={"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}},
+        requests={cond.request_id: cond, uncond.request_id: uncond},
+    )
+    output = _FakeSchedulerOutput(
+        num_scheduled_tokens={cond.request_id: cond_sched, uncond.request_id: uncond_sched},
+        total_num_scheduled_tokens=cond_sched + uncond_sched,
+    )
+    return fake, output, cond, uncond
+
+
+def test_equalize_pulls_the_ahead_member_back() -> None:
+    fake, output, cond, uncond = _equalize_fixture(10, 6, 5, 4)
+
+    _equalize_cfg_pair_progress(fake, output)
+
+    assert cond.num_computed_tokens == 6
+    assert uncond.num_computed_tokens == 6
+    assert output.num_scheduled_tokens == {cond.request_id: 1, uncond.request_id: 4}
+    assert output.total_num_scheduled_tokens == 5
+
+
+def test_equalize_is_a_noop_when_the_pair_is_already_aligned() -> None:
+    fake, output, cond, uncond = _equalize_fixture(8, 8, 4, 4)
+
+    _equalize_cfg_pair_progress(fake, output)
+
+    assert (cond.num_computed_tokens, uncond.num_computed_tokens) == (8, 8)
+    assert output.total_num_scheduled_tokens == 8
+
+
+def test_equalize_skips_an_infeasible_rollback() -> None:
+    # Clawing back 4 tokens would leave the cond row with 0 scheduled tokens,
+    # which the scheduler output cannot express; leave the pair to _drop_split_pairs.
+    fake, output, cond, uncond = _equalize_fixture(10, 6, 4, 4)
+
+    _equalize_cfg_pair_progress(fake, output)
+
+    assert cond.num_computed_tokens == 10
+    assert output.total_num_scheduled_tokens == 8
+
+
+# ───────── finish ─────────
+
+
+def test_finish_requests_finishes_the_partner(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+    scheduler.schedule()
+
+    scheduler.finish_requests(cond.request_id, "FINISHED_STOPPED")
+
+    assert scheduler.finished == [{cond.request_id, uncond.request_id}]
+    assert scheduler._cfg_pairs == {}
+    assert scheduler._cfg_req_to_pair == {}
+    assert scheduler.running == []
+
+
+def test_finish_all_requests_clears_the_registry(scheduler: Any) -> None:
+    _add_pair(scheduler)
+    scheduler.schedule()
+
+    scheduler.finish_requests(None, "FINISHED_ABORTED")
+
+    assert scheduler._cfg_pairs == {}
+    assert scheduler._cfg_req_to_pair == {}
+
+
+# ───────── non-CFG traffic ─────────
+
+
+def test_plain_requests_are_untouched(scheduler: Any) -> None:
+    scheduler.add_request(_FakeRequest("plain-1", None))
+    scheduler.add_request(_FakeRequest("plain-2", {"seed": 3}))
+
+    output = scheduler.schedule()
+
+    assert scheduler._cfg_pairs == {}
+    assert scheduler.observed_waiting == [["plain-1", "plain-2"]]
+    assert output.num_scheduled_tokens == {"plain-1": 1, "plain-2": 1}
+    assert output.total_num_scheduled_tokens == 2
+    # No pair ever existed, so the hold bookkeeping is never allocated.
+    assert not hasattr(scheduler, "_cfg_hold_counts")
+
+
+def test_default_suffix_matches_the_audex_stage_processor() -> None:
+    from vllm_omni.model_executor.stage_input_processors.audex import CFG_UNCOND_SUFFIX
+
+    assert CFG_UNCOND_SUFFIX in cfg_uncond_suffixes()
+
+
+def test_audex_cfg_module_still_exports_the_patch_entry_points() -> None:
+    from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches, cfg_patches_applied
+
+    assert apply_cfg_patches is cfg_pairing.apply_cfg_patches
+    assert cfg_patches_applied is cfg_pairing.cfg_patches_applied
+
+
+# ───────── install gate ─────────
+
+
+def test_gate_triggers_on_a_cfg_logits_processor() -> None:
+    from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+    class AudexCFGLogitsProcessor:
+        pass
+
+    assert _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[AudexCFGLogitsProcessor]))
+    # vLLM keeps unresolved entries as dotted strings.
+    assert _stage_declares_cfg_pairs(
+        SimpleNamespace(logits_processors=["vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"])
+    )
+
+
+def test_gate_triggers_on_a_declared_cfg_role_default() -> None:
+    from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+    assert _stage_declares_cfg_pairs(
+        SimpleNamespace(logits_processors=None, sampling_extra_args_keys=("cfg_role", "cfg_scale"))
+    )
+
+
+def test_gate_stays_off_for_non_cfg_stages() -> None:
+    from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+    assert not _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[], sampling_extra_args_keys=()))
+    assert not _stage_declares_cfg_pairs(
+        SimpleNamespace(logits_processors=["some.module:MinPLogitsProcessor"], sampling_extra_args_keys=("tta_rvq",))
+    )
+    assert not _stage_declares_cfg_pairs(SimpleNamespace())
+
+
+def test_sampling_extra_args_keys_are_derived_from_stage_defaults() -> None:
+    from vllm_omni.engine.stage_init_utils import _sampling_extra_args_keys
+
+    assert _sampling_extra_args_keys({"extra_args": {"cfg_scale": 1.5, "cfg_role": "cond"}}) == (
+        "cfg_role",
+        "cfg_scale",
+    )
+    assert _sampling_extra_args_keys({"temperature": 0.1}) == ()
+    assert _sampling_extra_args_keys(None) == ()

tools/nightly/run_nightly_jobs.sh

@@ -43,7 +43,7 @@
 #
 #   stability (when included in TEST_TYPE):
 #     From repo root: pytest -s -v --run-level full_model tests/dfx/stability/scripts/test_stability_*.py
-#     model_type: omni → qwen3_omni; tts → qwen3_tts + voxcpm2; diffusion → qwen_image + wan22 + hunyuan_image; all → all six
+#     model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22 + hunyuan_image; all → all five
 #     LABEL_SUBSTR: if set, script path / job key / filename must contain it
 #
 # Requirements: bash, python3, PyYAML (pip install pyyaml)
@@ -588,7 +588,6 @@ def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -
 STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [
     ("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)),
     ("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)),
-    ("stability_voxcpm2", "tests/dfx/stability/scripts/test_stability_voxcpm2.py", ("tts",)),
     ("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)),
     ("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)),
     ("stability_hunyuan_image", "tests/dfx/stability/scripts/test_stability_hunyuan_image.py", ("diffusion",)),

vllm_omni/benchmarks/data_modules/daily_omni_dataset.py

@@ -7,7 +7,8 @@ Dataset source: https://huggingface.co/datasets/liarliar/Daily-Omni
 
 Supports loading QA metadata from:
 - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
-- HuggingFace datasets (``dataset_path``): legacy online mode
+- HuggingFace Hub id (``dataset_path``): resolved with ``snapshot_download`` to the repo's
+  ``qa.json``, which reads the hub cache and therefore also works under ``HF_HUB_OFFLINE``
 
 Video/audio files normally come from extracted ``Videos.tar``. When ``--daily-omni-video-dir``
 is not set, the first request that needs on-disk media downloads that archive from the Hugging Face
@@ -24,9 +25,9 @@ Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``?
     This class therefore inherits only ``BenchmarkDataset`` (minimal: ``dataset_path``,
     ``random_seed``, ``self.data``) and implements **two explicit loaders**:
     ``_load_from_local_json`` (default path for air-gapped runs) and ``_load_from_huggingface``
-    (optional legacy path for users who prefer ``datasets`` + Hub cache). The latter is **not**
+    (Hub ids, resolved to the repo's ``qa.json`` through the hub cache). The latter is **not**
     inheritance; it is the same Hub rows as before, factored into a helper so one class can
-    serve both deployment modes without mandatory ``datasets`` when using ``qa_json_path``.
+    serve both deployment modes without mandatory ``datasets``.
 
 Usage:
     from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset
@@ -38,7 +39,7 @@ Usage:
         random_seed=42,
     )
 
-    # HuggingFace mode (legacy, requires network)
+    # Hub id mode (works offline once the repo is in the HF hub cache)
     dataset = DailyOmniDataset(
         dataset_path="liarliar/Daily-Omni",
         dataset_split="train",
@@ -210,6 +211,68 @@ def _resolve_hf_cache_snapshot(path: Path) -> Path:
     return path
 
 
+#: QA metadata pulled by :func:`ensure_daily_omni_hub_root`. ``Videos.tar`` is deliberately left
+#: out: it is ~60GB and only :func:`ensure_daily_omni_hub_videos_dir` needs it, lazily, once a
+#: request actually asks for media.
+_DAILY_OMNI_QA_ALLOW_PATTERNS = ["qa.json", "QA.json"]
+
+
+def ensure_daily_omni_hub_root(repo_id: str) -> Path:
+    """Resolve a Daily-Omni Hub id to its snapshot directory, downloading QA metadata if needed.
+
+    Mirrors :func:`...seed_tts_dataset.resolve_seed_tts_root`: ``snapshot_download`` resolves the
+    repo through the **hub cache**, so an air-gapped run (``HF_HUB_OFFLINE=1``) keeps working as
+    long as the dataset was pre-staged there. ``datasets.load_dataset`` cannot do this — it only
+    consults the separate *datasets* cache, which a ``hf download`` never populates.
+
+    Raises:
+        ImportError: if ``huggingface_hub`` is not installed.
+        ValueError: if ``repo_id`` is empty.
+    """
+    rid = (repo_id or "").strip()
+    if not rid:
+        raise ValueError("repo_id is required to download Daily-Omni from Hugging Face")
+
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as e:
+        raise ImportError(
+            "Install huggingface_hub to load Daily-Omni from the Hub, or pass a local "
+            "--daily-omni-qa-json / --daily-omni-video-dir mirror."
+        ) from e
+
+    cache = snapshot_download(
+        repo_id=rid,
+        repo_type="dataset",
+        allow_patterns=_DAILY_OMNI_QA_ALLOW_PATTERNS,
+    )
+    root = _resolve_hf_cache_snapshot(Path(cache).resolve())
+    logger.info("Daily-Omni Hub snapshot ready at %s (repo=%s)", root, rid)
+    return root
+
+
+def _daily_omni_qa_load_error(
+    repo_id: str | None,
+    snapshot_error: Exception | None,
+    datasets_error: Exception,
+) -> RuntimeError:
+    """Compose an actionable error after both QA-metadata sources failed.
+
+    Offline CI hits this when the dataset is missing from the hub cache *or* staged under a
+    different repo id, so the message names both attempts and the flags that bypass them.
+    """
+    parts = [f"Could not load Daily-Omni QA metadata for {repo_id!r}."]
+    if snapshot_error is not None:
+        parts.append(f"huggingface_hub snapshot_download failed: {type(snapshot_error).__name__}: {snapshot_error}")
+    parts.append(f"datasets.load_dataset failed: {type(datasets_error).__name__}: {datasets_error}")
+    parts.append(
+        "For offline runs, pre-stage the dataset in the HF hub cache "
+        "(`hf download --repo-type dataset <repo>`) and make sure the repo id matches, "
+        "or pass --daily-omni-qa-json / --daily-omni-video-dir pointing at a local mirror."
+    )
+    return RuntimeError(" ".join(parts))
+
+
 def resolve_daily_omni_local_root(dataset_path: str | None) -> Path | None:
     """Return the local Daily-Omni mirror root for ``dataset_path``, or ``None`` if not local.
 
@@ -433,7 +496,8 @@ class DailyOmniDataset(BenchmarkDataset):
 
     QA metadata can be loaded from:
     - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
-    - HuggingFace datasets (``dataset_path``): legacy online mode
+    - HuggingFace Hub id (``dataset_path``): the repo's ``qa.json`` via ``snapshot_download``
+      (hub-cache backed, offline-safe); ``datasets.load_dataset`` only as a parquet fallback
 
     Video/audio files normally come from extracted ``Videos.tar``. When ``video_dir`` is not set,
     the first sample that needs on-disk media downloads that archive from the Hugging Face dataset
@@ -444,7 +508,8 @@ class DailyOmniDataset(BenchmarkDataset):
         qa_json_path: Path to local qa.json file (offline mode, preferred). When provided,
             ``dataset_path`` and ``dataset_split`` are ignored.
         dataset_path: HuggingFace dataset path (e.g., "liarliar/Daily-Omni"). Used only if
-            ``qa_json_path`` is not provided (legacy online mode).
+            ``qa_json_path`` is not provided; resolved through the hub cache, so a pre-staged
+            repo also loads with ``HF_HUB_OFFLINE=1``.
         dataset_split: Dataset split to use (default: "train"). Used only in online mode.
         random_seed: Random seed for shuffling
         video_dir: Directory containing extracted video files (default: None; may be filled lazily
@@ -602,21 +667,42 @@ class DailyOmniDataset(BenchmarkDataset):
         self.data = _ListDatasetIterator(data)
 
     def _load_from_huggingface(self) -> None:
-        """Load QA rows via ``datasets.load_dataset`` (legacy / convenience path).
+        """Load QA rows for a Hub id, preferring the repo's ``qa.json`` over ``datasets``.
 
-        Kept for backward compatibility: callers can still pass ``dataset_path=liarliar/Daily-Omni``
-        and get the same parquet-backed rows as the Hub dataset card, with streaming (or
-        non-streaming if ``no_stream=True``) and shuffle.
+        The Hub repo ships ``qa.json`` (1,197 rows) next to ``Videos.tar``, so the primary path is
+        :func:`ensure_daily_omni_hub_root` → :meth:`_load_from_local_json`: ``snapshot_download``
+        reads the hub cache and therefore works under ``HF_HUB_OFFLINE``, the same contract as
+        Seed-TTS. ``datasets.load_dataset`` remains only as a fallback for mirrors
+        that publish parquet rows instead of ``qa.json``; it consults the *datasets* cache rather
+        than the hub cache and streams shards over HTTP, so it cannot serve air-gapped runs.
 
         This is intentionally **not** implemented by subclassing ``HuggingFaceDataset``: that base
         always runs Hub ``load_dataset`` from its constructor and expects a Hub id as the primary
         API; Daily-Omni instead chooses the source in ``load_data()`` (JSON vs Hub) while sharing
         one ``sample()`` / request-building implementation for both.
         """
+        snapshot_error: Exception | None = None
+        try:
+            root = ensure_daily_omni_hub_root(self.dataset_path)
+        except Exception as e:  # noqa: BLE001 - any failure just falls through to `datasets`
+            snapshot_error = e
+        else:
+            qa_json = daily_omni_local_qa_json(root)
+            if qa_json is not None:
+                self.qa_json_path = qa_json
+                logger.info("Loading Daily-Omni QA metadata from Hub snapshot: %s", qa_json)
+                self._load_from_local_json()
+                return
+            logger.info(
+                "Hub snapshot %s ships no qa.json; falling back to `datasets.load_dataset`",
+                root,
+            )
+
         if load_dataset is None:
-            raise ImportError(
-                "datasets library is required for HuggingFace mode. "
-                "Install with: pip install datasets, or use local JSON mode instead."
+            raise _daily_omni_qa_load_error(
+                self.dataset_path,
+                snapshot_error,
+                ImportError("datasets library is not installed (pip install datasets)"),
             )
 
         load_kw: dict[str, Any] = {
@@ -626,7 +712,10 @@ class DailyOmniDataset(BenchmarkDataset):
         }
         if self.dataset_subset is not None:
             load_kw["name"] = self.dataset_subset
-        ds = load_dataset(self.dataset_path, **load_kw)
+        try:
+            ds = load_dataset(self.dataset_path, **load_kw)
+        except Exception as e:
+            raise _daily_omni_qa_load_error(self.dataset_path, snapshot_error, e) from e
         if not getattr(self, "disable_shuffle", False):
             ds = ds.shuffle(seed=self.random_seed)
         self.data = ds

vllm_omni/config/__init__.py

@@ -2,7 +2,6 @@
 Configuration module for vLLM-Omni.
 """
 
-from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.config.omni_config import (
@@ -22,7 +21,6 @@ from vllm_omni.config.omni_config import (
     VllmOmniGenerationStageConfig,
     VllmOmniOrchestratorConfig,
 )
-from vllm_omni.config.pipeline_registry import register_pipeline
 from vllm_omni.config.stage_config import (
     PIPELINE_WIDE_ENGINE_FIELDS,
     DeployConfig,
@@ -42,6 +40,15 @@ from vllm_omni.config.yaml_util import (
     to_dict,
 )
 
+# StageConfigFactory / register_pipeline pull pipeline_registry, which eagerly
+# imports PI0_PIPELINE → diffusion.data. Keep those lazy so
+# `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still
+# loading) cannot close a circular import through DiffusionOutput.
+_LAZY_ATTRS = {
+    "StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"),
+    "register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"),
+}
+
 __all__ = [
     # Legacy model-level configs.
     "LoRAConfig",
@@ -82,3 +89,19 @@ __all__ = [
     "merge_configs",
     "to_dict",
 ]
+
+
+def __getattr__(name: str):
+    spec = _LAZY_ATTRS.get(name)
+    if spec is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr_name = spec
+    import importlib
+
+    value = getattr(importlib.import_module(module_name), attr_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(__all__) | set(globals()))

vllm_omni/config/model.py

@@ -146,6 +146,11 @@ class OmniModelConfig(ModelConfig):
     task_type: str | None = None
     enable_sleep_mode: bool = False
     has_sampling_extra_args: bool = False
+    # Key names (not values) of the stage's default sampling ``extra_args``.
+    # Engine-core code runs before any request arrives, so this is the only
+    # place it can learn which request-shaping conventions a stage uses (e.g.
+    # ``cfg_role`` for classifier-free-guidance request pairs).
+    sampling_extra_args_keys: tuple[str, ...] = ()
 
     @property
     def registry(self):

vllm_omni/config/omni_config.py

@@ -11,13 +11,24 @@ later PRs cut consumers over to these classes.
 from __future__ import annotations
 
 import copy
+import math
 from collections.abc import Mapping
-from dataclasses import dataclass, field, fields
+from dataclasses import InitVar, dataclass, field, fields
+from functools import wraps
+from inspect import Parameter, signature
 from pathlib import Path
 from typing import Any, Literal, TypeAlias, TypedDict, cast
 
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict, Field, model_validator
+from typing_extensions import Self
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
 from vllm.config.utils import config
+from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.config.stage_config import (
@@ -77,6 +88,49 @@ class _QuantizationEngineOverrides(TypedDict, total=False):
     quantization: str
 
 
+class _TrackExplicitConfigFields:
+    """Record constructor inputs without adding a serialized config field."""
+
+    @model_validator(mode="wrap")
+    @classmethod
+    def _record_explicit_config_fields(cls, value: Any, handler: Any) -> Any:
+        result = handler(value)
+        kwargs = getattr(value, "kwargs", None)
+        if kwargs is not None:
+            explicit_fields = frozenset(kwargs)
+        elif isinstance(value, cls):
+            explicit_fields = getattr(value, "_omni_explicit_fields", frozenset())
+        else:
+            explicit_fields = frozenset()
+        object.__setattr__(result, "_omni_explicit_fields", explicit_fields)
+        return result
+
+
+def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]:
+    """Make inherited Pydantic dataclass fields keyword-only as well."""
+    generated_init = cls.__init__
+    generated_signature = signature(cls)
+    keyword_only_signature = generated_signature.replace(
+        parameters=[
+            parameter.replace(kind=Parameter.KEYWORD_ONLY)
+            if parameter.kind in {Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD}
+            else parameter
+            for parameter in generated_signature.parameters.values()
+        ]
+    )
+
+    @wraps(generated_init)
+    def keyword_only_init(self: Any, *args: Any, **kwargs: Any) -> None:
+        if args:
+            raise TypeError(f"{cls.__name__}() accepts keyword arguments only; got {len(args)} positional argument(s)")
+        generated_init(self, **kwargs)
+
+    cls.__init__ = keyword_only_init
+    cls.__signature__ = keyword_only_signature
+    cls.__match_args__ = ()
+    return cls
+
+
 class _ModelEngineOverrides(TypedDict, total=False):
     model: str
     model_arch: str
@@ -106,7 +160,6 @@ class _ModelEngineOverrides(TypedDict, total=False):
     enforce_eager: bool
     max_cudagraph_capture_size: int
     enable_flashinfer_autotune: bool
-    compilation_config: dict[str, Any]
     enable_multithread_weight_load: bool
     num_weight_load_threads: int
     disable_autocast: bool
@@ -148,7 +201,6 @@ class _RuntimeEngineOverrides(TypedDict, total=False):
     num_gpus: int
     log_level: str
     log_stats: bool
-    profiler_config: dict[str, Any]
 
 
 class _ParallelConfigEngineOverrides(TypedDict, total=False):
@@ -192,6 +244,8 @@ class _StageEngineValues:
     runtime: _RuntimeEngineOverrides
     parallel: _ParallelEngineOverrides
     diffusion: _DiffusionEngineOverrides
+    compilation_config: Mapping[str, Any] | VllmCompilationConfig | None
+    profiler_config: Mapping[str, Any] | VllmProfilerConfig | None
 
 
 @dataclass(frozen=True)
@@ -341,7 +395,6 @@ class OmniStageModelConfig:
     enforce_eager: bool = False
     max_cudagraph_capture_size: int | None = Field(default=None, ge=0)
     enable_flashinfer_autotune: bool | None = None
-    compilation_config: dict[str, Any] | None = None
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = Field(default=4, ge=1)
     disable_autocast: bool = False
@@ -352,21 +405,21 @@ class OmniStageModelConfig:
     tokenizer_subdir: str | None = None
 
 
-@config
-class OmniStageLoadConfig:
-    """Per-stage loading behavior and resolved tokenizer input."""
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig):
+    """vLLM loading behavior plus Omni stage-specific tokenizer inputs."""
 
     tokenizer: str | None = None
-    download_dir: str | None = None
     skip_tokenizer_init: bool = False
-    load_format: str = "auto"
     tokenizer_mode: str = "auto"
     config_format: str | None = None
     skip_mm_profiling: bool | None = None
 
 
-@config
-class OmniStageCacheConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageCacheConfig(_TrackExplicitConfigFields, VllmCacheConfig):
     """Per-stage engine cache and memory behavior.
 
     This is separate from ``_DiffusionConfigProjection.cache_config``, which configures
@@ -383,17 +436,26 @@ class OmniStageCacheConfig:
     mamba_ssm_cache_dtype: str | None = None
 
 
-@config
-class OmniStageSchedulerConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig):
     """Per-stage request scheduling behavior."""
 
+    # Upstream receives max_model_len only while materializing SchedulerConfig.
+    # Omni retains it as unresolved stage input until the owning engine process.
     max_num_seqs: int | None = Field(default=None, ge=1)
     max_num_batched_tokens: int | None = Field(default=None, ge=1)
     max_model_len: int | None = Field(default=None, ge=-1)
+    is_encoder_decoder: InitVar[bool] = False
     enable_chunked_prefill: bool | None = None
     async_scheduling: bool | None = None
 
-    def __post_init__(self) -> None:
+    def __post_init__(self, is_encoder_decoder: bool = False) -> None:
+        # Upstream initializes these derived fields in its terminal post-init.
+        # Keep them serializable here without running model-dependent checks.
+        self.max_num_encoder_input_tokens = self.max_num_batched_tokens
+        self.encoder_cache_size = self.max_num_batched_tokens
+
         if (
             self.max_num_batched_tokens is not None
             and self.max_num_seqs is not None
@@ -432,24 +494,73 @@ class OmniStageRuntimeConfig:
     num_gpus: int = Field(default=1, ge=1)
     log_level: str = "info"
     log_stats: bool = False
-    profiler_config: dict[str, Any] | None = None
 
 
-@config
-class OmniStageParallelConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig):
     """Common per-stage distributed parallelism behavior."""
 
-    pipeline_parallel_size: int = Field(default=1, ge=1)
-    data_parallel_size: int = Field(default=1, ge=1)
-    tensor_parallel_size: int = Field(default=1, ge=1)
-    enable_expert_parallel: bool = False
-    world_size: int = Field(default=1, ge=1, init=False)
+    # EngineArgs intentionally leaves these unresolved. The upstream terminal
+    # ParallelConfig defaults (rank 0, local size 1, port 29550, worker
+    # ``"auto"``) must not turn into explicit engine inputs in the head
+    # process merely because this transport class inherits ParallelConfig.
+    data_parallel_size_local: int | None = Field(default=None, ge=0)
+    data_parallel_rank: int | None = Field(default=None, ge=0)
+    data_parallel_rpc_port: int | None = None
+    worker_cls: str | None = None
+
+    @model_validator(mode="after")
+    def _validate_parallel_config(self) -> Self:
+        """Run upstream validation without resolving deferred engine inputs."""
+        deferred_values = (
+            self.data_parallel_size_local,
+            self.data_parallel_rank,
+            self.data_parallel_rpc_port,
+            self.worker_cls,
+            self.all2all_backend,
+            self.disable_custom_all_reduce,
+        )
+        if self.data_parallel_size_local is not None and self.data_parallel_size_local > self.data_parallel_size:
+            raise ValueError(
+                f"data_parallel_size_local ({self.data_parallel_size_local}) "
+                f"must be <= data_parallel_size ({self.data_parallel_size})"
+            )
+        if self.data_parallel_rank is not None and not 0 <= self.data_parallel_rank < self.data_parallel_size:
+            raise ValueError(
+                f"data_parallel_rank ({self.data_parallel_rank}) must be in the range [0, {self.data_parallel_size})"
+            )
+        self.data_parallel_size_local = self.data_parallel_size
+        self.data_parallel_rank = 0
+        self.data_parallel_rpc_port = VllmParallelConfig.data_parallel_rpc_port
+        self.worker_cls = VllmParallelConfig.worker_cls
+        try:
+            VllmParallelConfig._validate_parallel_config(self)
+        finally:
+            (
+                self.data_parallel_size_local,
+                self.data_parallel_rank,
+                self.data_parallel_rpc_port,
+                self.worker_cls,
+                self.all2all_backend,
+                self.disable_custom_all_reduce,
+            ) = deferred_values
+        return self
 
     def __post_init__(self) -> None:
+        # Keep config construction transport-safe. Upstream runtime backend,
+        # rank, port, and platform resolution happens in the owning process.
+        self.data_parallel_index = self.data_parallel_rank
         self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size
 
+    @property
+    def world_size_across_dp(self) -> int:
+        # Omni's public world_size has historically included DP.
+        return self.world_size
 
-@config
+
+@_enforce_keyword_only_init
+@config(kw_only=True)
 class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     """Diffusion-stage distributed parallelism behavior."""
 
@@ -468,6 +579,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     hsdp_replicate_size: int = Field(default=1, ge=1)
 
     def __post_init__(self) -> None:
+        self.data_parallel_index = self.data_parallel_rank
         self.sequence_parallel_size = (
             self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
         )
@@ -576,6 +688,10 @@ 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"
@@ -638,6 +754,20 @@ 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):
@@ -814,6 +944,10 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
         "omni_kv_config",
         "cfg_kv_collect_func",
         "max_num_seqs",
+        "kv_cache_memory_bytes",
+        "gpu_memory_utilization",
+        "max_num_batched_tokens",
+        "max_model_len",
         "enable_sleep_mode",
         "enforce_eager",
         "enable_multithread_weight_load",
@@ -839,6 +973,58 @@ _DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMP
     "stage_id",
 }
 
+
+def _upstream_engine_field_map(
+    config_cls: type[Any],
+    *,
+    aliases: Mapping[str, str] = {},
+    exclude: frozenset[str] = frozenset(),
+) -> dict[str, str]:
+    """Map reusable upstream config fields to their EngineArgs inputs."""
+    engine_fields = frozenset(config_field.name for config_field in fields(VllmEngineArgs))
+    return {
+        config_field.name: engine_name
+        for config_field in fields(config_cls)
+        if config_field.init
+        and config_field.name not in exclude
+        and (engine_name := aliases.get(config_field.name, config_field.name)) in engine_fields
+    }
+
+
+_LOAD_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(VllmLoadConfig)
+_CACHE_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+    VllmCacheConfig,
+    aliases={"cache_dtype": "kv_cache_dtype"},
+)
+_SCHEDULER_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+    VllmSchedulerConfig,
+    aliases={"policy": "scheduling_policy"},
+    # ``scheduler_cls`` is selected by the immutable stage topology, while
+    # ``disable_hybrid_kv_cache_manager`` belongs to the cache concern in the
+    # Omni schema.  Keep both out of the dynamic SchedulerConfig projection so
+    # one input cannot acquire two owners.
+    exclude=frozenset(
+        {
+            "scheduler_cls",
+            "disable_hybrid_kv_cache_manager",
+        }
+    ),
+)
+_PARALLEL_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+    VllmParallelConfig,
+    aliases={"data_parallel_master_ip": "data_parallel_address"},
+    # Runtime owns worker selection. The private API-process fields are
+    # terminal vLLM internals rather than per-stage user inputs.
+    exclude=frozenset(
+        {
+            "distributed_executor_backend",
+            "worker_cls",
+            "_api_process_count",
+            "_api_process_rank",
+        }
+    ),
+)
+
 _QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__)
 _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__)
 _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__)
@@ -846,7 +1032,17 @@ _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__)
 _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__)
 _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__)
 _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__)
-_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__)
+_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"})
+_LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP.values())
+_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(
+    f.name for f in fields(OmniStageDiffusionParallelConfig)
+) & frozenset(_ParallelConfigEngineOverrides.__annotations__)
+_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP = {name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS}
+_LLM_PARALLEL_CONFIG_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP)
+_PARALLEL_CONFIG_ENGINE_FIELDS = _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
 _PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
 _COMMON_STAGE_ENGINE_FIELDS = (
     _QUANTIZATION_ENGINE_FIELDS
@@ -856,12 +1052,16 @@ _COMMON_STAGE_ENGINE_FIELDS = (
     | _SCHEDULER_ENGINE_FIELDS
     | _CONNECTOR_ENGINE_FIELDS
     | _RUNTIME_ENGINE_FIELDS
+    | _DIRECT_VLLM_CONFIG_ENGINE_FIELDS
+)
+_LLM_STAGE_ENGINE_FIELDS = (
+    _COMMON_STAGE_ENGINE_FIELDS
+    | _LLM_LOAD_ENGINE_FIELDS
+    | _LLM_CACHE_ENGINE_FIELDS
+    | _LLM_SCHEDULER_ENGINE_FIELDS
+    | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS
+    | {"parallel_config"}
 )
-_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageParallelConfig)) - {"world_size"}
-_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageDiffusionParallelConfig)) - {
-    "world_size"
-}
-_LLM_STAGE_ENGINE_FIELDS = _COMMON_STAGE_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
 _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
     _COMMON_STAGE_ENGINE_FIELDS
     | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
@@ -882,6 +1082,27 @@ _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
     StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
     StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS,
 }
+_PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = {
+    StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS,
+    StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS,
+    StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP),
+}
+
+_LOAD_STAGE_ENGINE_FIELD_MAP = {
+    **{name: name for name in _LOAD_ENGINE_FIELDS},
+    **_LOAD_CONFIG_ENGINE_FIELD_MAP,
+}
+_CACHE_STAGE_ENGINE_FIELD_MAP = {
+    **{name: name for name in _CACHE_ENGINE_FIELDS},
+    **_CACHE_CONFIG_ENGINE_FIELD_MAP,
+}
+_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {
+    **{name: name for name in _SCHEDULER_ENGINE_FIELDS},
+    **_SCHEDULER_CONFIG_ENGINE_FIELD_MAP,
+}
+_DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP = {name: name for name in _LOAD_ENGINE_FIELDS}
+_DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP = {name: name for name in _CACHE_ENGINE_FIELDS}
+_DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {name: name for name in _SCHEDULER_ENGINE_FIELDS}
 
 
 def _validate_stage_engine_override_ownership(
@@ -892,13 +1113,19 @@ def _validate_stage_engine_override_ownership(
     try:
         owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
         parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
+        parallel_config_fields = _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE[execution_type]
     except KeyError as exc:
         raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc
 
     unowned_fields = set(overrides) - owner_fields
     parallel_config = overrides.get("parallel_config")
     if isinstance(parallel_config, Mapping):
-        unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - parallel_owner_fields)
+        # Nested values traditionally use upstream config names, while flat
+        # EngineArgs uses aliases such as ``data_parallel_address``.  Accept
+        # either spelling at the boundary; the builder canonicalizes aliases
+        # before instantiating the inherited config.
+        accepted_parallel_names = parallel_config_fields | parallel_owner_fields
+        unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - accepted_parallel_names)
     if unowned_fields:
         names = ", ".join(sorted(unowned_fields))
         raise ValueError(
@@ -969,17 +1196,28 @@ def _stage_engine_values(
         topology.execution_type,
         engine,
     )
+    if topology.execution_type in {
+        StageExecutionType.LLM_AR,
+        StageExecutionType.LLM_GENERATION,
+    }:
+        load_engine_fields = _LLM_LOAD_ENGINE_FIELDS
+        cache_engine_fields = _LLM_CACHE_ENGINE_FIELDS
+        scheduler_engine_fields = _LLM_SCHEDULER_ENGINE_FIELDS
+    else:
+        load_engine_fields = _LOAD_ENGINE_FIELDS
+        cache_engine_fields = _CACHE_ENGINE_FIELDS
+        scheduler_engine_fields = _SCHEDULER_ENGINE_FIELDS
     return _StageEngineValues(
         quantization=cast(
             _QuantizationEngineOverrides,
             _select_engine_overrides(engine, _QUANTIZATION_ENGINE_FIELDS),
         ),
         model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)),
-        load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, _LOAD_ENGINE_FIELDS)),
-        cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, _CACHE_ENGINE_FIELDS)),
+        load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)),
+        cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)),
         scheduler=cast(
             _SchedulerEngineOverrides,
-            _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS),
+            _select_engine_overrides(engine, scheduler_engine_fields),
         ),
         connector=cast(
             _ConnectorEngineOverrides,
@@ -988,6 +1226,8 @@ def _stage_engine_values(
         runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)),
         parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)),
         diffusion=_DiffusionEngineOverrides.from_engine(engine),
+        compilation_config=_copy_value(engine.get("compilation_config")),
+        profiler_config=_copy_value(engine.get("profiler_config")),
     )
 
 
@@ -1042,6 +1282,8 @@ class BaseVllmOmniStageConfig:
     connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig)
     runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig)
     parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig)
+    compilation_config: VllmCompilationConfig | None = None
+    profiler_config: VllmProfilerConfig | None = None
     quantization_config: _QuantizationConfigType = None
 
     @property
@@ -1193,6 +1435,8 @@ def _build_common_stage_config_kwargs(
                 parallel_config,
             ),
             "parallel_config": parallel_config,
+            "compilation_config": _copy_value(engine.compilation_config),
+            "profiler_config": _copy_value(engine.profiler_config),
             "quantization_config": _copy_value(quantization_config),
         },
         input_proc,
@@ -1387,12 +1631,51 @@ def _build_load_config(
     return OmniStageLoadConfig(**kwargs)
 
 
+def _config_kwargs_from_engine_args(
+    engine: Mapping[str, Any],
+    field_map: Mapping[str, str],
+) -> dict[str, Any]:
+    """Convert EngineArgs names back to their upstream config field names."""
+    return {
+        config_name: _copy_value(engine[engine_name])
+        for config_name, engine_name in field_map.items()
+        if engine_name in engine and engine[engine_name] is not None
+    }
+
+
+def _normalize_config_mapping(
+    values: Mapping[str, Any],
+    field_map: Mapping[str, str],
+) -> dict[str, Any]:
+    """Canonicalize config and EngineArgs spellings to config field names."""
+    engine_to_config = {engine_name: config_name for config_name, engine_name in field_map.items()}
+    normalized: dict[str, Any] = {}
+    source_names: dict[str, str] = {}
+    for name, value in values.items():
+        if value is None:
+            continue
+        config_name = name if name in field_map else engine_to_config.get(name, name)
+        if config_name in normalized:
+            raise ValueError(
+                f"Config mapping specifies both {source_names[config_name]!r} and {name!r} "
+                f"for upstream field {config_name!r}"
+            )
+        normalized[config_name] = _copy_value(value)
+        source_names[config_name] = name
+    return normalized
+
+
+def _omni_config_kwargs(engine: Mapping[str, Any], fields_: frozenset[str]) -> dict[str, Any]:
+    return {name: _copy_value(engine[name]) for name in fields_ if name in engine and engine[name] is not None}
+
+
 def _build_cache_config(
     deploy: DeployConfig,
     engine: _CacheEngineOverrides,
     execution_type: StageExecutionType,
 ) -> OmniStageCacheConfig:
-    kwargs = _config_kwargs(engine)
+    kwargs = _config_kwargs_from_engine_args(engine, _CACHE_CONFIG_ENGINE_FIELD_MAP)
+    kwargs.update(_omni_config_kwargs(engine, _CACHE_ENGINE_FIELDS))
     if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None:
         kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching)
     if "disable_hybrid_kv_cache_manager" not in kwargs and execution_type == StageExecutionType.LLM_GENERATION:
@@ -1406,7 +1689,8 @@ def _build_scheduler_config(
     engine: _SchedulerEngineOverrides,
     execution_type: StageExecutionType,
 ) -> OmniStageSchedulerConfig:
-    kwargs = _config_kwargs(engine)
+    kwargs = _config_kwargs_from_engine_args(engine, _SCHEDULER_CONFIG_ENGINE_FIELD_MAP)
+    kwargs.update(_omni_config_kwargs(engine, _SCHEDULER_ENGINE_FIELDS))
     if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None:
         kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill)
     if "async_scheduling" not in kwargs and execution_type == StageExecutionType.LLM_AR:
@@ -1454,16 +1738,21 @@ def _build_parallel_config(
     config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig,
 ) -> OmniStageParallelConfig:
     parallel_config = _mapping_or_empty(engine.get("parallel_config"))
-    config_fields = {
-        config_field.name
-        for config_field in fields(config_cls)
-        if getattr(config_field.default, "init", None) is not False
-    }
-    kwargs = {
-        name: _copy_value(value) for name in _PARALLEL_CONFIG_ENGINE_FIELDS if (value := engine.get(name)) is not None
-    }
+    config_fields = frozenset(signature(config_cls).parameters)
+    is_diffusion = issubclass(config_cls, OmniStageDiffusionParallelConfig)
+    field_map = _DIFFUSION_PARALLEL_CONFIG_FIELD_MAP if is_diffusion else _PARALLEL_CONFIG_ENGINE_FIELD_MAP
+    owned_parallel_fields = frozenset(field_map)
+    kwargs = _config_kwargs(engine) if is_diffusion else _config_kwargs_from_engine_args(engine, field_map)
+    kwargs = {name: value for name, value in kwargs.items() if name in owned_parallel_fields}
     kwargs = {name: value for name, value in kwargs.items() if name in config_fields}
-    kwargs.update({name: value for name, value in _config_kwargs(parallel_config).items() if name in config_fields})
+    nested_kwargs = _normalize_config_mapping(parallel_config, field_map)
+    kwargs.update(
+        {
+            name: value
+            for name, value in nested_kwargs.items()
+            if name in owned_parallel_fields and name in config_fields
+        }
+    )
     if "pipeline_parallel_size" not in kwargs and deploy.pipeline_parallel_size is not None:
         kwargs["pipeline_parallel_size"] = _copy_value(deploy.pipeline_parallel_size)
     if "data_parallel_size" not in kwargs and deploy.data_parallel_size is not None:

vllm_omni/config/pipeline_registry.py

@@ -33,7 +33,7 @@ from vllm.logger import init_logger
 from vllm_omni.config.stage_config import (
     PipelineConfig,
 )
-from vllm_omni.diffusion.models.pi0.pipeline_pi0 import PI0_PIPELINE
+from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE
 from vllm_omni.model_executor.models.audex.pipeline import (
     AUDEX_S2S_PIPELINE,
     AUDEX_THINKER_ONLY_PIPELINE,
@@ -85,6 +85,7 @@ from vllm_omni.model_executor.models.ming_tts.pipeline import (
     MING_TTS_PIPELINE,
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
+from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE
 from vllm_omni.model_executor.models.moss_tts.pipeline import (
     MOSS_TTS_LOCAL_PIPELINE,
     MOSS_TTS_PIPELINE,
@@ -177,6 +178,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
     "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,
     "minicpmo_4_5": MINICPMO_4_5_PIPELINE,
+    "minimax_music3": MINIMAX_MUSIC3_PIPELINE,
     "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE,
     "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,
     "dynin_omni": DYNIN_OMNI_PIPELINE,

vllm_omni/config/stage_config.py

@@ -382,6 +382,10 @@ 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/data_entry_keys.py

@@ -91,6 +91,11 @@ class OmniPayloadMeta(TypedDict, total=False):
     ref_context_included: bool
     talker_prefill_offset: int
     omni_final_stage_id: int
+    # Per-request audio seed. Stages that draw their own noise (flow-matching
+    # / diffusion decoders) need the producing stage's seed to stay
+    # reproducible, and the producing stage's SamplingParams do not travel
+    # with the payload.
+    audio_seed: int
 
 
 class OmniPayload(TypedDict, total=False):
@@ -189,6 +194,7 @@ class MetaStruct(_StructBase):
     codec_left_context_frames: int | None = None
     code_flat_numel: int | None = None
     omni_final_stage_id: int | None = None
+    audio_seed: int | None = None
 
 
 class OmniPayloadStruct(_StructBase):

vllm_omni/deploy/minimax_music3.yaml

@@ -0,0 +1,130 @@
+# MiniMax-Music3 deploy: Stage 0 (AR talker) -> Stage 1 (acoustic decoder),
+# both on one 140 GB card, streaming 200-frame windows over shared memory.
+#
+# Auto-discovered default for model_type=minimax_music3.
+#
+# Classifier-free guidance is mandatory for this checkpoint: every request
+# decodes as a cond/uncond pair in stage 0, so stage-0 KV, max_num_seqs and
+# prefill budget are all sized for TWO rows per user request. The stage-0
+# extra_args block below is what turns on has_sampling_extra_args (so per-row
+# sampling extra args reach the model's forward) and is also how the CFG
+# scheduler patches learn this stage is paired.
+pipeline: minimax_music3
+async_chunk: true
+
+# The repo-root config.json reports model_type=minimax_music3, which upstream
+# transformers does not ship a config class for.
+trust_remote_code: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      # Stream stereo PCM back to the client as stage 1 finishes each window.
+      codec_streaming: true
+      # Window geometry is fixed by the checkpoint (200-frame windows on a
+      # 100-frame hop, models/minimax_music3/constants.py), so there are no
+      # codec_chunk_frames knobs here: changing the window changes the audio.
+      shm_threshold_bytes: 65536
+      connector_get_sleep_s: 0.01
+      # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer
+      # for a chunk than a per-step codec stage would. First chunk additionally
+      # covers stage-0 prefill of both CFG rows.
+      connector_get_max_wait_first_chunk: 30000
+      connector_get_max_wait: 10000
+
+stages:
+  - stage_id: 0
+    # 16 rows = 8 concurrent guided requests, since CFG doubles every request.
+    max_num_seqs: 16
+    # ~16 GB of bf16 Qwen3-8B weights plus KV for 16 full-length rows
+    # (36 layers x 8 KV heads x 128 head_dim ~= 144 KB/token, so ~1.4 GB per
+    # 10240-token row) fits comfortably inside 0.6 of a 140 GB card.
+    gpu_memory_utilization: 0.6
+    dtype: bfloat16
+    # The checkpoint's max_position_embeddings is 10240 and the whole song is
+    # decoded in-context. Note this cannot hold a max-length prompt AND a
+    # max-length song at the same time: the prompt cap is 5000 tokens and the
+    # frame budget is 9000, which overruns 10240. That is a property of the
+    # checkpoint, not of this file; long captions shorten the maximum song.
+    max_model_len: 10240
+    # One scheduler step must be able to prefill both rows of a pair
+    # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked.
+    max_num_batched_tokens: 10240
+    # Two rows at 8B is a launch-bound decode step, so the backbone is
+    # captured. FULL_DECODE_ONLY only ever dispatches a graph for a uniform
+    # decode batch; prefill and mixed steps stay eager, which is what the
+    # talker's frame-feedback fallback is for. The talker's own sampler and
+    # depth decoder run after the graph, not inside it.
+    #
+    # What makes the forward capturable: the frame feedback writes a
+    # fixed-size slab in prepare_runner_inputs (outside the captured region)
+    # and the forward selects over it with torch.where, instead of an
+    # index_copy against an index tensor whose address and length changed
+    # every step. The row-to-request binding the sampler reads is likewise
+    # rebuilt in prepare_runner_inputs, because a replayed step never runs
+    # the forward's Python body.
+    enforce_eager: false
+    compilation_config:
+      # No compilation. The talker walks the backbone's decoder layers itself,
+      # so the compiled Qwen3Model.forward is never entered and inductor has
+      # nothing to do; what the mode still decides is which kernels vLLM
+      # defaults to. Left at its non-eager default it picks the elementwise
+      # PyTorch norms over the fused vllm_c ones, which is both slower and
+      # numerically different: measured here, that alone moved the sampled
+      # codes off the eager trajectory by frame 4. Graph capture is
+      # independent of this and still happens.
+      mode: 0
+      # 16 rows is the stage ceiling, i.e. 8 guided requests.
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16]
+      cudagraph_mode: FULL_DECODE_ONLY
+      cudagraph_num_of_warmups: 1
+      # Belt and braces on the same point: the non-eager default is "none",
+      # which hands the norm and activation ops to elementwise PyTorch.
+      custom_ops: ["all"]
+    # A CFG pair must be sampled inside one engine step for the logits blend
+    # and the cond->uncond token copy to line up.
+    async_scheduling: false
+    # Off during bring-up: the uncond row is a long run of identical
+    # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than
+    # its cond partner, which is exactly the num_computed_tokens skew the CFG
+    # pair patches then have to undo.
+    enable_prefix_caching: false
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 1.0
+      # The frame budget is enforced by the model. This ceiling only has to
+      # leave room for the extra step that delivers the final window.
+      max_tokens: 9002
+      extra_args:
+        cfg_role: cond
+        cfg_scale: 1.5
+
+  - stage_id: 1
+    # One row per user request: the uncond companion is stage-0-final and
+    # never reaches the acoustic stage.
+    max_num_seqs: 8
+    # No KV cache on this stage; the budget only covers the DiT and DAV
+    # weights plus per-window latents and activations.
+    gpu_memory_utilization: 0.2
+    # The flow-matching solver degrades measurably in bfloat16
+    # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py).
+    dtype: float32
+    enforce_eager: true
+    # Async scheduling builds per-request payloads against the NEXT step's
+    # input batch, which breaks when chunk arrival staggers requests.
+    async_scheduling: false
+    enable_prefix_caching: false
+    # One placeholder prompt token per arriving window, so a six-minute song
+    # (90 windows) needs a tiny context.
+    max_model_len: 4096
+    max_num_batched_tokens: 4096
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 4096
+      seed: 0

vllm_omni/deploy/minimax_music3_2gpu.yaml

@@ -0,0 +1,97 @@
+# MiniMax-Music3 two-GPU layout: AR talker on device 0, acoustic decoder on
+# device 1.
+#
+# Splitting the stages removes the memory-budget compromise of the single-card
+# file: the talker's CFG-doubled KV no longer competes with the DiT and DAV
+# weights, and the flow-matching solver no longer competes with AR decode for
+# SMs. Everything else (window geometry, CFG contract, sampling defaults)
+# matches minimax_music3.yaml, so the two files stay behaviorally identical.
+pipeline: minimax_music3
+async_chunk: true
+
+# The repo-root config.json reports model_type=minimax_music3, which upstream
+# transformers does not ship a config class for.
+trust_remote_code: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      # Stream stereo PCM back to the client as stage 1 finishes each window.
+      codec_streaming: true
+      # Window geometry is fixed by the checkpoint (200-frame windows on a
+      # 100-frame hop, models/minimax_music3/constants.py), so there are no
+      # codec_chunk_frames knobs here: changing the window changes the audio.
+      shm_threshold_bytes: 65536
+      connector_get_sleep_s: 0.01
+      # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer
+      # for a chunk than a per-step codec stage would. First chunk additionally
+      # covers stage-0 prefill of both CFG rows.
+      connector_get_max_wait_first_chunk: 30000
+      connector_get_max_wait: 10000
+
+stages:
+  - stage_id: 0
+    # 32 rows = 16 concurrent guided requests, since CFG doubles every request.
+    # Twice the single-card figure: this stage now owns its whole device.
+    max_num_seqs: 32
+    gpu_memory_utilization: 0.85
+    dtype: bfloat16
+    # The checkpoint's max_position_embeddings is 10240 and the whole song is
+    # decoded in-context. Note this cannot hold a max-length prompt AND a
+    # max-length song at the same time: the prompt cap is 5000 tokens and the
+    # frame budget is 9000, which overruns 10240. That is a property of the
+    # checkpoint, not of this file; long captions shorten the maximum song.
+    max_model_len: 10240
+    # One scheduler step must be able to prefill both rows of a pair
+    # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked.
+    max_num_batched_tokens: 10240
+    # The talker owns its sampler and drives a depth decoder per AR step;
+    # leave graph capture off until that path is verified graph-safe.
+    enforce_eager: true
+    # A CFG pair must be sampled inside one engine step for the logits blend
+    # and the cond->uncond token copy to line up.
+    async_scheduling: false
+    # Off during bring-up: the uncond row is a long run of identical
+    # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than
+    # its cond partner, which is exactly the num_computed_tokens skew the CFG
+    # pair patches then have to undo.
+    enable_prefix_caching: false
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 1.0
+      # The frame budget is enforced by the model. This ceiling only has to
+      # leave room for the extra step that delivers the final window.
+      max_tokens: 9002
+      extra_args:
+        cfg_role: cond
+        cfg_scale: 1.5
+
+  - stage_id: 1
+    # One row per user request: the uncond companion is stage-0-final and
+    # never reaches the acoustic stage.
+    max_num_seqs: 16
+    # No KV cache on this stage, but it now has a device to itself, so the
+    # solver can hold more windows in flight than the shared-card profile.
+    gpu_memory_utilization: 0.6
+    # The flow-matching solver degrades measurably in bfloat16
+    # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py).
+    dtype: float32
+    enforce_eager: true
+    # Async scheduling builds per-request payloads against the NEXT step's
+    # input batch, which breaks when chunk arrival staggers requests.
+    async_scheduling: false
+    enable_prefix_caching: false
+    # One placeholder prompt token per arriving window, so a six-minute song
+    # (90 windows) needs a tiny context.
+    max_model_len: 4096
+    max_num_batched_tokens: 4096
+    devices: "1"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 4096
+      seed: 0

vllm_omni/diffusion/attention/backends/abstract.py

@@ -81,6 +81,19 @@ class AttentionBackend(ABC):
         supported_head_sizes = cls.get_supported_head_sizes()
         return (not supported_head_sizes) or head_size in supported_head_sizes
 
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        """Whether this backend reads K/V pages by the runtime block stride.
+
+        Returning ``True`` means the physical cache layout has ``num_blocks``
+        as its outer stride, so native vLLM may safely use page-size padding
+        when it unifies cache layouts across layers. Dense diffusion backends
+        conservatively keep the default ``False``; a paged backend should
+        override this only when its kernel actually follows that layout.
+        """
+
+        return False
+
 
 @dataclass(frozen=True, slots=True)
 class QueryRange:

vllm_omni/diffusion/attention/layer.py

@@ -11,8 +11,10 @@ from dataclasses import replace
 
 import torch
 import torch.nn as nn
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.models.utils import extract_layer_index
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
@@ -61,12 +63,23 @@ class Attention(nn.Module):
         # perf for this layer (e.g. Wan2.2 cross-attn has short sequences and
         # block-FP8 quant offers no win). Default False = follow global config.
         disable_kv_quant: bool = False,
+        # Opt-in marker for Scheduler-managed paged KV. Unmarked diffusion
+        # attention remains dense and contributes no native KVCacheSpec.
+        paged_kv_cache_role: str | None = None,
+        paged_kv_cache_dtype: torch.dtype | None = None,
     ):
         super().__init__()
 
         self.role = role
         self.role_category = role_category
         self.qkv_layout = qkv_layout
+        if paged_kv_cache_role == "":
+            raise ValueError("paged_kv_cache_role must be non-empty when provided")
+        self.paged_kv_cache_role = paged_kv_cache_role
+        self.paged_kv_cache_dtype = paged_kv_cache_dtype
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.head_size = head_size
 
         # Resolve backend via role-aware config.
         # The global diffusion config is set during model init via
@@ -171,6 +184,25 @@ class Attention(nn.Module):
         self._disable_kv_quant: bool = disable_kv_quant
         self._init_kv_cache_quantization(config)
 
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
+        """Return native rank-local geometry for an opted-in paged cache."""
+
+        if self.paged_kv_cache_role is None:
+            return None
+        dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
+        # Keep backend layout discovery under the same config context used by
+        # upstream vLLM's attention-spec collector.
+        with set_current_vllm_config(vllm_config):
+            indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride()
+        return FullAttentionSpec(
+            block_size=vllm_config.cache_config.block_size,
+            num_kv_heads=self.num_kv_heads,
+            head_size=self.head_size,
+            dtype=dtype,
+            indexes_kv_by_block_stride=indexes_kv_by_block_stride,
+            non_causal=not self.causal,
+        )
+
     def _get_active_parallel_strategy(self):
         """Get the parallel strategy based on current SP active state.
 

vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py

@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused MiniMax H3 modulation with FP32 accumulation."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _indexed_scale_shift_kernel(
+    output_ptr,
+    x_ptr,
+    shift_ptr,
+    scale_ptr,
+    indices_ptr,
+    hidden_size,
+    stride_x_row,
+    stride_shift_row,
+    stride_scale_row,
+    stride_indices,
+    block_n: tl.constexpr,
+):
+    row = tl.program_id(0)
+    columns = tl.arange(0, block_n)
+    mask = columns < hidden_size
+    index = tl.load(indices_ptr + row * stride_indices)
+
+    x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
+    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+
+    tl.store(
+        output_ptr + row * stride_x_row + columns,
+        x * (1.0 + scale) + shift,
+        mask=mask,
+    )
+
+
+@triton.jit
+def _indexed_gate_kernel(
+    output_ptr,
+    x_ptr,
+    gate_ptr,
+    other_ptr,
+    indices_ptr,
+    hidden_size,
+    stride_output_row,
+    stride_x_row,
+    stride_gate_row,
+    stride_other_row,
+    stride_indices,
+    block_n: tl.constexpr,
+):
+    row = tl.program_id(0)
+    columns = tl.arange(0, block_n)
+    mask = columns < hidden_size
+    index = tl.load(indices_ptr + row * stride_indices)
+
+    x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
+    gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32)
+    other = tl.load(other_ptr + row * stride_other_row + columns, mask=mask, other=0.0).to(tl.float32)
+
+    tl.store(
+        output_ptr + row * stride_output_row + columns,
+        x + gate * other,
+        mask=mask,
+    )
+
+
+@triton.jit
+def _rms_norm_indexed_scale_shift_kernel(
+    output_ptr,
+    x_ptr,
+    weight_ptr,
+    shift_ptr,
+    scale_ptr,
+    indices_ptr,
+    hidden_size: tl.constexpr,
+    eps: tl.constexpr,
+    stride_x_row,
+    stride_shift_row,
+    stride_scale_row,
+    stride_indices,
+    block_n: tl.constexpr,
+):
+    row = tl.program_id(0)
+    columns = tl.arange(0, block_n)
+    mask = columns < hidden_size
+    index = tl.load(indices_ptr + row * stride_indices)
+
+    x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
+    weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
+    variance = tl.sum(x * x, axis=0) / hidden_size
+    normalized = x * tl.rsqrt(variance + eps) * weight
+
+    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+    tl.store(
+        output_ptr + row * hidden_size + columns,
+        normalized * (1.0 + scale) + shift,
+        mask=mask,
+    )
+
+
+@triton.jit
+def _indexed_gate_rms_norm_scale_shift_kernel(
+    residual_out_ptr,
+    modulated_out_ptr,
+    residual_ptr,
+    gate_ptr,
+    branch_ptr,
+    weight_ptr,
+    shift_ptr,
+    scale_ptr,
+    indices_ptr,
+    hidden_size: tl.constexpr,
+    eps: tl.constexpr,
+    stride_residual_row,
+    stride_gate_row,
+    stride_branch_row,
+    stride_shift_row,
+    stride_scale_row,
+    stride_indices,
+    block_n: tl.constexpr,
+):
+    row = tl.program_id(0)
+    columns = tl.arange(0, block_n)
+    mask = columns < hidden_size
+    index = tl.load(indices_ptr + row * stride_indices)
+
+    residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0).to(tl.float32)
+    gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32)
+    branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0).to(tl.float32)
+    updated = residual + gate * branch
+    tl.store(residual_out_ptr + row * hidden_size + columns, updated, mask=mask)
+
+    weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
+    variance = tl.sum(updated * updated, axis=0) / hidden_size
+    normalized = updated * tl.rsqrt(variance + eps) * weight
+    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+    tl.store(
+        modulated_out_ptr + row * hidden_size + columns,
+        normalized * (1.0 + scale) + shift,
+        mask=mask,
+    )
+
+
+def indexed_scale_shift_(
+    x: torch.Tensor,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    indices: torch.Tensor,
+) -> torch.Tensor:
+    """Apply indexed scale/shift in-place to a disposable contiguous input."""
+    if x.is_cpu:
+        x.copy_((x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(x.dtype))
+        return x
+    rows, hidden_size = x.shape
+    if rows == 0:
+        return x
+    _indexed_scale_shift_kernel[(rows,)](
+        x,
+        x,
+        shift,
+        scale,
+        indices,
+        hidden_size,
+        x.stride(0),
+        shift.stride(0),
+        scale.stride(0),
+        indices.stride(0),
+        block_n=triton.next_power_of_2(hidden_size),
+        num_warps=8,
+    )
+    return x
+
+
+def indexed_gate(
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    other: torch.Tensor,
+    indices: torch.Tensor,
+) -> torch.Tensor:
+    """Return ``x + gate[indices] * other`` without indexed temporaries."""
+    if x.is_cpu:
+        return (x + gate.index_select(0, indices) * other).to(x.dtype)
+    output = torch.empty_like(x)
+    rows, hidden_size = x.shape
+    if rows == 0:
+        return output
+    _indexed_gate_kernel[(rows,)](
+        output,
+        x,
+        gate,
+        other,
+        indices,
+        hidden_size,
+        output.stride(0),
+        x.stride(0),
+        gate.stride(0),
+        other.stride(0),
+        indices.stride(0),
+        block_n=triton.next_power_of_2(hidden_size),
+        num_warps=8,
+    )
+    return output
+
+
+def rms_norm_indexed_scale_shift(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    indices: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Fuse H3 RMSNorm with its indexed AdaLN affine transform."""
+    if x.is_cpu:
+        input_dtype = x.dtype
+        normalized = x.float()
+        variance = normalized.pow(2).mean(-1, keepdim=True)
+        normalized = normalized * torch.rsqrt(variance + eps)
+        normalized = (weight.float() * normalized).to(input_dtype)
+        return (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(input_dtype)
+    output = torch.empty_like(x)
+    rows, hidden_size = x.shape
+    if rows:
+        _rms_norm_indexed_scale_shift_kernel[(rows,)](
+            output,
+            x,
+            weight,
+            shift,
+            scale,
+            indices,
+            hidden_size,
+            eps,
+            x.stride(0),
+            shift.stride(0),
+            scale.stride(0),
+            indices.stride(0),
+            block_n=triton.next_power_of_2(hidden_size),
+            num_warps=8,
+        )
+    return output
+
+
+def indexed_gate_rms_norm_scale_shift(
+    residual: torch.Tensor,
+    gate: torch.Tensor,
+    branch: torch.Tensor,
+    weight: torch.Tensor,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    indices: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse gated residual, the following RMSNorm, and AdaLN affine."""
+    if residual.is_cpu:
+        input_dtype = residual.dtype
+        residual_out = (residual + gate.index_select(0, indices) * branch).to(input_dtype)
+        normalized = residual_out.float()
+        variance = normalized.pow(2).mean(-1, keepdim=True)
+        normalized = normalized * torch.rsqrt(variance + eps)
+        normalized = (weight.float() * normalized).to(input_dtype)
+        modulated_out = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(
+            input_dtype
+        )
+        return residual_out, modulated_out
+    residual_out = torch.empty_like(residual)
+    modulated_out = torch.empty_like(residual)
+    rows, hidden_size = residual.shape
+    if rows:
+        _indexed_gate_rms_norm_scale_shift_kernel[(rows,)](
+            residual_out,
+            modulated_out,
+            residual,
+            gate,
+            branch,
+            weight,
+            shift,
+            scale,
+            indices,
+            hidden_size,
+            eps,
+            residual.stride(0),
+            gate.stride(0),
+            branch.stride(0),
+            shift.stride(0),
+            scale.stride(0),
+            indices.stride(0),
+            block_n=triton.next_power_of_2(hidden_size),
+            num_warps=8,
+        )
+    return residual_out, modulated_out
+
+
+__all__ = [
+    "indexed_gate",
+    "indexed_gate_rms_norm_scale_shift",
+    "indexed_scale_shift_",
+    "rms_norm_indexed_scale_shift",
+]

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1259,12 +1259,11 @@ def extract_minimax_h3_context(
     preprocessing via ``_embed``, cacheable ``blocks`` loop, and
     ``final_layer`` postprocessing with row selection and update masks.
     """
+    from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import indexed_scale_shift_
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
-        _BF16_DTYPE,
         _FORWARD_SUPPORTED_KWARGS,
         MINIMAX_H3_ADALN_MODALITY_NUM,
         _build_rope_table,
-        _modulate_scale_shift,
         _required_kwarg,
     )
 
@@ -1341,13 +1340,12 @@ def extract_minimax_h3_context(
     cu_seqlens = cu_seqlens.to(device)
 
     shift_msa, scale_msa, *_ = module.blocks[0].adaln_proj(t_emb)
-    modulated_hidden = module.blocks[0].norm1(decoder_input)
-    modulated_input = _modulate_scale_shift(
-        modulated_hidden,
+    modulated_input = module.blocks[0].norm1(decoder_input)
+    modulated_input = indexed_scale_shift_(
+        modulated_input,
         shift_msa,
         scale_msa,
         combined_indices,
-        dtype=_BF16_DTYPE,
     )
 
     def run_transformer_blocks() -> tuple[torch.Tensor, ...]:

vllm_omni/diffusion/data.py

@@ -708,6 +708,16 @@ 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
 
@@ -851,6 +861,14 @@ class OmniDiffusionConfig:
     # Maximum number of sequences to generate in a batch
     max_num_seqs: int = 1
 
+    # Native vLLM KV-cache sizing inputs used only by paged_scheduler mode.
+    # These mirror the structured stage cache/scheduler config until all
+    # diffusion runtime construction consumes VllmOmniConfig directly.
+    kv_cache_memory_bytes: int | None = None
+    gpu_memory_utilization: float = 0.9
+    max_num_batched_tokens: int | None = None
+    max_model_len: int | None = None
+
     # Request-mode batch admission: wait briefly for compatible requests to
     # accumulate in the scheduler waiting queue before the first schedule() of
     # a wave.  Improves fused forward batch sizes under bursty HTTP ingress.
@@ -928,6 +946,18 @@ 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', "
@@ -936,6 +966,14 @@ class OmniDiffusionConfig:
         if not isinstance(self.diffusion_compile_dynamic, bool):
             raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+        if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0:
+            raise ValueError("kv_cache_memory_bytes must be non-negative")
+        if not 0.0 < self.gpu_memory_utilization <= 1.0:
+            raise ValueError("gpu_memory_utilization must be in (0, 1]")
+        if self.max_num_batched_tokens is not None and self.max_num_batched_tokens <= 0:
+            raise ValueError("max_num_batched_tokens must be positive")
+        if self.max_model_len is not None and self.max_model_len != -1 and self.max_model_len <= 0:
+            raise ValueError("max_model_len must be positive or -1")
 
         if self.omni_kv_config is None:
             self.omni_kv_config = {}

vllm_omni/diffusion/diffusion_engine.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import copy
 import inspect
 import queue
 import threading
@@ -17,18 +18,23 @@ from typing import TYPE_CHECKING, Any
 import numpy as np
 import PIL.Image
 import torch
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.kv_cache_interface import KVCacheConfig
 
 from vllm_omni.diffusion.data import (
     DiffusionOutput,
     DiffusionRequestAbortedError,
     OmniDiffusionConfig,
 )
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.io_support import (
     get_dummy_run_num_frames,
+    get_dummy_run_num_image_inputs,
     image_color_format,
     supports_audio_output,
     supports_multimodal_input,
@@ -184,12 +190,39 @@ class DiffusionEngine:
                 from the resolved execution mode.
         """
         self.od_config = od_config
+        # Set after the paged-KV profile request has gone through model-owned
+        # preprocessing. Real requests are admitted only within this measured
+        # activation envelope: (max execution sequences, max seq_len,
+        # max target_len).
+        self._diffusion_kv_profile_limits: tuple[int, int, int] | None = None
 
         self._init_process_hooks(od_config)
         self.execution_mode = self._resolve_execution_mode(od_config)
         self._init_executor(od_config)
         try:
-            self._init_scheduler(od_config, scheduler)
+            profile_requests = self._prepare_diffusion_kv_profile_requests()
+            kv_control_plane = initialize_diffusion_kv_control_plane(
+                self.executor,
+                od_config,
+                profile_requests=profile_requests,
+            )
+            if kv_control_plane is None:
+                self._init_scheduler(od_config, scheduler)
+            else:
+                (
+                    kv_cache_config,
+                    scheduler_block_size,
+                    hash_block_size,
+                    kv_vllm_config,
+                ) = kv_control_plane
+                self._init_scheduler(
+                    od_config,
+                    scheduler,
+                    kv_cache_config,
+                    scheduler_block_size=scheduler_block_size,
+                    hash_block_size=hash_block_size,
+                    kv_vllm_config=kv_vllm_config,
+                )
             self._init_runtime_state()
             self._init_execute_fn()
             self._log_execution_mode(od_config)
@@ -242,6 +275,11 @@ class DiffusionEngine:
         self,
         od_config: OmniDiffusionConfig,
         scheduler: BaseScheduler | None = None,
+        kv_cache_config: KVCacheConfig | None = None,
+        *,
+        scheduler_block_size: int | None = None,
+        hash_block_size: int | None = None,
+        kv_vllm_config: VllmConfig | None = None,
     ) -> None:
         if scheduler is not None:
             self.scheduler = scheduler
@@ -249,7 +287,16 @@ class DiffusionEngine:
             self.scheduler = StepScheduler()
         else:
             self.scheduler = RequestScheduler()
-        self.scheduler.initialize(od_config)
+        if kv_cache_config is None:
+            self.scheduler.initialize(od_config)
+        else:
+            self.scheduler.initialize(
+                od_config,
+                kv_cache_config=kv_cache_config,
+                scheduler_block_size=scheduler_block_size,
+                hash_block_size=hash_block_size,
+                kv_vllm_config=kv_vllm_config,
+            )
 
     def _init_runtime_state(self) -> None:
         # DP multi-concurrency: allow batching dp_size requests so each
@@ -703,9 +750,36 @@ class DiffusionEngine:
         """Run model-owned preprocessing once, before entering Engine locks."""
 
         pre_process_func = getattr(self, "pre_process_func", None)
-        if pre_process_func is None:
-            return request
-        return pre_process_func(request)
+        if pre_process_func is not None:
+            request = pre_process_func(request)
+        self._validate_diffusion_kv_profile_limits(request)
+        return request
+
+    def _validate_diffusion_kv_profile_limits(self, request: OmniDiffusionRequest) -> None:
+        """Keep admitted paged-KV requests within the profiled activation shape."""
+
+        profile_limits = getattr(self, "_diffusion_kv_profile_limits", None)
+        if profile_limits is None:
+            return
+        kv_requests = request.diffusion_kv_requests
+        if not kv_requests:
+            return
+
+        request_limits = (
+            len(kv_requests),
+            max(kv_request.seq_len for kv_request in kv_requests),
+            max(kv_request.target_len for kv_request in kv_requests),
+        )
+        if any(actual > profiled for actual, profiled in zip(request_limits, profile_limits, strict=True)):
+            request_sequences, request_seq_len, request_target_len = request_limits
+            profile_sequences, profile_seq_len, profile_target_len = profile_limits
+            raise ValueError(
+                f"Diffusion KV request {request.request_id!r} exceeds the startup memory-profile envelope: "
+                f"sequences={request_sequences} (profiled={profile_sequences}), "
+                f"max_seq_len={request_seq_len} (profiled={profile_seq_len}), "
+                f"max_target_len={request_target_len} (profiled={profile_target_len}). "
+                "Reduce the request shape or extend the model's paged-KV profile recipe."
+            )
 
     def _add_prepared_request(self, request: OmniDiffusionRequest) -> str:
         """Admit a request whose model-owned preprocessing is complete."""
@@ -869,47 +943,111 @@ class DiffusionEngine:
             self.close()
             raise e
 
-    def _dummy_run(self):
-        """A dummy run to warm up the model."""
-        num_inference_steps = 1
-        height = 512
-        width = 512
-        prompt: OmniTextPrompt = {"prompt": "dummy run"}
+    def _make_dummy_request(
+        self,
+        *,
+        height: int,
+        width: int,
+        guidance_scale: float,
+        num_image_inputs: int = 1,
+    ) -> OmniDiffusionRequest | None:
+        """Build a one-step model request for startup profiling or warmup."""
 
+        prompt: OmniTextPrompt = {"prompt": "dummy run"}
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
         if supports_image_input:
-            # Provide a dummy image input if the model supports it
             color_format = image_color_format(self.od_config.model_class_name)
-            dummy_image = PIL.Image.new(color_format, (width, height))
-            prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
+            images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)]
+            prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images
 
         if supports_audio_input:
             audio_sr = 16000
-            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
-            prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
+            prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
 
         num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
         if num_frames <= 0:
-            logger.info("Skipping dummy warmup run (num_frames=0)")
-            return
-        req = OmniDiffusionRequest(
+            return None
+        return OmniDiffusionRequest(
             prompt=prompt,
             request_id=DUMMY_DIFFUSION_REQUEST_ID,
             sampling_params=OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
-                num_inference_steps=num_inference_steps,
+                num_inference_steps=1,
                 num_frames=num_frames,
-                # Keep warmup path minimal and robust across text encoders.
-                # Some models may fail when warmup implicitly triggers
-                # classifier-free guidance with an empty negative prompt.
-                guidance_scale=0.0,
+                guidance_scale=guidance_scale,
                 num_outputs_per_prompt=1,
-                # Disable CFG for warmup to avoid triggering CFG parallel
-                # validation when cfg_parallel_size > 1.
                 extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
             ),
         )
+
+    def _prepare_diffusion_kv_profile_requests(self) -> list[OmniDiffusionRequest] | None:
+        """Prepare the per-rank request batch used to profile paged-KV headroom.
+
+        The profile executes directly on each Worker before the Scheduler and
+        its KV manager exist. It uses the maximum number of requests that one
+        rank can execute together. DLO+DP is the exception because each rank
+        executes one request from the collective wave.
+
+        Hunyuan is currently the only model integrated with
+        ``paged_scheduler``; 1024x1024, enabled CFG, and the maximum advertised
+        reference-image count exercise its first-step activation peak.
+        Admission compares each preprocessed request's CFG count and tokenized
+        sequence/target shape with the resulting per-request profile envelope.
+        Future paged model integrations must extend this recipe for their
+        serving limits (for example, video frame count) rather than reusing it
+        silently.
+        """
+
+        if (
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is not DiffusionKVCacheMode.PAGED_SCHEDULER
+        ):
+            return None
+        request = self._make_dummy_request(
+            height=1024,
+            width=1024,
+            guidance_scale=5.0,
+            num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name),
+        )
+        if request is None:
+            raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request")
+        request = self._prepare_request_for_admission(request)
+        kv_requests = request.diffusion_kv_requests
+        if not kv_requests:
+            raise RuntimeError("paged_scheduler profile preprocessing must produce Diffusion KV requests")
+        self._diffusion_kv_profile_limits = (
+            len(kv_requests),
+            max(kv_request.seq_len for kv_request in kv_requests),
+            max(kv_request.target_len for kv_request in kv_requests),
+        )
+        dlo_dp_request_mode = self.execution_mode is DiffusionExecutionMode.REQUEST_BATCH and _uses_dlo_dp_concurrency(
+            self.od_config
+        )
+        profile_batch_size = 1 if dlo_dp_request_mode else _max_num_seqs(self.od_config)
+        profile_requests: list[OmniDiffusionRequest] = []
+        for index in range(profile_batch_size):
+            profile_request = copy.copy(request)
+            profile_request.request_id = f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-{index}"
+            profile_request.sampling_params = copy.deepcopy(request.sampling_params)
+            # Preprocessing owns request geometry and prepared model inputs,
+            # but native KV requests remain Scheduler-only mutable state. The
+            # profile bypasses Scheduler admission and must not send them to
+            # the Worker.
+            profile_request.diffusion_kv_requests = None
+            profile_requests.append(profile_request)
+        return profile_requests
+
+    def _dummy_run(self):
+        """A dummy run to warm up the model."""
+        req = self._make_dummy_request(
+            height=512,
+            width=512,
+            guidance_scale=0.0,
+        )
+        if req is None:
+            logger.info("Skipping dummy warmup run (num_frames=0)")
+            return
         logger.info("dummy run to warm up the model")
         output = self.add_req_and_wait_for_response(req)
         if output.error:
@@ -1127,4 +1265,7 @@ class DiffusionEngine:
         if runner_output is not None and runner_output.async_output_id is not None:
             return DiffusionOutput(async_output_id=runner_output.async_output_id)
 
+        if state.status == DiffusionRequestStatus.FINISHED_ERROR and state.error:
+            return DiffusionOutput(error=state.error)
+
         return DiffusionOutput(error=missing_result_error)

vllm_omni/diffusion/diffusion_kv/initialization.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from vllm.config import VllmConfig
+from vllm.v1.core.kv_cache_utils import (
+    generate_scheduler_kv_cache_config,
+    get_kv_cache_configs,
+    resolve_kv_cache_block_sizes,
+)
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+from vllm_omni.platforms import current_omni_platform
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+def build_native_kv_cache_configs(
+    vllm_config: VllmConfig,
+    worker_specs: list[dict[str, KVCacheSpec]],
+    available_memory: list[int],
+) -> tuple[list[KVCacheConfig], KVCacheConfig]:
+    """Build rank-local and Scheduler-native configs using vLLM utilities."""
+
+    from vllm.v1.core import single_type_kv_cache_manager
+
+    single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
+    worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)
+    return worker_configs, generate_scheduler_kv_cache_config(worker_configs)
+
+
+def initialize_diffusion_kv_control_plane(
+    executor: DiffusionExecutor,
+    od_config: OmniDiffusionConfig,
+    *,
+    profile_requests: list[OmniDiffusionRequest] | None = None,
+    device: torch.device | None = None,
+) -> tuple[KVCacheConfig, int, int, VllmConfig] | None:
+    """Run the native Worker-spec to Scheduler-config initialization chain."""
+
+    if (
+        getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+        is not DiffusionKVCacheMode.PAGED_SCHEDULER
+    ):
+        return None
+    if not profile_requests:
+        raise ValueError("paged_scheduler Diffusion KV initialization requires prepared memory profile requests")
+    if device is None:
+        device = current_omni_platform.get_torch_device(0)
+    vllm_config = create_diffusion_vllm_config(device, od_config)
+    worker_specs = executor.get_kv_cache_specs()
+    available_memory = executor.determine_available_kv_memory(profile_requests)
+    if len(worker_specs) != od_config.num_gpus or len(available_memory) != od_config.num_gpus:
+        raise ValueError(
+            "Diffusion KV initialization rank count mismatch: "
+            f"expected={od_config.num_gpus}, specs={len(worker_specs)}, memory={len(available_memory)}"
+        )
+    worker_configs, scheduler_kv_cache_config = build_native_kv_cache_configs(
+        vllm_config,
+        worker_specs,
+        available_memory,
+    )
+    scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
+        scheduler_kv_cache_config,
+        vllm_config,
+    )
+    executor.set_kv_cache_configs(worker_configs)
+    return (
+        scheduler_kv_cache_config,
+        scheduler_block_size,
+        hash_block_size,
+        vllm_config,
+    )

vllm_omni/diffusion/diffusion_kv/manager.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+from vllm_omni.diffusion.diffusion_kv.metadata import (
+    DiffusionKVMetadata,
+    DiffusionKVSequenceMetadata,
+)
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+
+
+class DiffusionKVAdmissionError(ValueError):
+    """A request-specific KV admission failure safe to return to the caller."""
+
+
+class DiffusionKVCacheManager:
+    """Atomic public-request facade over native vLLM KVCacheManager."""
+
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        *,
+        max_model_len: int,
+        scheduler_block_size: int,
+        hash_block_size: int,
+        max_in_flight_tokens: int | None = None,
+    ) -> None:
+        if max_model_len <= 0:
+            raise ValueError(f"max_model_len must be positive, got {max_model_len}")
+        if kv_cache_config.num_blocks <= 0 or not kv_cache_config.kv_cache_groups:
+            raise ValueError("Diffusion KVCacheConfig must contain a positive block pool and at least one group")
+        if scheduler_block_size <= 0 or hash_block_size <= 0:
+            raise ValueError("scheduler_block_size and hash_block_size must be positive")
+        if max_in_flight_tokens is not None and max_in_flight_tokens <= 0:
+            raise ValueError(f"max_in_flight_tokens must be positive when provided, got {max_in_flight_tokens}")
+        self.native_manager = KVCacheManager(
+            kv_cache_config=kv_cache_config,
+            max_model_len=max_model_len,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            max_in_flight_tokens=max_in_flight_tokens,
+            enable_caching=False,
+        )
+        self.max_model_len = max_model_len
+        # Native vLLM may reserve a null block, so an idle BlockPool does not
+        # necessarily report ``kv_cache_config.num_blocks`` free blocks.
+        self._empty_pool_num_free_blocks = self.native_manager.block_pool.get_num_free_blocks()
+        self._requests: dict[str, tuple[DiffusionKVRequest, ...]] = {}
+        self._metadata: dict[str, DiffusionKVMetadata] = {}
+        self._internal_request_ids: set[str] = set()
+        self._next_allocation_generation = 1
+
+    def _get_empty_pool_required_blocks(self, requests: Sequence[DiffusionKVRequest]) -> int:
+        """Return the native full-sequence reservation for one public request.
+
+        vLLM performs this calculation inside ``allocate_slots`` when
+        ``full_sequence_must_fit`` is enabled. Diffusion CFG adds one semantic
+        requirement that native requests do not have: every execution sequence
+        must fit atomically. Sum the native per-sequence result so a request
+        that can never fit is rejected independently of current pool usage.
+        """
+
+        coordinator = self.native_manager.coordinator
+        empty_blocks = self.native_manager.empty_kv_cache_blocks.blocks
+        return sum(
+            coordinator.get_num_blocks_to_allocate(
+                request_id=request.request_id,
+                num_tokens=request.num_tokens,
+                new_computed_blocks=empty_blocks,
+                num_encoder_tokens=0,
+                total_computed_tokens=0,
+                num_local_computed_tokens=0,
+                num_tokens_main_model=request.num_tokens,
+                apply_admission_cap=True,
+            )
+            for request in requests
+        )
+
+    def has_request(self, public_request_id: str) -> bool:
+        return public_request_id in self._requests
+
+    def reserve_request(
+        self,
+        public_request_id: str,
+        kv_requests: Sequence[DiffusionKVRequest],
+    ) -> DiffusionKVMetadata | None:
+        """Reserve all CFG sequences or roll back the complete public request."""
+
+        if not public_request_id:
+            raise ValueError("public_request_id must be non-empty")
+        if public_request_id in self._requests:
+            raise ValueError(f"Diffusion KV request {public_request_id!r} is already allocated")
+        requests = tuple(kv_requests)
+        if not requests:
+            raise ValueError("Diffusion KV allocation requires at least one sequence")
+
+        sequence_ids = [request.sequence_id for request in requests]
+        internal_ids = [request.request_id for request in requests]
+        expected_sequence_ids = list(range(len(requests)))
+        if sequence_ids != expected_sequence_ids:
+            raise ValueError(
+                "Diffusion KV requests must be ordered by contiguous sequence_id values: "
+                f"expected={expected_sequence_ids}, got={sequence_ids}"
+            )
+        if len(internal_ids) != len(set(internal_ids)):
+            raise ValueError(f"Diffusion KV internal request IDs must be unique, got {internal_ids}")
+        conflicts = self._internal_request_ids.intersection(internal_ids)
+        if conflicts:
+            raise ValueError(f"Diffusion KV internal request IDs are already allocated: {sorted(conflicts)}")
+        for request in requests:
+            if request.kv_contexts:
+                raise DiffusionKVAdmissionError("independent DiffusionKVContext allocation is not implemented yet")
+            if request.seq_len > self.max_model_len:
+                raise DiffusionKVAdmissionError(
+                    f"Diffusion KV sequence {request.request_id!r} exceeds max_model_len: "
+                    f"seq_len={request.seq_len}, max_model_len={self.max_model_len}"
+                )
+
+        required_blocks = self._get_empty_pool_required_blocks(requests)
+        if required_blocks > self._empty_pool_num_free_blocks:
+            raise DiffusionKVAdmissionError(
+                f"Diffusion KV request {public_request_id!r} cannot fit even when the block pool is empty: "
+                f"required_blocks={required_blocks}, available_blocks={self._empty_pool_num_free_blocks}; "
+                "increase KV cache capacity or reduce the request sequence count/length"
+            )
+
+        allocated: list[DiffusionKVRequest] = []
+        sequence_metadata: list[DiffusionKVSequenceMetadata] = []
+        try:
+            for request in requests:
+                blocks = self.native_manager.allocate_slots(
+                    request,
+                    num_new_tokens=request.seq_len,
+                    delay_cache_blocks=True,
+                    full_sequence_must_fit=True,
+                )
+                if blocks is None:
+                    self._rollback(allocated)
+                    allocated.clear()
+                    return None
+                allocated.append(request)
+                sequence_metadata.append(
+                    DiffusionKVSequenceMetadata(
+                        sequence_id=request.sequence_id,
+                        prefix_len=request.prefix_len,
+                        target_len=request.target_len,
+                        seq_len=request.seq_len,
+                        block_ids=blocks.get_block_ids(),
+                    )
+                )
+        except Exception:
+            self._rollback(allocated)
+            raise
+
+        metadata = DiffusionKVMetadata(
+            request_id=public_request_id,
+            allocation_generation=self._next_allocation_generation,
+            sequences=tuple(sequence_metadata),
+        )
+        self._next_allocation_generation += 1
+        self._requests[public_request_id] = requests
+        self._metadata[public_request_id] = metadata
+        self._internal_request_ids.update(internal_ids)
+        return metadata
+
+    def get_metadata(self, public_request_id: str) -> DiffusionKVMetadata:
+        try:
+            return self._metadata[public_request_id]
+        except KeyError as exc:
+            raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc
+
+    def free_request(self, public_request_id: str) -> None:
+        requests = self._requests.pop(public_request_id, ())
+        self._metadata.pop(public_request_id, None)
+        for request in reversed(requests):
+            self.native_manager.free(request)
+            self._internal_request_ids.discard(request.request_id)
+
+    def close(self) -> None:
+        for public_request_id in tuple(self._requests):
+            self.free_request(public_request_id)
+
+    def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None:
+        for request in reversed(requests):
+            self.native_manager.free(request)

vllm_omni/diffusion/diffusion_kv/request.py

@@ -105,7 +105,7 @@ class DiffusionKVRequest:
         self.status = RequestStatus.WAITING
         self.num_preemptions = 0
         self.num_in_flight_tokens = 0
-        # vLLM 0.26 uses this optional boundary when publishing cached blocks;
+        # vLLM 0.26+ uses this optional boundary when publishing cached blocks;
         # zero means that no sparse-retention boundary is active.
         self.shared_prefix_boundary = 0
 

vllm_omni/diffusion/executor/abstract.py

@@ -5,10 +5,12 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 
 if TYPE_CHECKING:
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
     from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
     from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 
@@ -115,6 +117,39 @@ class DiffusionExecutor(ABC):
         """
         return None
 
+    def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
+        """Collect rank-local native specs after every Worker loads its model."""
+
+        result = self.collective_rpc(
+            "get_kv_cache_specs",
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, list):
+            raise TypeError(f"get_kv_cache_specs returned {type(result).__name__}, expected list")
+        return result
+
+    def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]:
+        """Profile and collect the KV memory budget on every Worker rank."""
+
+        result = self.collective_rpc(
+            "determine_available_kv_memory",
+            args=(profile_requests,),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, list) or not all(isinstance(value, int) for value in result):
+            raise TypeError("determine_available_kv_memory must return list[int]")
+        return result
+
+    def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None:
+        """Send each native rank-local config back to all Workers."""
+
+        # The default control-plane RPC mode executes on every rank and has
+        # rank 0 return the gathered rank statuses, so failures on nonzero
+        # ranks are not silently dropped.
+        self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs,))
+
     @abstractmethod
     def shutdown(self) -> None:
         """Shutdown the executor and release resources."""

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -465,16 +465,14 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         for new_req in new_reqs:
             validate_new_request_data_identity(new_req)
 
-        has_diffusion_kv_metadata = any(new_req.diffusion_kv_metadata is not None for new_req in new_reqs)
-
         # DP multi-concurrency: when DLO+AllGather is active and multiple
-        # requests are scheduled, send ALL requests in one broadcast RPC.
-        # Each rank picks req[rank % len(reqs)] and computes independently.
+        # requests are scheduled, send every complete NewRequestData envelope
+        # in one broadcast RPC. Each rank picks one envelope, keeping its
+        # request and Diffusion KV metadata inseparable.
         # All ranks reply (unique_reply_rank=None) so we collect dp_size
         # responses and match by dp_rank.
         if (
             len(new_reqs) > 1
-            and not has_diffusion_kv_metadata
             and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
             and getattr(self.od_config, "dlo_use_allgather", True)
         ):
@@ -507,12 +505,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     f"empty prompt request IDs: {empty_prompt_ids}."
                 )
 
-            reqs_list = [nr.req for nr in new_reqs]
             try:
                 results = self.collective_rpc(
                     "execute_model",
                     timeout=_DLO_DP_WAVE_TIMEOUT_S,
-                    args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job),
+                    args=(new_reqs, self.od_config, scheduler_output.kv_prefetch_job),
                     unique_reply_rank=None,
                     exec_all_ranks=True,
                 )

vllm_omni/diffusion/io_support.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import inspect
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.registry import DiffusionModelRegistry
 
 
@@ -47,3 +48,9 @@ def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool)
     if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
         return int(getattr(model_cls, "dummy_run_num_frames"))
     return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1
+
+
+def get_dummy_run_num_image_inputs(model_class_name: str) -> int:
+    """Return the maximum advertised image-input count for profiling."""
+
+    return get_diffusion_model_metadata(model_class_name).max_multimodal_image_inputs or 1

vllm_omni/diffusion/layers/activation.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Custom activation functions for diffusion models."""
+
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+from vllm_omni.platforms import current_omni_platform
+
+
+class SiluAndMul(CustomOp):
+    """SwiGLU activation dispatched by the current vLLM-Omni platform."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        if current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa():
+            self.op = torch.ops._C.silu_and_mul
+
+    @staticmethod
+    def forward_native(x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        return F.silu(x[..., :d]) * x[..., d:]
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        if x.device.type != "cuda":
+            return self.forward_native(x)
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        self.op(out, x)
+        return out
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)

vllm_omni/diffusion/layers/rope.py

@@ -423,16 +423,16 @@ class WanS2VRotaryPosEmbed(torch.nn.Module):
                         h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
                         w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
 
-                        freqs_0 = freqs_split[0][f_sam] if f_o >= 0 else freqs_split[0][f_sam].conj()
+                        freqs_0 = torch.index_select(freqs_split[0] if f_o >= 0 else freqs_split[0].conj(), 0, f_sam)
                         freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
 
                         freqs_i = torch.cat(
                             [
                                 freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
-                                freqs_split[1][h_sam]
+                                torch.index_select(freqs_split[1], 0, h_sam)
                                 .view(1, seq_h_int, 1, -1)
                                 .expand(seq_f_int, seq_h_int, seq_w_int, -1),
-                                freqs_split[2][w_sam]
+                                torch.index_select(freqs_split[2], 0, w_sam)
                                 .view(1, 1, seq_w_int, -1)
                                 .expand(seq_f_int, seq_h_int, seq_w_int, -1),
                             ],
@@ -523,14 +523,18 @@ class RotaryEmbeddingS2VGrid(torch.nn.Module):
                         h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
                         w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
 
-                        freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj()
+                        freqs_0 = torch.index_select(freqs[0] if f_o >= 0 else freqs[0].conj(), 0, f_sam)
                         freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
 
                         freqs_i = torch.cat(
                             [
                                 freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
-                                freqs[1][h_sam].view(1, seq_h_int, 1, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
-                                freqs[2][w_sam].view(1, 1, seq_w_int, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                torch.index_select(freqs[1], 0, h_sam)
+                                .view(1, seq_h_int, 1, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                torch.index_select(freqs[2], 0, w_sam)
+                                .view(1, 1, seq_w_int, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
                             ],
                             dim=-1,
                         ).reshape(seg_len, 1, -1)

vllm_omni/diffusion/lora/manager.py

@@ -9,7 +9,7 @@ from typing import get_args
 
 import torch
 import torch.nn as nn
-from vllm.config.lora import MaxLoRARanks
+from vllm.config.lora import LoRAConfig, MaxLoRARanks
 from vllm.logger import init_logger
 from vllm.lora.layers import BaseLayerWithLoRA
 from vllm.lora.lora_model import LoRAModel
@@ -23,7 +23,6 @@ from vllm.lora.utils import (
 )
 from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear
 
-from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.diffusion.lora.utils import (
     _expand_expected_modules_for_packed_layers,
     _match_target_modules,

vllm_omni/diffusion/lora/utils.py

@@ -5,8 +5,8 @@ from __future__ import annotations
 
 import torch.nn as nn
 from transformers import PretrainedConfig
+from vllm.config.lora import LoRAConfig
 
-from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.diffusion.lora.layers import (
     DiffusionColumnParallelLinearWithLoRA,
     DiffusionMergedColumnParallelLinearWithLoRA,

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -14,6 +14,7 @@ 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,
@@ -440,10 +441,13 @@ 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:
+                if _dist_offload and not _select_host_weight_cache:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -455,6 +459,16 @@ 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
 
@@ -482,16 +496,22 @@ class DiffusersPipelineLoader:
                         )
                 else:
                     if _dist_offload and _use_ag and _has_online_quant:
-                        raise ValueError(
-                            "Online quantization is incompatible with DLO+AllGather: "
-                            "the sharding + AllGather mechanism flattens weights by "
-                            "dtype, which breaks quantized weight/scale layouts. "
-                            "Please use --dlo-no-use-allgather or disable online "
-                            "quantization."
+                        unsupported_methods = self._unsupported_dlo_allgather_online_quant_methods(model)
+                        if unsupported_methods:
+                            raise ValueError(
+                                "DLO+AllGather supports online quantization only for "
+                                "per-tensor FP8 linears; unsupported online methods: "
+                                f"{', '.join(unsupported_methods)}. Please use "
+                                "--dlo-no-use-allgather or disable online quantization."
+                            )
+                        logger.info(
+                            "Online per-tensor FP8 with DLO+AllGather: using the "
+                            "ordinary loader before sharding finalized weights and scales"
                         )
                     if _dist_offload and plan_result is not None:
                         logger.info(
-                            "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
+                            "DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s",
+                            plan_result.fallback_code or "unknown",
                             plan_result.fallback_reason,
                         )
                     logger.debug("Loading weights on %s ...", load_device)
@@ -516,7 +536,110 @@ class DiffusersPipelineLoader:
                 logger.info("Quantization complete, offloaded model back to CPU")
 
         self._apply_skip_softmax_calibration(model)
-        return model.eval()
+        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",
+            )
 
     @staticmethod
     def _request_offload_after_quant(model: nn.Module) -> int:
@@ -548,6 +671,29 @@ class DiffusersPipelineLoader:
         ``uses_meta_device=True``, e.g. online FP8)."""
         return has_online_quantization(model)
 
+    @staticmethod
+    def _unsupported_dlo_allgather_online_quant_methods(model: nn.Module) -> tuple[str, ...]:
+        """Return unsupported online-quant methods for DLO AllGather.
+
+        Per-tensor online FP8 is safe after the ordinary loader has finalized
+        its weight and scale parameters. DLO shards those runtime tensors by
+        dtype and reconstructs their recorded shapes and strides before the
+        kernel consumes them. Other online methods may create different scale,
+        packing, or aliasing layouts and remain fail-closed until validated.
+        """
+        from vllm.model_executor.layers.quantization.online.fp8 import (
+            Fp8PerTensorOnlineLinearMethod,
+        )
+
+        unsupported: set[str] = set()
+        for module in model.modules():
+            quant_method = getattr(module, "quant_method", None)
+            if not getattr(quant_method, "uses_meta_device", False):
+                continue
+            if not isinstance(quant_method, Fp8PerTensorOnlineLinearMethod):
+                unsupported.add(type(quant_method).__name__)
+        return tuple(sorted(unsupported))
+
     def _apply_skip_softmax_calibration(self, model: nn.Module) -> None:
         from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
             apply_skip_softmax_calibration,

vllm_omni/diffusion/model_loader/host_weight_cache.py

@@ -0,0 +1,574 @@
+# 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

@@ -0,0 +1,287 @@
+# 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,6 +9,7 @@ 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
@@ -28,7 +29,7 @@ TensorTransform = Callable[[torch.Tensor], torch.Tensor]
 class TensorBinding:
     """One runtime tensor backed by one safetensors entry."""
 
-    checkpoint_key: str
+    storage_key: str
     file_path: str
     transform: TensorTransform | None = None
 
@@ -37,9 +38,16 @@ class TensorBinding:
 class HostWeightPlan:
     """Complete, prevalidated host backing consumed by an offload backend."""
 
-    backing_kind: str
+    backing_kind: Literal["checkpoint_mmap", "host_weight_cache"]
     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)
@@ -48,6 +56,7 @@ class HostWeightPlanResult:
 
     plan: HostWeightPlan | None
     fallback_reason: str | None = None
+    fallback_code: str | None = None
 
 
 class _PlanIncompatibleError(RuntimeError):
@@ -258,11 +267,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.checkpoint_key not in available:
+                if binding.storage_key not in available:
                     raise _PlanIncompatibleError(
-                        f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
+                        f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}"
                     )
-                source = handle.get_slice(binding.checkpoint_key)
+                source = handle.get_slice(binding.storage_key)
                 source_shape = tuple(source.get_shape())
                 source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
                 if source_shape != tuple(target.shape):
@@ -287,11 +296,19 @@ 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")
+        return HostWeightPlanResult(
+            None,
+            f"TP={tensor_parallel_size} requires the ordinary loader",
+            "unsupported_tp",
+        )
     if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp")
     if online_quantization:
-        return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
+        return HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+            "unsupported_online_quantization",
+        )
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
     if not callable(remap_fn):
@@ -328,14 +345,14 @@ def build_checkpoint_mmap_plan(
                 )
             checkpoint_key, file_path = model_to_ckpt[runtime_name]
             bindings[runtime_name] = TensorBinding(
-                checkpoint_key=checkpoint_key,
+                storage_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))
+        return HostWeightPlanResult(None, str(exc), "checkpoint_plan_incompatible")
 
     return HostWeightPlanResult(
         HostWeightPlan(

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -913,7 +913,7 @@ class HunYuanRotary2DEmbedder:
         return q, k
 
 
-class ImageKVCacheManager:
+class ImageKVCacheManager(nn.Module):
     """
     Manages specialized caching and updating of KV-Cache for image tokens in multimodal models.
     """
@@ -932,6 +932,7 @@ class ImageKVCacheManager:
             image_token_len: Number of tokens per image (including special placeholders),
             default 4097 (timestamp + 4096 image tokens).
         """
+        super().__init__()
         # attention related
         self.num_heads = num_heads
         self.num_kv_heads = num_kv_heads
@@ -953,6 +954,10 @@ class ImageKVCacheManager:
             softmax_scale=self.scaling,
             num_kv_heads=self.num_kv_heads,
             prefix=f"{prefix}.attn" if prefix else "",
+            paged_kv_cache_role="primary",
+            # Hunyuan image RoPE produces bfloat16 K and denoising runs under
+            # bfloat16 autocast, independent of the weight-loading dtype.
+            paged_kv_cache_dtype=torch.bfloat16,
         )
 
     @staticmethod
@@ -1095,7 +1100,7 @@ class ImageKVCacheManager:
         )
         return key, value
 
-    def __call__(
+    def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -28,11 +28,18 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
+    indexed_gate,
+    indexed_gate_rms_norm_scale_shift,
+    indexed_scale_shift_,
+    rms_norm_indexed_scale_shift,
+)
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
+from vllm_omni.diffusion.layers.activation import SiluAndMul
 from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
@@ -189,30 +196,6 @@ def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> RMSNorm
     return RMSNorm(size, eps=eps, dtype=dtype)
 
 
-def _modulate_scale_shift(
-    x: torch.Tensor,
-    shift: torch.Tensor,
-    scale: torch.Tensor,
-    indices: torch.Tensor,
-    *,
-    dtype: torch.dtype,
-) -> torch.Tensor:
-    # Apply per-index affine modulation: x * (1 + scale[idx]) + shift[idx].
-    return (x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(dtype)
-
-
-def _modulate_gate(
-    x: torch.Tensor,
-    gate: torch.Tensor,
-    other: torch.Tensor,
-    indices: torch.Tensor,
-    *,
-    dtype: torch.dtype,
-) -> torch.Tensor:
-    # Apply the per-index gated residual: x + gate[idx] * other.
-    return (x + gate.index_select(0, indices) * other).to(dtype)
-
-
 def _sequence_parallel_local_span(
     seq_len: int,
     *,
@@ -586,6 +569,7 @@ class MiniMaxH3MLP(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.fc1",
         )
+        self.act_fn = SiluAndMul()
         # Chunk the fused fc1 output as [gate, up], then compute
         # silu(gate) * up.
         self.fc2 = RowParallelLinear(
@@ -600,8 +584,7 @@ class MiniMaxH3MLP(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         hidden, _ = self.fc1(x)
-        gate, up = hidden.chunk(2, dim=-1)
-        hidden = nn.functional.silu(gate) * up
+        hidden = self.act_fn(hidden)
         out, _ = self.fc2(hidden)
         return out
 
@@ -794,8 +777,14 @@ class MiniMaxH3DiTBlock(nn.Module):
         ) = self.adaln_proj(t_emb)
 
         residual = x
-        h = self.norm1(x)
-        h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE)
+        h = rms_norm_indexed_scale_shift(
+            x,
+            self.norm1.weight,
+            shift_msa,
+            scale_msa,
+            combined_indices,
+            self.norm1.variance_epsilon,
+        )
         h = self.attn(
             h,
             rope_table=rope_table,
@@ -805,13 +794,19 @@ class MiniMaxH3DiTBlock(nn.Module):
             sp_seq_lens=sp_seq_lens,
             video_layout=video_layout,
         )
-        x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE)
-
+        x, h = indexed_gate_rms_norm_scale_shift(
+            residual,
+            gate_msa,
+            h,
+            self.norm2.weight,
+            shift_mlp,
+            scale_mlp,
+            combined_indices,
+            self.norm2.variance_epsilon,
+        )
         residual = x
-        h = self.norm2(x)
-        h = _modulate_scale_shift(h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE)
         h = self.mlp(h)
-        return _modulate_gate(residual, gate_mlp, h, combined_indices, dtype=_BF16_DTYPE)
+        return indexed_gate(residual, gate_mlp, h, combined_indices)
 
 
 class MiniMaxH3FinalLayer(nn.Module):
@@ -866,7 +861,7 @@ class MiniMaxH3FinalLayer(nn.Module):
         """
         shift, scale = self.adaln_proj(t_emb)
         h = self.norm(x)
-        h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE)
+        h = indexed_scale_shift_(h, shift, scale, inverse_indices)
         # Preserve full precision through both final output projections.
         h = h.to(_FP32_DTYPE)
         video, _ = self.video_out(h)

vllm_omni/diffusion/models/pi0/pipeline_pi0.py

@@ -22,15 +22,11 @@ import torch
 from torch import nn
 from vllm.logger import init_logger
 
-from vllm_omni.config.stage_config import (
-    PipelineConfig,
-    StageExecutionType,
-    StagePipelineConfig,
-)
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.pi0.config import Pi0Config
 from vllm_omni.diffusion.models.pi0.modeling_pi0 import Pi0ForActionPrediction
 from vllm_omni.diffusion.models.pi0.processor_pi0 import build_model_inputs
+from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE as PI0_PIPELINE
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 logger = init_logger(__name__)
@@ -54,32 +50,6 @@ def get_pi0_post_process_func(od_config: OmniDiffusionConfig):
     return _pi0_post_process
 
 
-# Single-stage diffusion topology for online OpenPI serving.
-#
-# π0 is one diffusion stage (robot observation -> action chunk via flow
-# matching). It is registered in ``OMNI_PIPELINES`` (rather than relying on the
-# default single-stage fallback) because π0 is *online-served* via
-# ``vllm serve --deploy-config pi0.yaml``; that path
-# (``OmniStagePipelineFactory._create_from_registry``) requires a registered
-# pipeline so the deploy yaml's ``policy_server_config`` reaches the OpenPI
-# websocket layer. Offline-only models (e.g. ``internvla_a1``) need no topology.
-PI0_PIPELINE = PipelineConfig(
-    model_type="pi0",
-    model_arch="Pi0Pipeline",
-    stages=(
-        StagePipelineConfig(
-            stage_id=0,
-            model_stage="diffusion",
-            execution_type=StageExecutionType.DIFFUSION,
-            input_sources=(),
-            final_output=True,
-            final_output_type="action",
-            model_arch="Pi0Pipeline",
-        ),
-    ),
-)
-
-
 class Pi0Pipeline(nn.Module):
     """π0 VLA pipeline: raw robot obs → continuous action chunk.
 

vllm_omni/diffusion/models/pi0_pipeline_config.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lightweight pipeline topology for π0 (Pi-Zero)."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+# π0 is one diffusion stage (robot observation -> action chunk via flow
+# matching). It is registered in ``OMNI_PIPELINES`` because π0 is online-served
+# via ``vllm serve --deploy-config pi0.yaml``. Keep this topology outside the
+# ``pi0`` package: importing a package submodule executes ``pi0.__init__``,
+# which imports the runtime pipeline and ``diffusion.data``.
+PI0_PIPELINE = PipelineConfig(
+    model_type="pi0",
+    model_arch="Pi0Pipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="action",
+            model_arch="Pi0Pipeline",
+        ),
+    ),
+)

vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py

@@ -98,7 +98,7 @@ class StableDiffusionXLPipeline(
             prefetch_list=sdxl_subfolders,
             local_files_only=local_files_only,
             torch_dtype=dtype,
-        )
+        ).to(self.device)
         self.text_encoder_2 = from_pretrained_with_prefetch(
             CLIPTextModelWithProjection.from_pretrained,
             model,
@@ -106,7 +106,7 @@ class StableDiffusionXLPipeline(
             prefetch_list=sdxl_subfolders,
             local_files_only=local_files_only,
             torch_dtype=dtype,
-        )
+        ).to(self.device)
         self.unet = SDXLUNet2DConditionModel(od_config=od_config)
         self.vae = from_pretrained_with_prefetch(
             DistributedAutoencoderKL.from_pretrained,
@@ -146,7 +146,7 @@ class StableDiffusionXLPipeline(
         )
         text_input_ids = text_inputs.input_ids
 
-        text_encoder_device = next(text_encoder.parameters()).device
+        text_encoder_device = self.device
         outputs = text_encoder(text_input_ids.to(text_encoder_device), output_hidden_states=True)
         prompt_embeds = outputs.hidden_states[-2].to(dtype=self.od_config.dtype, device=self.device)
 

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -17,7 +17,6 @@ Key integration points:
 
 from __future__ import annotations
 
-import json
 import math
 import os
 from collections.abc import Iterable
@@ -28,6 +27,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torchvision.transforms as T
+from huggingface_hub import snapshot_download
 from PIL import Image
 from transformers import AutoTokenizer
 from vllm.logger import init_logger
@@ -40,6 +40,9 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.transformers_utils.configs.sensenova_u1 import (
+    SenseNovaU1Config,
+)
 
 from .sensenova_u1_transformer import (
     SenseNovaU1ForCausalLM,
@@ -86,16 +89,12 @@ SYSTEM_MESSAGE_FOR_GEN = (
 
 
 # ---------------------------------------------------------------------------
-# Config parsed from config.json
+# Config
 # ---------------------------------------------------------------------------
-
-
 def _resolve_model_path(model_path: str) -> str:
     """Resolve a HuggingFace model ID or local path to a local directory."""
     if os.path.isdir(model_path):
         return model_path
-    from huggingface_hub import snapshot_download
-
     return snapshot_download(model_path)
 
 
@@ -182,73 +181,6 @@ def _load_image_native(
     return flatten_pv, grid_hw
 
 
-def _load_sensenova_u1_config(model_path: str):
-    """Parse config.json and return (top_config, llm_config, vision_config) namespaces."""
-    local_path = _resolve_model_path(model_path)
-    cfg_path = os.path.join(local_path, "config.json")
-    with open(cfg_path, encoding="utf-8") as f:
-        raw = json.load(f)
-
-    llm_raw = raw.get("llm_config", {})
-    vis_raw = raw.get("vision_config", {})
-
-    # LLM config needs these extra fields for 3D RoPE
-    llm_cfg = SimpleNamespace(
-        hidden_size=llm_raw.get("hidden_size", 4096),
-        intermediate_size=llm_raw.get("intermediate_size", 11008),
-        num_hidden_layers=llm_raw.get("num_hidden_layers", 32),
-        num_attention_heads=llm_raw.get("num_attention_heads", 32),
-        num_key_value_heads=llm_raw.get("num_key_value_heads", llm_raw.get("num_attention_heads", 32)),
-        head_dim=llm_raw.get("head_dim", llm_raw.get("hidden_size", 4096) // llm_raw.get("num_attention_heads", 32)),
-        hidden_act=llm_raw.get("hidden_act", "silu"),
-        rms_norm_eps=llm_raw.get("rms_norm_eps", 1e-6),
-        vocab_size=llm_raw.get("vocab_size", 152064),
-        rope_theta=llm_raw.get("rope_theta", 1000000.0),
-        rope_theta_hw=llm_raw.get("rope_theta_hw", 10000.0),
-        max_position_embeddings=llm_raw.get("max_position_embeddings", 40960),
-        max_position_embeddings_hw=llm_raw.get("max_position_embeddings_hw", 10000),
-        attention_bias=llm_raw.get("attention_bias", True),
-        layer_types=llm_raw.get("layer_types", ["full_attention"] * llm_raw.get("num_hidden_layers", 32)),
-        tie_word_embeddings=llm_raw.get("tie_word_embeddings", False),
-    )
-
-    vis_cfg = SimpleNamespace(
-        num_channels=vis_raw.get("num_channels", 3),
-        patch_size=vis_raw.get("patch_size", 16),
-        hidden_size=vis_raw.get("hidden_size", 1024),
-        llm_hidden_size=vis_raw.get("llm_hidden_size", [llm_cfg.hidden_size]),
-        downsample_ratio=vis_raw.get("downsample_ratio", [0.5]),
-        rope_theta_vision=vis_raw.get("rope_theta_vision", 10000.0),
-        max_position_embeddings_vision=vis_raw.get("max_position_embeddings_vision", 10000),
-        output_hidden_states=vis_raw.get("output_hidden_states", False),
-        use_return_dict=vis_raw.get("use_return_dict", True),
-    )
-
-    top_cfg = SimpleNamespace(
-        downsample_ratio=raw.get("downsample_ratio", 0.5),
-        template=raw.get("template", "neo1_0"),
-        fm_head_layers=raw.get("fm_head_layers", 2),
-        fm_head_dim=raw.get("fm_head_dim", 4096),
-        fm_head_mlp_ratio=raw.get("fm_head_mlp_ratio", 1.0),
-        use_pixel_head=raw.get("use_pixel_head", False),
-        noise_scale=raw.get("noise_scale", 1.0),
-        noise_scale_mode=raw.get("noise_scale_mode", "none"),
-        noise_scale_base_image_seq_len=raw.get("noise_scale_base_image_seq_len", 256),
-        noise_scale_max_value=raw.get("noise_scale_max_value", 10.0),
-        add_noise_scale_embedding=raw.get("add_noise_scale_embedding", False),
-        time_schedule=raw.get("time_schedule", "standard"),
-        time_shift_type=raw.get("time_shift_type", "exponential"),
-        base_shift=raw.get("base_shift", 0.5),
-        max_shift=raw.get("max_shift", 1.15),
-        base_image_seq_len=raw.get("base_image_seq_len", 256),
-        max_image_seq_len=raw.get("max_image_seq_len", 4096),
-        concat_time_token_num=raw.get("concat_time_token_num", 0),
-        t_eps=raw.get("t_eps", 0.02),
-    )
-
-    return top_cfg, llm_cfg, vis_cfg
-
-
 # ---------------------------------------------------------------------------
 # Vision Embedding — 2D RoPE + Conv patch embed (no transformer blocks)
 # ---------------------------------------------------------------------------
@@ -549,7 +481,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         model_path = od_config.model
         self.local_model_path = _resolve_model_path(model_path)
 
-        self.top_cfg, self.llm_cfg, self.vis_cfg = _load_sensenova_u1_config(model_path)
+        self.model_cfg = SenseNovaU1Config.from_pretrained(self.local_model_path)
+        self.llm_cfg = self.model_cfg.llm_config
+        self.vis_cfg = self.model_cfg.vision_config
 
         self.tokenizer = AutoTokenizer.from_pretrained(self.local_model_path)
         self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
@@ -573,13 +507,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         # FM modules
         llm_hidden_size = self.llm_cfg.hidden_size
         patch_size = self.vis_cfg.patch_size
-        merge_size = int(1 / self.top_cfg.downsample_ratio)
+        merge_size = int(1 / self.model_cfg.downsample_ratio)
         output_dim = 3 * (patch_size * merge_size) ** 2
 
         vision_model_mot_gen = NEOVisionModel(self.vis_cfg)
         timestep_embedder = TimestepEmbedder(llm_hidden_size)
 
-        if self.top_cfg.use_pixel_head:
+        if self.model_cfg.use_pixel_head:
             fm_head = ConvDecoder(llm_hidden_size)
         else:
             # The reference implementation uses a fixed 2-layer GELU MLP with a 4096-d
@@ -601,13 +535,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             }
         )
 
-        if self.top_cfg.add_noise_scale_embedding:
+        if self.model_cfg.add_noise_scale_embedding:
             self.fm_modules["noise_scale_embedder"] = TimestepEmbedder(llm_hidden_size)
 
         # Config shortcuts
         self.patch_size = patch_size
         self.merge_size = merge_size
-        self.downsample_ratio = self.top_cfg.downsample_ratio
+        self.downsample_ratio = self.model_cfg.downsample_ratio
 
         # Weight sources for diffusers_loader
         self.weights_sources = [
@@ -727,6 +661,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         t,
         z,
         image_token_num,
+        t_eps: float,
         image_size=None,
         cache_dit_skip=False,
         **_kw,
@@ -748,7 +683,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         )
         last_hidden_state = outputs.hidden_states
 
-        if self.top_cfg.use_pixel_head:
+        if self.model_cfg.use_pixel_head:
             merge_size = self.merge_size
             token_h = image_size[1] // (self.patch_size * merge_size)
             token_w = image_size[0] // (self.patch_size * merge_size)
@@ -765,7 +700,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         else:
             x_pred = self.fm_modules["fm_head"](last_hidden_state[:, -image_token_num:].view(B, L, -1)).view(B, L, -1)
 
-        v_pred = (x_pred - z) / (1 - t).clamp_min(self.top_cfg.t_eps)
+        v_pred = (x_pred - z) / (1 - t).clamp_min(t_eps)
         return v_pred
 
     def _apply_time_schedule(self, t, image_seq_len, timestep_shift):
@@ -1055,14 +990,14 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         token_w = p.image_size[0] // (self.patch_size * merge_size)
         grid_hw = torch.tensor([[grid_h, grid_w]] * p.batch_size, device=self.device)
 
-        noise_scale = self.top_cfg.noise_scale
-        if self.top_cfg.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"):
-            base = float(self.top_cfg.noise_scale_base_image_seq_len)
+        noise_scale = self.model_cfg.noise_scale
+        if self.model_cfg.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"):
+            base = float(self.model_cfg.noise_scale_base_image_seq_len)
             scale = math.sqrt((grid_h * grid_w) / (merge_size**2) / base)
-            noise_scale = scale * float(self.top_cfg.noise_scale)
-            if self.top_cfg.noise_scale_mode == "dynamic_sqrt":
+            noise_scale = scale * float(self.model_cfg.noise_scale)
+            if self.model_cfg.noise_scale_mode == "dynamic_sqrt":
                 noise_scale = math.sqrt(noise_scale)
-        noise_scale = min(noise_scale, self.top_cfg.noise_scale_max_value)
+        noise_scale = min(noise_scale, self.model_cfg.noise_scale_max_value)
 
         generator = torch.Generator(self.device).manual_seed(p.seed)
         image_prediction = noise_scale * torch.randn(
@@ -1101,6 +1036,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             z=z,
             image_token_num=ns.token_h * ns.token_w,
             image_size=p.image_size,
+            t_eps=p.t_eps,
         )
         if cache_dit_skip:
             kwargs["cache_dit_skip"] = True
@@ -1241,7 +1177,6 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
     @torch.inference_mode()
     def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         p = self._parse_request(req)
-        self.top_cfg.t_eps = p.t_eps
 
         input_images = self._extract_input_images(p.first_prompt)
         modalities = p.first_prompt.get("modalities", []) if isinstance(p.first_prompt, dict) else []
@@ -1438,8 +1373,8 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
                 ns.token_h * ns.token_w,
                 -1,
             )
-            if self.top_cfg.add_noise_scale_embedding:
-                ns_tensor = torch.full_like(t_expanded, ns.noise_scale / self.top_cfg.noise_scale_max_value)
+            if self.model_cfg.add_noise_scale_embedding:
+                ns_tensor = torch.full_like(t_expanded, ns.noise_scale / self.model_cfg.noise_scale_max_value)
                 ns_emb = self.fm_modules["noise_scale_embedder"](ns_tensor).view(
                     p.batch_size,
                     ns.token_h * ns.token_w,

vllm_omni/diffusion/offloader/__init__.py

@@ -71,8 +71,9 @@ 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, if ordinary
-            weight materialization was skipped.
+        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.
 
     Returns:
         OffloadBackend instance or None if offloading disabled
@@ -85,10 +86,14 @@ 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:
-        raise RuntimeError(
-            "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
-        )
+        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
 
     # Return None if no offloading requested
     if config.strategy == OffloadStrategy.NONE:
@@ -96,7 +101,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 host_weight_plan is not None:
+        if plan_requires_backend:
             raise RuntimeError(
                 "The loader skipped ordinary weight materialization for DLO, "
                 "but this platform cannot create the required offload backend"
@@ -112,7 +117,7 @@ def get_offload_backend(
         try:
             device = current_omni_platform.get_torch_device()
         except (NotImplementedError, AttributeError) as exc:
-            if host_weight_plan is not None:
+            if plan_requires_backend:
                 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,6 +31,9 @@ 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":
@@ -101,6 +104,7 @@ 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:
@@ -136,6 +140,7 @@ 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

@@ -0,0 +1,280 @@
+# 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,7 +18,9 @@ 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
@@ -27,6 +29,7 @@ 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,
 )
@@ -34,6 +37,12 @@ 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,
@@ -93,6 +102,7 @@ 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()
@@ -444,6 +454,26 @@ 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]]],
@@ -456,18 +486,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                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()}"
-                    )
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -481,6 +500,32 @@ 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]
@@ -514,17 +559,27 @@ 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:
-            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
+            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
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -734,6 +789,7 @@ 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]] = []
@@ -804,6 +860,14 @@ 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]
@@ -901,8 +965,11 @@ 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."""
@@ -1023,7 +1090,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     framework="pt",
                     device="cpu",
                 )
-            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
+            tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
             if is_parameter:
                 replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                 parent._parameters[leaf_name] = replacement
@@ -1103,6 +1170,179 @@ 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.
 
@@ -1443,17 +1683,41 @@ 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":
+            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:
                 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: checkpoint pages are "
-                    "node-shared; each worker owns only two bounded host staging slots"
+                    "DLO rank-local mmap storage enabled: file-backed pages are node-shared; "
+                    "transfer setup will select registered direct H2D or bounded host staging"
                 )
         else:
             remaining_meta = [
@@ -1623,6 +1887,9 @@ 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:
@@ -1632,17 +1899,14 @@ 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:
+        if self._using_rank_local_mmap and not self._using_registered_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1724,6 +1988,7 @@ 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
@@ -1746,8 +2011,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         if not self.enabled and not hasattr(self, "_mmap_file_cache"):
             return
 
-        if self._using_rank_local_mmap:
-            current_omni_platform.synchronize()
+        # 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()
 
         for blocks in self._blocks:
             for block in blocks:
@@ -1763,9 +2031,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
-        self._release_mmap_handles()
+        if registration_released:
+            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

@@ -0,0 +1,62 @@
+# 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",
+]