Add normalized online FP8 cache for DLO

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.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
@@ -81,6 +105,31 @@ 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.
 
+### Phase-I normalized online-FP8 cache
+
+The direct-checkpoint plan above cannot quantize a BF16 source online. Phase I
+adds an explicit, opt-in normalized cache for the narrower case of one TP1
+transformer with per-tensor online FP8 and DLO AllGather:
+
+```text
+cache miss:  BF16 source -> ordinary online FP8 loader -> atomic FP8 cache
+cache hit:   FP8 cache   -> direct mmap plan -> DLO host shard -> AllGather
+```
+
+The cache stores the serialized runtime representation, including finalized FP8
+weights, generated scales, and persistent transformer buffers. Its manifest
+records the schema, component, quantization contract, tensor count, source
+files, and a source fingerprint. A missing, stale, incomplete, incompatible, or
+invalid artifact is a cache miss; cache publication is best-effort and falls
+back to the ordinary loader on failure.
+
+Phase I keeps the cache binding deliberately simple: one source per transformer
+component, with existing per-tensor transforms applied before publication. The
+cache is not selected for DLO-disabled, no-AllGather, TP>1, or HSDP deployments,
+and it normalizes the transformer only. Cache sharing and compatibility for
+other components, transformed TP coordinates, and broader DP/SP layouts remain
+in [RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -181,13 +230,16 @@ requirements.
 
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
-ordinary loader. Online quantization remains incompatible with DLO AllGather;
-use `--dlo-no-use-allgather` or disable online quantization.
+ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
+ordinary loader finalizes their runtime weights and scales; DLO then shards and
+reconstructs those tensors with their recorded layouts. Other online methods
+must use `--dlo-no-use-allgather` or disable online quantization until their
+runtime layouts are validated.
 
-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).
+The Phase-I normalized online-FP8 cache is the supported TP1 transformer case.
+The general cache-compatibility contract for additional components, transformed
+TP coordinates, and cross-DP/SP reuse remains in
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
 
 ## Validation coverage
 
@@ -197,6 +249,8 @@ Current source-level validation includes:
 - HSDP + DLO without AllGather acceptance at configuration level;
 - loader preflight fallback for TP, HSDP, online quantization, unknown custom
   loaders, missing keys, and shape/dtype mismatches;
+- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
+  sharding of finalized weights and scales;
 - exact loader-to-backend plan transfer and ordinary-loader fallback;
 - rank-local mmap source retention, bounded two-slot staging, and adapter
   transforms without parameter-side flags;

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/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

@@ -42,6 +42,12 @@ vllm serve /path/to/model --omni \
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --usp 4
+
+# Opt-in normalized online-FP8 cache (TP1 + DLO AllGather)
+vllm serve /path/to/model --omni \
+  --quantization fp8 \
+  --enable-distributed-layerwise-offload \
+  --dlo-fp8-cache-dir /path/to/node-local/fp8-cache
 ```
 
 ```python
@@ -63,6 +69,7 @@ omni = Omni(
 | `--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-fp8-cache-dir PATH` | Opt into the Phase-I normalized online-FP8 cache for the transformer | unset |
 
 ## Host-weight loading
 
@@ -104,6 +111,24 @@ 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.
 
+## Normalized online-FP8 cache
+
+`--dlo-fp8-cache-dir` is an opt-in Phase-I cache for per-tensor online FP8 with
+TP1 and DLO AllGather. On a cache miss, the ordinary loader quantizes the
+transformer once and publishes finalized FP8 weights and scales as an atomic,
+serialized safetensors artifact. On a valid cache hit, the loader selects that
+serialized source before DLO planning, so the transformer can be loaded through
+the direct checkpoint-mmap path without re-materializing the full BF16 source.
+
+The cache is keyed to the source files through the manifest's schema, component,
+quantization contract, and source fingerprint. Incomplete, incompatible, or
+stale artifacts are treated as cache misses. Cache creation is best-effort: a
+publication failure logs a warning and keeps the ordinary online-FP8 path.
+
+Phase I normalizes the transformer only. Other pipeline components keep their
+ordinary loading and quantization behavior. The cache is not selected when DLO
+is disabled, AllGather is disabled, TP is greater than one, or HSDP is enabled.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -138,16 +163,18 @@ must enter each collective.
   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.
+- Per-tensor online FP8 linears use the ordinary loader and can run with either
+  DLO transfer path. With AllGather, every rank temporarily materializes the
+  complete FP8 model in host memory before DLO retains only its shard. Other
+  online quantization methods require no-AllGather until their runtime layouts
+  are validated.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
-[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
-the direct-checkpoint path.
+Sharing transformed TP or quantized runtime layouts beyond this TP1 transformer
+cache remains a follow-up design in
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

docs/user_guide/diffusion/parallelism/hsdp.md

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

docs/user_guide/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():
@@ -610,8 +733,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
 
 def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
     cfg = OmniStageDiffusionParallelConfig(
-        pipeline_parallel_size=2,
-        ulysses_degree=2,
+        ulysses_degree=4,
         use_hsdp=True,
         hsdp_shard_size=-1,
         hsdp_replicate_size=2,
@@ -622,10 +744,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
 
 
 def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with TP"):
         OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
-    with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+    with pytest.raises(ValueError, match="not compatible with DP"):
         OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
 
 
@@ -1078,3 +1200,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/distributed/test_distributed_vae_executor.py

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

tests/diffusion/distributed/test_expert_parallel_layout.py

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

tests/diffusion/distributed/test_hsdp.py

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

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -14,7 +14,7 @@ from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
 
 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 +91,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"])
@@ -394,6 +405,78 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     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_fp8_cache.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the Phase-I normalized online-FP8 cache contract."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.fp8_cache import (
+    CACHE_COMPLETE_NAME,
+    CACHE_COMPONENT_SUBFOLDER,
+    CACHE_MANIFEST_NAME,
+    _source_fingerprint,
+    _transformer_quant_config,
+    build_normalized_fp8_cache,
+    is_online_fp8_config,
+    load_fp8_cache_manifest,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _write_valid_cache(tmp_path):
+    source = tmp_path / "source.safetensors"
+    save_file({"weight": torch.ones(2, 2)}, str(source))
+
+    cache = tmp_path / "cache"
+    component = cache / CACHE_COMPONENT_SUBFOLDER
+    component.mkdir(parents=True)
+    shard = component / "part-00000.safetensors"
+    save_file({"weight": torch.ones(2, 2, dtype=torch.float8_e4m3fn)}, str(shard))
+    (component / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {
+                "metadata": {"total_size": 4},
+                "weight_map": {"weight": shard.name},
+            }
+        )
+    )
+
+    manifest = {
+        "schema_version": 1,
+        "cache_kind": "normalized_fp8_runtime",
+        "component": "transformer",
+        "component_subfolder": CACHE_COMPONENT_SUBFOLDER,
+        "source_prefix": "transformer.",
+        "quantization": {"method": "fp8", "serialized": True, "activation_scheme": "dynamic"},
+        "tensor_count": 1,
+        "source_files": [str(source)],
+        "source_fingerprint": _source_fingerprint({str(source)}),
+    }
+    cache.mkdir(exist_ok=True)
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+    (cache / CACHE_COMPLETE_NAME).write_text("phase1\n")
+    return cache, source, manifest
+
+
+def test_online_and_serialized_fp8_configs_are_classified():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    assert is_online_fp8_config(Fp8Config())
+    assert not is_online_fp8_config(Fp8Config(is_checkpoint_fp8_serialized=True))
+
+    cache_config = _transformer_quant_config()
+    assert not is_online_fp8_config(cache_config)
+    assert cache_config.resolve("transformer").is_checkpoint_fp8_serialized
+    assert cache_config.resolve("text_encoder") is None
+
+
+def test_load_fp8_cache_manifest_validates_artifact_and_sources(tmp_path):
+    cache, source, expected = _write_valid_cache(tmp_path)
+
+    assert load_fp8_cache_manifest(cache) == expected
+
+    source.write_bytes(source.read_bytes() + b"changed")
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_wrong_serialization_contract(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+
+    manifest["quantization"]["serialized"] = False
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_path_traversal(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+    index_path = cache / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
+    index_path.write_text(json.dumps({"weight_map": {"weight": "../outside.safetensors"}}))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_non_object_json(tmp_path):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / CACHE_MANIFEST_NAME).write_text("[]")
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_build_normalized_fp8_cache_serializes_weight_and_scale(tmp_path):
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        Fp8PerTensorOnlineLinearMethod,
+    )
+
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    source_file = source_dir / "model.safetensors"
+    save_file({"block.weight": torch.ones((3, 2), dtype=torch.bfloat16)}, str(source_file))
+
+    model = nn.Module()
+    model._dit_modules = ("transformer",)
+    model.transformer = nn.Module()
+    model.transformer.block = nn.Module()
+    model.transformer.block.weight = nn.Parameter(torch.ones((2, 3), dtype=torch.float8_e4m3fn), requires_grad=False)
+    model.transformer.block.weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
+    model.transformer.block.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
+
+    source = SimpleNamespace(
+        model_or_path=str(source_dir),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    cache_dir = tmp_path / "cache"
+
+    manifest = build_normalized_fp8_cache(
+        model,
+        sources=(source,),
+        cache_dir=cache_dir,
+    )
+
+    assert manifest["cache_kind"] == "normalized_fp8_runtime"
+    tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00000.safetensors"))
+    scale_tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00001.safetensors"))
+    assert tensors["block.weight"].dtype == torch.float8_e4m3fn
+    assert tuple(tensors["block.weight"].shape) == (3, 2)
+    assert torch.equal(scale_tensors["block.weight_scale"], torch.tensor([0.5]))
+
+
+@pytest.mark.parametrize("missing", [CACHE_COMPLETE_NAME, CACHE_MANIFEST_NAME])
+def test_load_fp8_cache_manifest_requires_completion_marker(tmp_path, missing):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / missing).unlink()
+
+    assert load_fp8_cache_manifest(cache) is None

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_distributed_layerwise_backend.py

@@ -185,6 +185,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))

tests/diffusion/test_diffusion_config_propagation.py

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

@@ -58,6 +58,23 @@ def test_default_stage_devices_from_sequence_parallel():
     assert stage_cfg["runtime"]["devices"] == "0,1,2,3"
 
 
+def test_default_stage_devices_and_dp_from_num_gpus():
+    """Resolve omitted DP before deriving devices from an explicit GPU count."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "num_gpus": 8,
+            "tensor_parallel_size": 2,
+            "ulysses_degree": 2,
+        }
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config.data_parallel_size == 2
+    assert parallel_config.world_size == 8
+    assert stage_cfg["engine_args"]["num_gpus"] == 8
+    assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
+
+
 def test_default_stage_config_uses_parallel_size_kwargs():
     """Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -76,8 +93,8 @@ def test_default_stage_config_uses_parallel_size_kwargs():
     assert parallel_config.enable_expert_parallel is True
 
 
-def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
-    """Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
+def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
+    """Keep omitted DP unresolved until the runtime WORLD size is known."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
         {
             "pipeline_parallel_size": None,
@@ -92,7 +109,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
 
     parallel_config = stage_cfg["engine_args"]["parallel_config"]
     assert parallel_config.pipeline_parallel_size == 1
-    assert parallel_config.data_parallel_size == 1
+    assert parallel_config.data_parallel_size is None
     assert parallel_config.tensor_parallel_size == 1
     assert parallel_config.enable_expert_parallel is False
     assert stage_cfg["engine_args"]["enforce_eager"] is False
@@ -330,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--dlo-no-use-allgather",
             "--dlo-resident-layers",
             "20",
+            "--dlo-fp8-cache-dir",
+            "/tmp/dlo-fp8-cache",
         ]
     )
 
@@ -340,9 +359,11 @@ 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_resident_layers == 20
+    assert args.dlo_fp8_cache_dir == "/tmp/dlo-fp8-cache"
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
     assert engine_args["dlo_resident_layers"] == 20
+    assert engine_args["dlo_fp8_cache_dir"] == "/tmp/dlo-fp8-cache"
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

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

@@ -12,12 +12,22 @@ from __future__ import annotations
 
 import copy
 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 +87,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 +159,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 +200,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 +243,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 +394,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 +404,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 +435,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 +493,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 +578,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
         )
@@ -489,12 +600,17 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             * self.cfg_parallel_size
         )
         if self.use_hsdp:
-            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
-                raise ValueError(
-                    "HSDP (use_hsdp=True) cannot be used with TP or DP "
-                    f"(tensor_parallel_size={self.tensor_parallel_size}, "
-                    f"data_parallel_size={self.data_parallel_size})"
-                )
+            incompatible = []
+            if self.tensor_parallel_size > 1:
+                incompatible.append("TP")
+            if self.data_parallel_size > 1:
+                incompatible.append("DP")
+            if self.pipeline_parallel_size > 1:
+                incompatible.append("PP")
+            if self.enable_expert_parallel:
+                incompatible.append("EP")
+            if incompatible:
+                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
             if self.hsdp_shard_size == -1:
                 if other_parallel_world_size == 1:
                     raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1")
@@ -577,6 +693,7 @@ class _DiffusionConfigProjection:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = Field(default=0, ge=0)
+    dlo_fp8_cache_dir: str | None = None
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -814,6 +931,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 +960,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 +1019,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 +1039,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 +1069,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 +1100,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 +1183,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 +1213,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 +1269,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 +1422,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 +1618,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 +1676,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 +1725,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

@@ -383,6 +383,7 @@ class StageDeployConfig:
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
     dlo_resident_layers: int | None = None
+    dlo_fp8_cache_dir: str | 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

@@ -154,8 +154,8 @@ class DiffusionParallelConfig:
     pipeline_parallel_size: int = 1
     """Number of pipeline parallel stages."""
 
-    data_parallel_size: int = 1
-    """Number of data parallel groups."""
+    data_parallel_size: int | None = None
+    """Number of data parallel groups. Inferred from WORLD when omitted."""
 
     tensor_parallel_size: int = 1
     """Number of tensor parallel groups."""
@@ -243,7 +243,7 @@ class DiffusionParallelConfig:
     def _validate_parallel_config(self) -> Self:
         """Validates the config relationships among the parallel strategies."""
         assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0"
-        assert self.data_parallel_size > 0, "Data parallel size must be > 0"
+        assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0"
         assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0"
         assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
         assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
@@ -283,10 +283,12 @@ class DiffusionParallelConfig:
                 self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
             )
 
-        # Calculate world_size from other parallelism dimensions
+        # Until the runtime WORLD size is known, an omitted DP dimension means
+        # one replica. OmniDiffusionConfig resolves it against num_gpus below.
+        data_parallel_size = self.data_parallel_size or 1
         other_parallel_world_size = (
             self.pipeline_parallel_size
-            * self.data_parallel_size
+            * data_parallel_size
             * self.tensor_parallel_size
             * self.sequence_parallel_size
             * self.cfg_parallel_size
@@ -297,13 +299,17 @@ class DiffusionParallelConfig:
         # 1. Standalone: when other parallelism is all 1, HSDP determines world_size
         # 2. Combined: HSDP overlays on top of other parallelism
         if self.use_hsdp:
-            if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
-                raise ValueError(
-                    "HSDP (use_hsdp=True) cannot be used with TP or DP "
-                    f"(tensor_parallel_size={self.tensor_parallel_size}, "
-                    f"data_parallel_size={self.data_parallel_size}). "
-                    "Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP."
-                )
+            incompatible = []
+            if self.tensor_parallel_size > 1:
+                incompatible.append("TP")
+            if data_parallel_size > 1:
+                incompatible.append("DP")
+            if self.pipeline_parallel_size > 1:
+                incompatible.append("PP")
+            if self.enable_expert_parallel:
+                incompatible.append("EP")
+            if incompatible:
+                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
             if self.hsdp_shard_size == -1:
                 # Auto-calculate: use other_parallel_world_size as shard_size
                 if self.hsdp_replicate_size <= 0:
@@ -338,6 +344,42 @@ class DiffusionParallelConfig:
         else:
             self.world_size = other_parallel_world_size
 
+    def resolve_data_parallel_size(self, world_size: int) -> int:
+        """Resolve or validate diffusion DP against the actual WORLD size."""
+        if world_size <= 0:
+            raise ValueError(f"WORLD size must be > 0, but got {world_size}")
+
+        if self.use_hsdp:
+            if self.data_parallel_size not in (None, 1):
+                raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1")
+            expected_world_size = self.hsdp_replicate_size * self.hsdp_shard_size
+            if world_size != expected_world_size:
+                raise ValueError(
+                    f"WORLD size ({world_size}) must equal HSDP size "
+                    f"({self.hsdp_replicate_size} x {self.hsdp_shard_size} = {expected_world_size})"
+                )
+            self.data_parallel_size = 1
+            self.world_size = world_size
+            return 1
+
+        non_dp_size = (
+            self.pipeline_parallel_size
+            * self.tensor_parallel_size
+            * self.sequence_parallel_size
+            * self.cfg_parallel_size
+        )
+        if world_size % non_dp_size != 0:
+            raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})")
+        inferred_data_parallel_size = world_size // non_dp_size
+        if self.data_parallel_size is not None and self.data_parallel_size != inferred_data_parallel_size:
+            raise ValueError(
+                f"data_parallel_size ({self.data_parallel_size}) does not match WORLD-derived value "
+                f"({inferred_data_parallel_size})"
+            )
+        self.data_parallel_size = inferred_data_parallel_size
+        self.world_size = world_size
+        return inferred_data_parallel_size
+
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig":
         """
@@ -710,6 +752,10 @@ class OmniDiffusionConfig:
     dlo_use_allgather: bool = True
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
+    # Optional Phase-I normalized online-FP8 cache. A cache miss builds the
+    # artifact after the ordinary online loader; a valid cache hit uses the
+    # serialized FP8 checkpoint-mmap path.
+    dlo_fp8_cache_dir: str | None = None
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 
@@ -851,6 +897,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.
@@ -936,6 +990,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 = {}
@@ -983,10 +1045,7 @@ class OmniDiffusionConfig:
             else:
                 self.num_gpus = 1
 
-        if self.num_gpus < self.parallel_config.world_size:
-            raise ValueError(
-                f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})"
-            )
+        self.parallel_config.resolve_data_parallel_size(self.num_gpus)
 
         if self.diffusion_compile_granularity == "full":
             incompatible_features = []

vllm_omni/diffusion/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/distributed/__init__.py

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

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

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

vllm_omni/diffusion/distributed/group_coordinator.py

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

vllm_omni/diffusion/distributed/hsdp.py

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

vllm_omni/diffusion/distributed/parallel_state.py

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

vllm_omni/diffusion/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/checkpoint_adapters/direct_mmap.py

@@ -41,6 +41,8 @@ class _MiniMaxH3DirectMmapAdapter:
         target: torch.Tensor,
     ) -> DirectMmapTensorPolicy:
         del target
+        if getattr(self.pipeline, "_dlo_fp8_cache_normalized", False):
+            return DirectMmapTensorPolicy(allow_custom_loader=True)
         transform = None
         suffix = ".qkv_proj.weight"
         if runtime_name.endswith(suffix):

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -8,7 +8,7 @@ import re
 import time
 from collections.abc import Generator, Iterable, Sequence
 from pathlib import Path
-from typing import cast
+from typing import Any, cast
 
 import torch
 from torch import nn
@@ -34,6 +34,12 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
+from vllm_omni.diffusion.model_loader.fp8_cache import (
+    _transformer_quant_config,
+    build_normalized_fp8_cache,
+    is_online_fp8_config,
+    load_fp8_cache_manifest,
+)
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     build_checkpoint_mmap_plan,
@@ -141,6 +147,10 @@ class DiffusersPipelineLoader:
         self.quant_config = od_config.quantization_config
         self.parallel_config = od_config.parallel_config
         self.host_weight_plan: HostWeightPlan | None = None
+        self._fp8_cache_dir: Path | None = None
+        self._fp8_cache_requested = False
+        self._fp8_cache_hit = False
+        self._fp8_cache_manifest: dict[str, Any] | None = None
 
     def take_host_weight_plan(self) -> HostWeightPlan | None:
         """Transfer the loader-produced plan to the offload backend."""
@@ -367,12 +377,61 @@ class DiffusersPipelineLoader:
         )
 
     def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]:
-        return tuple(
+        sources = tuple(
             cast(
                 Iterable[DiffusersPipelineLoader.ComponentSource],
                 getattr(model, "weights_sources", ()),
             )
         )
+        if not self._fp8_cache_hit or self._fp8_cache_dir is None or self._fp8_cache_manifest is None:
+            return sources
+
+        cache_prefix = str(self._fp8_cache_manifest["source_prefix"])
+        cache_subfolder = str(self._fp8_cache_manifest.get("component_subfolder", "transformer"))
+        return tuple(
+            dataclasses.replace(
+                source,
+                model_or_path=str(self._fp8_cache_dir),
+                subfolder=cache_subfolder,
+                revision=None,
+            )
+            if source.prefix == cache_prefix
+            else source
+            for source in sources
+        )
+
+    def _prepare_fp8_cache(self) -> None:
+        """Select the Phase-I cache path before constructing the pipeline."""
+        cache_dir = getattr(self.od_config, "dlo_fp8_cache_dir", None)
+        if not cache_dir or not is_online_fp8_config(self.quant_config):
+            return
+        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because distributed layerwise offload is disabled")
+            return
+        if not getattr(self.od_config, "dlo_use_allgather", True):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because DLO AllGather is disabled")
+            return
+        if int(getattr(self.parallel_config, "tensor_parallel_size", 1)) != 1:
+            logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache requires TP=1")
+            return
+        if bool(getattr(self.parallel_config, "use_hsdp", False)):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache does not support HSDP")
+            return
+
+        self._fp8_cache_dir = Path(cache_dir)
+        self._fp8_cache_requested = True
+        self._fp8_cache_manifest = load_fp8_cache_manifest(self._fp8_cache_dir)
+        if self._fp8_cache_manifest is None:
+            logger.info("DLO normalized FP8 cache miss: %s", self._fp8_cache_dir)
+            return
+
+        self._fp8_cache_hit = True
+        self.quant_config = _transformer_quant_config()
+        self.od_config.quantization_config = self.quant_config
+        logger.info(
+            "DLO normalized FP8 cache hit: %s; selecting serialized FP8 checkpoint path",
+            self._fp8_cache_dir,
+        )
 
     def _get_expected_parameter_names(self, model: nn.Module) -> set[str]:
         """Return parameter names that should be covered by strict load checks."""
@@ -403,6 +462,7 @@ class DiffusersPipelineLoader:
         self.host_weight_plan = None
         if load_format is None:
             load_format = "default"
+        self._prepare_fp8_cache()
         # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
         # weights are already quantized in the checkpoint — load directly on CPU.
         # For online quantization, load on device so quantization can run on accelerator,
@@ -413,7 +473,9 @@ class DiffusersPipelineLoader:
             is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr(
                 quant_cfg, "is_checkpoint_quantized", False
             )
-            if not is_offline:
+            if self._fp8_cache_hit:
+                logger.info("Normalized FP8 cache is already serialized; keeping model initialization on CPU")
+            elif not is_offline:
                 load_device = device.type
                 offload_after_quant = True
                 logger.info(
@@ -431,6 +493,11 @@ class DiffusersPipelineLoader:
                 )
             else:
                 model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
+                if self._fp8_cache_hit:
+                    # The direct-mmap adapter must see this before plan
+                    # construction so normalized QKV tensors bypass the raw
+                    # checkpoint grouped-QKV transform.
+                    setattr(model, "_dlo_fp8_cache_normalized", True)
 
                 _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
                 _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
@@ -482,12 +549,17 @@ 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(
@@ -515,6 +587,20 @@ class DiffusersPipelineLoader:
                 model.to("cpu")
                 logger.info("Quantization complete, offloaded model back to CPU")
 
+        if self._fp8_cache_requested and self._fp8_cache_dir is not None and not self._fp8_cache_hit:
+            try:
+                build_normalized_fp8_cache(
+                    model,
+                    sources=self._get_weight_sources(model),
+                    cache_dir=self._fp8_cache_dir,
+                )
+            except Exception as exc:
+                logger.warning(
+                    "DLO normalized FP8 cache build failed; retaining ordinary-loader path: %s",
+                    exc,
+                    exc_info=True,
+                )
+
         self._apply_skip_softmax_calibration(model)
         return model.eval()
 
@@ -548,6 +634,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/fp8_cache.py

@@ -0,0 +1,349 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Phase-I normalized FP8 cache support for DLO checkpoint mmap."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import uuid
+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_plan import (
+    _build_source_map,
+    _collect_required_dit_tensors,
+)
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+logger = init_logger(__name__)
+
+CACHE_SCHEMA_VERSION = 1
+CACHE_KIND = "normalized_fp8_runtime"
+CACHE_MANIFEST_NAME = "manifest.json"
+CACHE_COMPLETE_NAME = "COMPLETE"
+CACHE_COMPONENT_SUBFOLDER = "transformer"
+
+
+def _transformer_quant_config() -> ComponentQuantizationConfig:
+    """Build the serialized FP8 config used by a normalized cache."""
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    return ComponentQuantizationConfig(
+        {
+            "transformer": Fp8Config(
+                is_checkpoint_fp8_serialized=True,
+                activation_scheme="dynamic",
+            )
+        }
+    )
+
+
+def is_online_fp8_config(quant_config: Any) -> bool:
+    """Return whether the transformer component requests online FP8."""
+    if quant_config is None:
+        return False
+    if hasattr(quant_config, "resolve"):
+        quant_config = quant_config.resolve("transformer")
+    get_name = getattr(quant_config, "get_name", None)
+    if not callable(get_name) or str(get_name()).lower() != "fp8":
+        return False
+    return not bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
+
+
+def _source_fingerprint(files: set[str]) -> str:
+    entries: list[dict[str, int | str]] = []
+    for filename in sorted(files):
+        stat = os.stat(filename)
+        entries.append(
+            {
+                "path": os.path.abspath(filename),
+                "size": stat.st_size,
+                "mtime_ns": stat.st_mtime_ns,
+            }
+        )
+    payload = json.dumps(entries, sort_keys=True, separators=(",", ":")).encode()
+    return hashlib.sha256(payload).hexdigest()
+
+
+def load_fp8_cache_manifest(cache_dir: str | os.PathLike[str] | None) -> dict[str, Any] | None:
+    """Return a validated Phase-I manifest, or ``None`` for a cache miss."""
+    if cache_dir is None:
+        return None
+    root = Path(cache_dir)
+    manifest_path = root / CACHE_MANIFEST_NAME
+    if not (root / CACHE_COMPLETE_NAME).is_file() or not manifest_path.is_file():
+        return None
+    try:
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    if not isinstance(manifest, dict):
+        return None
+    if manifest.get("schema_version") != CACHE_SCHEMA_VERSION:
+        return None
+    if manifest.get("cache_kind") != CACHE_KIND:
+        return None
+    if manifest.get("component") != CACHE_COMPONENT_SUBFOLDER:
+        return None
+    quantization = manifest.get("quantization")
+    if not isinstance(quantization, dict) or quantization != {
+        "method": "fp8",
+        "serialized": True,
+        "activation_scheme": "dynamic",
+    }:
+        return None
+    if manifest.get("component_subfolder") != CACHE_COMPONENT_SUBFOLDER:
+        return None
+    if manifest.get("source_prefix") != f"{CACHE_COMPONENT_SUBFOLDER}.":
+        return None
+
+    index_path = root / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
+    if not index_path.is_file():
+        return None
+    try:
+        index = json.loads(index_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    if not isinstance(index, dict):
+        return None
+    weight_map = index.get("weight_map")
+    if not isinstance(weight_map, dict) or not weight_map or not all(isinstance(key, str) for key in weight_map):
+        return None
+    if manifest.get("tensor_count") != len(weight_map):
+        return None
+    for filename in weight_map.values():
+        if not isinstance(filename, str):
+            return None
+        relative_path = Path(filename)
+        if (
+            relative_path.is_absolute()
+            or relative_path.name != filename
+            or relative_path.suffix != ".safetensors"
+            or not (root / CACHE_COMPONENT_SUBFOLDER / relative_path).is_file()
+        ):
+            return None
+
+    source_files = manifest.get("source_files")
+    source_fingerprint = manifest.get("source_fingerprint")
+    if not isinstance(source_files, list) or not source_files or not isinstance(source_fingerprint, str):
+        return None
+    try:
+        if _source_fingerprint({str(filename) for filename in source_files}) != source_fingerprint:
+            return None
+    except (OSError, TypeError):
+        return None
+    return manifest
+
+
+def _cache_key_for_runtime_name(runtime_name: str, prefix: str) -> str:
+    if not runtime_name.startswith(prefix):
+        raise ValueError(f"runtime tensor {runtime_name!r} is outside cache source prefix {prefix!r}")
+    return runtime_name[len(prefix) :]
+
+
+def build_normalized_fp8_cache(
+    model: nn.Module,
+    *,
+    sources: tuple[object, ...],
+    cache_dir: str | os.PathLike[str],
+) -> dict[str, Any]:
+    """Write a normalized FP8 cache from a completed online-FP8 model.
+
+    The builder writes tensors one at a time into individual safetensors shards,
+    so it does not create a second full-model state dict. The first cache build
+    still follows the existing online loader; warm launches consume the cache
+    through the direct DLO mmap path.
+    """
+    cache_root = Path(cache_dir)
+    existing_manifest = load_fp8_cache_manifest(cache_root)
+    if existing_manifest is not None:
+        return existing_manifest
+    if cache_root.exists() and any(cache_root.iterdir()):
+        raise FileExistsError(f"FP8 cache directory is not empty and is not valid: {cache_root}")
+    cache_root.parent.mkdir(parents=True, exist_ok=True)
+
+    dit_names = tuple(getattr(model, "_dit_modules", ()))
+    dit_names = tuple(name for name in dit_names if hasattr(model, name))
+    if len(dit_names) != 1:
+        raise ValueError(f"Phase-I FP8 cache requires exactly one DiT module, got {dit_names}")
+    dit_name = dit_names[0]
+    if dit_name != CACHE_COMPONENT_SUBFOLDER:
+        raise ValueError(
+            f"Phase-I FP8 cache requires a top-level {CACHE_COMPONENT_SUBFOLDER!r} module, got {dit_name!r}"
+        )
+    source_prefix = f"{dit_name}."
+    dit_module = model.get_submodule(dit_name)
+    planned_sources = tuple(source for source in sources if getattr(source, "prefix", "") == source_prefix)
+    if len(planned_sources) != 1:
+        raise ValueError(f"Phase-I FP8 cache requires one source for {source_prefix!r}, got {len(planned_sources)}")
+
+    source_map = _build_source_map(
+        sources=planned_sources,
+        model_path=None,
+        remap_fn=getattr(type(model), "_remap_ckpt_key", None),
+    )
+    required = _collect_required_dit_tensors(((dit_name, dit_module),))
+    quantized_modules: dict[str, nn.Module] = {}
+    try:
+        from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
+    except ImportError as exc:  # pragma: no cover - vLLM 0.27 provides this class
+        raise RuntimeError("Phase-I FP8 cache requires vLLM's online FP8 implementation") from exc
+
+    for module_name, module in model.named_modules():
+        if not module_name.startswith(source_prefix):
+            continue
+        if isinstance(getattr(module, "quant_method", None), Fp8PerTensorOnlineLinearMethod):
+            if not hasattr(module, "weight_scale"):
+                raise ValueError(f"online FP8 module {module_name!r} has no generated weight_scale")
+            quantized_modules[module_name] = module
+
+    tmp_root = cache_root.parent / f".{cache_root.name}.tmp-{uuid.uuid4().hex}"
+    component_root = tmp_root / CACHE_COMPONENT_SUBFOLDER
+    component_root.mkdir(parents=True, exist_ok=False)
+    weight_map: dict[str, str] = {}
+    source_files: set[str] = set()
+    total_size = 0
+    part_index = 0
+
+    file_cache: dict[str, Any] = {}
+
+    def write_tensor(key: str, tensor: torch.Tensor) -> None:
+        nonlocal part_index, total_size
+        if key in weight_map:
+            raise ValueError(f"normalized FP8 cache has duplicate tensor key {key!r}")
+        tensor = tensor.detach().to("cpu").contiguous()
+        filename = f"part-{part_index:05d}.safetensors"
+        part_index += 1
+        save_file({key: tensor}, str(component_root / filename))
+        weight_map[key] = filename
+        total_size += tensor.numel() * tensor.element_size()
+
+    try:
+        for runtime_name, target in required.items():
+            binding_info = source_map.get(runtime_name)
+            module_name = runtime_name.removesuffix(".weight")
+            quant_module = quantized_modules.get(module_name) if runtime_name.endswith(".weight") else None
+            if runtime_name.endswith(".weight_scale"):
+                owner_name = runtime_name.removesuffix(".weight_scale")
+                if owner_name in quantized_modules:
+                    # Generated online scales are published in the dedicated
+                    # pass below; they have no raw checkpoint binding.
+                    continue
+            if quant_module is not None:
+                if binding_info is None:
+                    raise ValueError(f"no source binding for quantized weight {runtime_name!r}")
+                qweight = quant_module.weight.detach()
+                if qweight.ndim != 2:
+                    raise ValueError(f"Phase-I FP8 cache only supports 2D weights, got {runtime_name!r}")
+                # Online FP8 stores the final runtime weight as qweight.T. The
+                # serialized FP8 method expects the pre-process qweight shape.
+                qweight = qweight.t().contiguous()
+                # Normalized caches contain the post-checkpoint-adapter QKV
+                # order. The cache-hit adapter bypasses the raw-checkpoint
+                # grouped-QKV transform, then the serialized FP8 finalizer
+                # performs only the standard weight transpose.
+                write_tensor(binding_info[0], qweight)
+                source_files.add(binding_info[1])
+                continue
+
+            if binding_info is not None:
+                checkpoint_key, file_path = binding_info
+                source_files.add(file_path)
+                handle = file_cache.get(file_path)
+                if handle is None:
+                    handle = safe_open(file_path, framework="pt", device="cpu")
+                    file_cache[file_path] = handle
+                write_tensor(checkpoint_key, handle.get_tensor(checkpoint_key))
+            else:
+                # Constructor-derived persistent buffers have no checkpoint
+                # entry; preserve them as canonical cache tensors.
+                write_tensor(_cache_key_for_runtime_name(runtime_name, source_prefix), target)
+
+        for module_name, module in quantized_modules.items():
+            scale = module.weight_scale.detach()
+            logical_widths = getattr(module, "logical_widths", None)
+            if scale.numel() == 1 and logical_widths is not None and len(logical_widths) > 1:
+                # Serialized Fp8LinearMethod represents a fused per-tensor
+                # weight with one scale slot per logical shard. Online FP8
+                # produces one scalar for the fused tensor, so replicate that
+                # scalar to the offline runtime shape without changing its
+                # value.
+                scale = scale.reshape(1).expand(len(logical_widths)).contiguous()
+            write_tensor(
+                _cache_key_for_runtime_name(f"{module_name}.weight_scale", source_prefix),
+                scale,
+            )
+
+        index = {
+            "metadata": {"total_size": total_size},
+            "weight_map": dict(sorted(weight_map.items())),
+        }
+        (component_root / "model.safetensors.index.json").write_text(
+            json.dumps(index, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        manifest = {
+            "schema_version": CACHE_SCHEMA_VERSION,
+            "cache_kind": CACHE_KIND,
+            "component": "transformer",
+            "component_subfolder": CACHE_COMPONENT_SUBFOLDER,
+            "source_prefix": source_prefix,
+            "quantization": {
+                "method": "fp8",
+                "serialized": True,
+                "activation_scheme": "dynamic",
+            },
+            "tensor_count": len(weight_map),
+            "source_fingerprint": _source_fingerprint(source_files),
+            "source_files": sorted(source_files),
+        }
+        (tmp_root / CACHE_MANIFEST_NAME).write_text(
+            json.dumps(manifest, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        (tmp_root / CACHE_COMPLETE_NAME).write_text("phase1\n", encoding="utf-8")
+        for handle in file_cache.values():
+            close = getattr(handle, "__exit__", None)
+            if callable(close):
+                close(None, None, None)
+        try:
+            os.replace(tmp_root, cache_root)
+        except FileExistsError:
+            # Multiple DLO workers may build the same cache concurrently. If
+            # another worker published a complete artifact first, treat this
+            # publisher as a successful cache race rather than a build error.
+            published = load_fp8_cache_manifest(cache_root)
+            if published is None:
+                raise
+            shutil.rmtree(tmp_root, ignore_errors=True)
+            logger.info("Normalized FP8 cache was published concurrently at %s", cache_root)
+            return published
+        logger.info("Published normalized FP8 cache at %s (%d tensors)", cache_root, len(weight_map))
+        return manifest
+    except BaseException:
+        for handle in file_cache.values():
+            close = getattr(handle, "__exit__", None)
+            if callable(close):
+                close(None, None, None)
+        shutil.rmtree(tmp_root, ignore_errors=True)
+        raise
+
+
+__all__ = [
+    "CACHE_SCHEMA_VERSION",
+    "CACHE_KIND",
+    "CACHE_COMPONENT_SUBFOLDER",
+    "build_normalized_fp8_cache",
+    "is_online_fp8_config",
+    "load_fp8_cache_manifest",
+    "_transformer_quant_config",
+]

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/minimax_h3/pipeline_minimax_h3.py

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

vllm_omni/diffusion/models/minimax_h3/vae.py

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

vllm_omni/diffusion/models/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/base.py

@@ -62,7 +62,7 @@ class OffloadConfig:
         dp_size = 1
         if parallel_config is not None:
             dp_size = getattr(parallel_config, "data_parallel_size", 1)
-            # HSDP's fully_shard_degree also contributes to effective DP
+            # HSDP shard and replica sizes determine the effective group size.
             hsdp_shard_size = getattr(parallel_config, "hsdp_shard_size", -1) if use_hsdp else -1
             hsdp_replicate_size = getattr(parallel_config, "hsdp_replicate_size", 1) if use_hsdp else 1
             if use_hsdp and hsdp_shard_size > 0:

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1034,6 +1034,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names))
 
+        if getattr(pipeline, "_dlo_fp8_cache_normalized", False):
+            self._process_normalized_fp8_cache(pipeline)
+
         # Keep file handles open — _shard_and_pin will read from the mmap views.
         # They will be released after _shard_and_pin completes (when params are
         # replaced with offload placeholders).
@@ -1103,6 +1106,52 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 f"(first 5: {remaining_meta[:5]})."
             )
 
+    @staticmethod
+    def _process_normalized_fp8_cache(pipeline: nn.Module) -> None:
+        """Finalize serialized FP8 parameters after mmap replacement.
+
+        Phase-I cache files use the serialized FP8 checkpoint layout: the
+        weight is stored before the quantization method's final transpose, and
+        ``weight_scale`` is stored as an explicit cache tensor.  The ordinary
+        loader normally runs this method before DLO owns the plan; mmap loading
+        must run the same finalizer after replacing meta parameters.
+        """
+        from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+        from vllm.model_executor.utils import replace_parameter
+
+        for module in pipeline.modules():
+            quant_method = getattr(module, "quant_method", None)
+            if not isinstance(quant_method, Fp8LinearMethod):
+                continue
+            if getattr(module, "_already_called_process_weights_after_loading", False):
+                continue
+
+            # ``Fp8LinearMethod.process_weights_after_loading`` also handles
+            # non-serialized checkpoints.  Its common path calls
+            # ``process_fp8_weight_tensor_strategy`` to requantize fused
+            # tensors, which is a CUDA-only operation.  A normalized cache
+            # already contains the serialized FP8 values and the generated
+            # scales, so only perform the layout transition and the kernel's
+            # own preprocessing here.  This keeps the warm path CPU/mmap
+            # compatible and preserves the cache bytes exactly.
+            if quant_method.block_quant:
+                raise ValueError("Phase-I normalized FP8 cache does not support block-quantized weights")
+            replace_parameter(module, "weight", module.weight.t().data)
+            # The online quantizer records one scalar.  A serialized FP8
+            # checkpoint may expose one repeated slot per logical shard in a
+            # fused module, but those slots must agree before they can be
+            # collapsed to the canonical per-tensor ``[1, 1]`` form.
+            weight_scale = module.weight_scale.reshape(-1)
+            if weight_scale.numel() > 1 and not torch.all(weight_scale == weight_scale[0]):
+                raise ValueError("Phase-I normalized FP8 cache contains non-uniform fused weight scales")
+            weight_scale = weight_scale[:1].reshape(1, 1)
+            replace_parameter(module, "weight_scale", weight_scale.data)
+            module.input_scale = None
+            if quant_method.use_marlin and hasattr(quant_method.fp8_linear, "marlin_input_dtype"):
+                quant_method.fp8_linear.marlin_input_dtype = quant_method.marlin_input_dtype
+            quant_method.fp8_linear.process_weights_after_loading(module)
+            module._already_called_process_weights_after_loading = True
+
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.