[Cleanup][Full Duplex] Defer handoff recovery to PR 6170

Signed-off-by: Sy03 <1370724210@qq.com>

docs/design/feature/attention_backend_selection.md

@@ -61,8 +61,7 @@ unstructured keyword dictionaries:
 - `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
   validation;
 - `skip_softmax` is consumed by TRTLLM; and
-- `block_sparse` is consumed by block-sparse backends such as RainFusion; and
-- `fastvideo_vsa_topk` is consumed by FastVideo VSA.
+- `block_sparse` is consumed by block-sparse backends such as RainFusion.
 
 A backend reads only the fields it owns and rejects incompatible values. New
 options should be added to a shared typed spec only when more than one backend
@@ -88,19 +87,3 @@ An implementation change is complete only when it:
    corresponding user guide; and
 6. updates the overview matrix without moving algorithm details back into the
    landing page.
-
-## FastVideo VSA model metadata
-
-FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in
-the original latent video grid. Wan integrations therefore attach the
-post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The
-backend validates that `T * H * W` equals the sequence length, derives the
-runtime block count, and selects the configured top-k key/value blocks for
-each query block.
-
-Checkpoint-specific compensation remains model-owned. Wan layers pass an
-optional learned `gate_compress` projection through attention metadata when
-the checkpoint contains those weights. Checkpoints without the projection do
-not allocate or execute it. When top-k selects every block, native checkpoints
-route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path
-to retain checkpoint semantics.

docs/user_guide/diffusion/attention_backends.md

@@ -20,7 +20,6 @@ For the internal selector, registry, and platform contract, see
 | Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
 | Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
-| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
 
 ## Backend options
 
@@ -36,7 +35,6 @@ For the internal selector, registry, and platform contract, see
 | `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
-| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
 
 ## Configuration
 
@@ -107,7 +105,6 @@ Backend-specific typed blocks are documented with their consumers:
   and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
 - `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
 - `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
-- `fastvideo_vsa_topk`: [FastVideo VSA](attention_backends/fastvideo_vsa.md#choose-top-k).
 
 ## Platform defaults
 

docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md

@@ -1,125 +0,0 @@
-# FastVideo VSA
-
-FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's
-self-attention by partitioning the post-patch latent grid into spatiotemporal
-blocks. For every query block, VSA scores the key/value blocks and computes
-attention only against the selected top-k blocks.
-
-The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
-
-VSA is a CUDA-only, explicitly selected backend. It requires the
-`fastvideo-kernel` package and currently supports non-causal self-attention
-with equal query and key/value sequence lengths. Unsupported shapes, masks,
-dtypes, sequence-parallel execution, or kernel failures fall back to
-`TORCH_SDPA` and emit a warning with the reason.
-
-## Enable the backend
-
-For online serving, select the backend with the existing attention backend
-flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained
-for every query block:
-
-```bash
-vllm-omni serve <model> \
-  --diffusion-attention-backend FASTVIDEO_VSA \
-  --fastvideo-vsa-topk 64
-```
-
-The backwards-compatible environment variable selects the backend with the
-default `topk=64`:
-
-```bash
-export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA
-vllm-omni serve <model>
-```
-
-To tune top-k, pass the CLI backend and top-k flags together as shown above,
-or use the equivalent structured configuration:
-
-```bash
-vllm-omni serve <model> \
-  --diffusion-attention-config \
-  '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}'
-```
-
-Do not combine `--diffusion-attention-backend` with an explicit
-`diffusion_attention_config.default.backend`. The top-k value must be positive
-and is valid only when the default backend is `FASTVIDEO_VSA`.
-
-For a deploy YAML stage, use either the shorthand fields:
-
-```yaml
-stages:
-  - stage_id: 0
-    diffusion_attention_backend: FASTVIDEO_VSA
-    fastvideo_vsa_topk: 64
-```
-
-or the structured configuration:
-
-```yaml
-stages:
-  - stage_id: 0
-    diffusion_attention_config:
-      default:
-        backend: FASTVIDEO_VSA
-        fastvideo_vsa_topk: 64
-```
-
-## Choose top-k
-
-At runtime the backend logs the sequence shape and derived block count:
-
-```text
-FASTVIDEO_VSA routing: seq_len=27280, dit_seq_shape=(31, 22, 40),
-block_size=(4, 8, 8), num_blocks=120, topk=64, keep_ratio=53.3%,
-checkpoint_mode=native, route=VSA
-```
-
-Use `num_blocks` as the upper bound when tuning top-k. Top-k is the number of
-key/value blocks retained for each query block, not the number of tokens or
-the total number of blocks processed by the layer.
-
-- A smaller top-k increases sparsity and may reduce attention computation, but
-  it can remove relevant blocks and reduce visual quality or temporal
-  consistency. Routing and padding overhead can also make a smaller value
-  slower for some shapes.
-- A larger top-k retains more context and generally approaches dense-attention
-  quality, but reduces the potential speedup and increases memory traffic.
-- `topk > num_blocks` is invalid for the runtime shape and falls back to SDPA.
-- For a native checkpoint, `topk == num_blocks` routes to SDPA because scoring
-  every block provides no sparsity benefit.
-- For a FastVideo DMD checkpoint, `topk == num_blocks` stays on the VSA
-  all-block path so the checkpoint keeps its trained compensation semantics.
-
-There is no universally optimal value. Resolution, frame count, GPU, kernel
-version, and checkpoint all affect both quality and latency. Start from the
-logged `num_blocks`, test several keep ratios on the target workload, and
-compare output quality against the same checkpoint running dense attention.
-
-## Checkpoint behavior
-
-The backend does not expose a user-selectable gate mode. Wan checkpoints that
-contain learned `to_gate_compress` weights use that projection automatically.
-When those weights are absent, the unused projection is removed and VSA uses
-the sparse branch without learned compensation.
-
-FastVideo DMD checkpoints use their fixed distilled timestep schedule. Native
-Wan checkpoints keep their normal scheduler and inference-step configuration;
-selecting VSA does not turn a native checkpoint into a distilled model.
-
-## Verify routing and fallback
-
-Check the startup and first-forward logs instead of assuming that selecting
-the backend guarantees sparse execution:
-
-- `route=VSA` means top-k block selection is active.
-- `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all
-  blocks through the VSA kernel.
-- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
-  executed; the warning includes the reason.
-
-The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
-standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
-and no active sequence-parallel context. NPU and XPU paths do not execute the
-FastVideo VSA CUDA kernel.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -80,7 +80,7 @@ between the legacy and direct planar paths
 ([full result](https://github.com/vllm-project/vllm-omni/pull/6288#issuecomment-5337546499)):
 
 | Metric | Legacy | Direct planar | Change |
-| --- | ---: | ---: | ---: |
+|---|---:|---:|---:|
 | Median wall time | 1.805 s | 1.394 s | -22.8% |
 | Median process CPU time | 3.613 s | 3.207 s | -11.2% |
 | Peak RSS | 3182 MiB | 2794 MiB | -387 MiB (-12.2%) |
@@ -183,7 +183,7 @@ vllm serve "${MODEL}" \
 Use the profile that matches the per-GPU memory capacity:
 
 | Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
-| --- | ---: | ---: | ---: | --- | --- | --- |
+|---|---:|---:|---:|---|---|---|
 | `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
 | `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
 
@@ -202,8 +202,8 @@ increasing it on a different request shape.
 At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
 2 x RTX 5090 without OOM:
 
-| Shape    | Frames        | Client E2E | Sampled peak/GPU       | Output validation                                            |
-| -------: | ------------: | ---------: | ---------------------: | -----------------------------------------------------------: |
+| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
+|---:|---:|---:|---:|---|
 | 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
 
 This is a single end-to-end validation run, not a warmed multi-run latency
@@ -272,7 +272,7 @@ vllm serve "${MODEL}" \
 Do not add `--enforce-eager` to this performance configuration. The first
 request includes regional compilation; warm the server once before measuring
 steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
-`1`. The H3 VAE supports its native `tile` mode, not
+1. The H3 VAE supports its native `tile` mode, not
 `spatial_shard_height` or `spatial_shard_width`.
 
 ### Attention Backends
@@ -530,7 +530,7 @@ vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the
 `0.26.0+rocm723` wheel (HIP 7.2).
 
 | Workload | Configuration | Observed result |
-| ---------- | --------------- | ----------------- |
+|----------|---------------|-----------------|
 | T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup |
@@ -712,7 +712,7 @@ at most 15 seconds combined.
 ## Official input matrix and limits
 
 | Task | Supported references | Limits |
-| ------ | ---------------------- | -------- |
+|------|----------------------|--------|
 | T2VA | text only | prompt must be non-empty |
 | FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` |
 | Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference |
@@ -743,17 +743,13 @@ Omitting `quality` preserves the startup default: it uses the reference path
 normally, or the server-configured profile when the server was started with
 `--cache-backend cache_dit`.
 
-Turbo is independent of this quality switch: `quality` selects a Cache-DiT
-policy, while Turbo changes the active LoRA weights and sampling schedule. See
-[LoRA](#lora) below.
-
 The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
 TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
 `lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
 balanced switch order.
 
 | `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
-| --- | ---: | ---: | ---: | ---: | --- |
+|---|---:|---:|---:|---:|---|
 | `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
 | `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
 
@@ -763,49 +759,10 @@ balanced switch order.
 > workload. The values above apply to this deployment and are not universal
 > guarantees. `lossless` remains the exact reference path.
 
-## LoRA
-
-### Turbo LoRA
-
-Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported:
-
-```text
-lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
-```
-
-Download only that file:
-
-```bash
-export TURBO_DIR=/path/to/minimax-h3-turbo
-export TURBO_FILE=minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
-hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
-export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
-```
-
-Start from a non-offloaded FL2VA server command and add
-`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
-`--lora-path` preloads the adapter; each request still activates it and uses
-the published sampling settings:
-
-```bash
--F 'num_inference_steps=5' \
--F 'flow_shift=6' \
--F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \
--F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}"
-```
-
-For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
-ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
-dynamic-only and does not support prefusion, DLO, or LoRA composition.
-It also rejects model-level, layerwise, and distributed layerwise offload;
-the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
-The five requested sigma points produce the four denoiser evaluations expected
-by the Turbo artifact.
-
 ## Key parameters
 
 | Parameter | Recommended value | Notes |
-| ----------- | ------------------- | ------- |
+|-----------|-------------------|-------|
 | `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
 | `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
 | `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
@@ -830,9 +787,9 @@ Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 servi
 The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
 excluded warmup followed by three requests.
 
-| Workload                               | Configuration                               | Observed result                      |
-| -------------------------------------- | ------------------------------------------- | ------------------------------------ |
-| FL2VA, 209 frames, 1248x768            | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency    |
+| Workload | Configuration | Observed result |
+|----------|---------------|-----------------|
+| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
 | Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
 
 These measurements describe the validated shapes rather than a general

tests/config/test_config_factory.py

@@ -1304,7 +1304,6 @@ stages:
             ({"model_type": "step_audio_2"}, None, "step_audio_2"),
             (None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"),
             (None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"),
-            (None, {"_class_name": "WanDMDPipeline"}, "wan2_2_ti2v"),
         ],
     )
     def test_migrated_models_are_discovered_without_explicit_deploy(
@@ -1854,18 +1853,6 @@ class TestQwen3TTSPipeline:
         # Stage 1 uses its per-stage override
         assert stages[1].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav"
 
-    @pytest.mark.parametrize(
-        "deploy_name",
-        ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"],
-    )
-    def test_default_async_deploy_enables_chunk_ramp(self, deploy_name):
-        deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
-
-        assert deploy.async_chunk is True
-        assert deploy.connectors is not None
-        extra = deploy.connectors["connector_of_shared_memory"]["extra"]
-        assert extra["codec_chunk_ramp"] == [4, 4, 8, 16, 25]
-
     def test_subtalker_sampling_params_deep_merge_preserves_base_keys(self):
         """Verify subtalker sampling params participate in stage deep-merge."""
         base = {

tests/config/test_omni_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for the additive structured Omni config."""
 
 from __future__ import annotations
@@ -1152,35 +1152,6 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
     assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
 
 
-def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch):
-    from vllm_omni.platforms import current_omni_platform
-
-    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
-    deploy_path = tmp_path / "dreamzero_fastvideo_vsa.yaml"
-    deploy_path.write_text(
-        "\n".join(
-            [
-                "pipeline: dreamzero",
-                "async_chunk: false",
-                "stages:",
-                "  - stage_id: 0",
-                "    diffusion_attention_backend: FASTVIDEO_VSA",
-                "    fastvideo_vsa_topk: 96",
-            ]
-        )
-    )
-
-    stage = _from_pipeline_key(
-        "dreamzero",
-        deploy_config_path=str(deploy_path),
-    ).stage_by_id(0)
-
-    attention_config = stage.diffusion_config.diffusion_attention_config
-    assert attention_config.default is not None
-    assert attention_config.default.backend == "FASTVIDEO_VSA"
-    assert attention_config.default.fastvideo_vsa_topk == 96
-
-
 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
     deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml"
     deploy_path.write_text(

tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py

@@ -0,0 +1,89 @@
+"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
+
+``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
+handing control to upstream, because upstream ``Scheduler.schedule()`` raises
+``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
+admits in a finished state -- which kills the stage's engine core, not just
+the request.
+
+Regression for the sweep missing ``skipped_waiting``, the third queue
+upstream admits from. An aborted duplex session parked there was re-selected
+by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
+stage.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_POLICIES = [
+    pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
+    pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
+]
+
+
+class _StubRequest:
+    """Minimal ``Request`` stub with the surface the sweep exercises.
+
+    Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
+    relies on so the same stub works under both scheduling policies.
+    """
+
+    def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
+        self.request_id = request_id
+        self.status = status
+        self.priority = 0
+        self.arrival_time = arrival_time
+
+    def __lt__(self, other: _StubRequest) -> bool:
+        return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
+
+
+def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
+    """Bare scheduler with just the surface the sweep reads."""
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(policy)
+    for req in waiting:
+        scheduler.waiting.add_request(req)
+    scheduler.skipped_waiting = create_request_queue(policy)
+    for req in skipped_waiting:
+        scheduler.skipped_waiting.add_request(req)
+    scheduler.running = list(running)
+    return scheduler
+
+
+@pytest.mark.parametrize("policy", _POLICIES)
+def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
+    monkeypatch: pytest.MonkeyPatch,
+    policy: SchedulingPolicy,
+) -> None:
+    """The public schedule path must remove aborted skipped requests first."""
+    aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
+    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.num_waiting_for_streaming_input = 0
+    scheduler._process_pending_omni_inputs = lambda model_mode: None
+    scheduler._resync_streaming_input_counter = lambda: None
+    scheduler._should_defer_waiting_admission = lambda: False
+    scheduler._restore_omni_wait_queues = lambda: None
+    scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
+    scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
+    scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
+
+    sentinel = object()
+
+    def upstream_schedule(_self, _throttle_prefills: bool = False):
+        assert list(scheduler.skipped_waiting) == []
+        return sentinel
+
+    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
+
+    assert OmniARScheduler.schedule(scheduler) is sentinel

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -28,7 +28,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
     sched = OmniARScheduler.__new__(OmniARScheduler)
     sched._new_prompt_len_snapshot = {}
-    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
+    sched.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_id=stage_id, max_model_len=None),
+    )
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
@@ -210,6 +212,81 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.chunk_transfer_adapter.save_async.assert_not_called()
 
 
+def test_stale_async_frame_is_dropped_before_output_processing() -> None:
+    session = _make_request()
+    session.status = RequestStatus.RUNNING
+    session.num_in_flight_tokens = 2
+    session.num_computed_tokens = session.num_tokens + 1
+    session.num_output_placeholders = 1
+    session.async_tokens_to_discard = 1
+    session.sampling_params = SimpleNamespace(num_logprobs=1)
+    num_computed_tokens = session.num_computed_tokens
+    num_output_placeholders = session.num_output_placeholders
+
+    sched = MagicMock()
+    sched.requests = {session.request_id: session}
+    sched.perf_metrics = None
+    sched.structured_output_manager.should_advance.return_value = False
+
+    def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
+        request.async_tokens_to_discard = 0
+        return token_ids, False
+
+    sched._update_request_with_output.side_effect = discard_stale_output
+    sched._process_kv_transfer_trigger.return_value = False
+    sched.chunk_transfer_adapter = MagicMock()
+    sched.running = [session]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.multimodal_outputs = None
+    model_runner_output.inter_stage_outputs = [{"hidden": object()}]
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {session.request_id: 0}
+    model_runner_output.routed_experts = None
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert session.async_tokens_to_discard == 0
+    assert session.status == RequestStatus.RUNNING
+    assert session.num_computed_tokens == num_computed_tokens
+    assert session.num_output_placeholders == num_output_placeholders
+    sched.chunk_transfer_adapter.save_async.assert_not_called()
+
+    session.sampling_params = SamplingParams(max_tokens=8)
+    next_payload = {"hidden": object()}
+    sched._update_request_with_output.side_effect = None
+    sched._update_request_with_output.return_value = ([43], False)
+    model_runner_output.sampled_token_ids = [[43]]
+    model_runner_output.inter_stage_outputs = [next_payload]
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+    sched.chunk_transfer_adapter.save_async.assert_called_once_with(
+        next_payload,
+        session,
+        False,
+    )
+
+
 def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
     sched = _make_scheduler(stage_id=0)
     session = _make_request()
@@ -354,3 +431,34 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
     assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
     assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
     assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+
+
+def test_native_duplex_prompt_reset_prevents_stage1_context_overflow() -> None:
+    sched = _make_scheduler(stage_id=1)
+    sched.vllm_config.model_config.max_model_len = 4096
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        receives_chunks=False,
+        segment_finished_requests=set(),
+    )
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.resumable = True
+    session.prompt_token_ids = [0] * 3
+    session._all_token_ids.clear()
+    session._all_token_ids.extend([0] * 4095)
+    session._output_token_ids.extend([1] * 4080)
+    session.num_prompt_tokens = 3
+    session.num_computed_tokens = 4095
+
+    update = _make_update([0] * 16)
+    update.model_intermediate_buffer = {"native_duplex": True}
+
+    sched._update_request_as_session(session, update)
+
+    assert session.prompt_token_ids == [0] * 16
+    assert list(session._all_token_ids) == [0] * 16
+    assert session._output_token_ids == []
+    assert session.num_prompt_tokens == 16
+    assert session.num_computed_tokens == 0
+    assert session.model_intermediate_buffer == update.model_intermediate_buffer
+    assert session.status == RequestStatus.WAITING

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -76,6 +76,9 @@ def _make_scheduler(
         else:
             hf_config.voxcpm2_runtime_config = runtime_config
     sched.vllm_config = SimpleNamespace(model_config=model_config)
+    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
+    # abort sweep at the top of ``schedule()`` reads it.
+    sched.skipped_waiting = _MockQueue()
     return sched
 
 

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -0,0 +1,151 @@
+"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
+``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
+and in ``OmniARScheduler.schedule``.
+
+Upstream keeps that counter incrementally (``+1`` when a resumable request
+parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
+next update arrives or the request is finished) and
+``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
+queues. Omni rewrites ``request.status`` outside both hooks -- the
+chunk-transfer adapter's park/restore and
+``_realign_request_status_to_queues`` -- so a counted request can leave
+``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
+counter for the lifetime of the process.
+
+An inflated counter does not degrade throughput, it hangs the stage:
+``EngineCore.has_work()`` reads false while a live request sits in
+``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
+``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
+server answered the first conversation after a boot and was silent from the
+second on while still reporting healthy.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _StubRequest:
+    def __init__(self, request_id: str, status: RequestStatus) -> None:
+        self.request_id = request_id
+        self.status = status
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+
+_SCHEDULER_PARAMS = [
+    pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
+    pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
+]
+
+
+def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.requests = requests
+    scheduler.running = running
+    scheduler.waiting = waiting
+    scheduler.skipped_waiting = skipped if skipped is not None else []
+    scheduler.num_waiting_for_streaming_input = counter
+    return scheduler
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_clears_counter_left_by_a_stomped_status(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    """The abort case that wedged the talker.
+
+    The request was counted while parked for streaming input, then had its
+    status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
+    ``finish_requests`` took the running branch and never decremented. With
+    no parked requests left, the counter must land at 0.
+    """
+    aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={},  # upstream already dropped it
+        running=[],
+        waiting=[],
+        counter=1,
+    )
+
+    monkeypatch.setattr(
+        scheduler_mod.VLLMScheduler,
+        "finish_requests",
+        lambda self, request_ids, finished_status: [aborted],
+    )
+
+    scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
+
+    assert scheduler.num_waiting_for_streaming_input == 0
+
+
+def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
+    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
+    so the resync has to look at both or it under-counts and reintroduces the
+    spin the counter exists to prevent."""
+    parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
+    scheduler = _make_scheduler(
+        OmniARScheduler,
+        requests={parked.request_id: parked},
+        running=[],
+        waiting=[],
+        counter=0,
+        skipped=[parked],
+    )
+
+    scheduler._resync_streaming_input_counter()
+
+    assert scheduler.num_waiting_for_streaming_input == 1
+
+
+def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Pins placement in ``schedule``: the chunk-transfer adapter rewrites
+    status during ``process_pending_chunks``, so the resync has to run after
+    that and before upstream schedules, while ``self.waiting`` is still the
+    real queue."""
+    parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.requests = {parked.request_id: parked}
+    scheduler.running = [parked]
+    scheduler.waiting = []
+    scheduler.skipped_waiting = []
+    scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
+    scheduler.policy = "fcfs"
+
+    seen: dict[str, int] = {}
+
+    monkeypatch.setattr(
+        ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
+    )
+    monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
+    monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
+
+    def fake_schedule(self, throttle_prefills=False):
+        seen["counter"] = self.num_waiting_for_streaming_input
+        raise _StopScheduleError
+
+    class _StopScheduleError(Exception):
+        pass
+
+    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
+
+    with pytest.raises(_StopScheduleError):
+        OmniARScheduler.schedule(scheduler)
+
+    assert seen["counter"] == 0

tests/diffusion/attention/test_attention_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """Tests for per-role attention backend configuration (RFC: per-role-attention-backend).
 
@@ -90,14 +90,6 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match=match):
             AttentionSpec(**spec)
 
-    def test_fastvideo_vsa_topk_serialized(self):
-        spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=96)
-        assert spec.backend_kwargs() == {"topk": 96}
-
-    def test_fastvideo_vsa_topk_rejected_for_other_backend(self):
-        with pytest.raises(ValueError, match="only supported by the FASTVIDEO_VSA"):
-            AttentionSpec(backend="TORCH_SDPA", fastvideo_vsa_topk=96)
-
     def test_block_sparse_defaults_applied_when_backend_selected(self):
         spec = AttentionSpec(backend="RAINFUSION_ATTN")
         assert spec.block_sparse.sparsity == 0.8

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -1,203 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import sys
-import types
-from typing import Any
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
-    FastVideoVSAImpl,
-)
-from vllm_omni.diffusion.attention.backends.registry import (
-    DiffusionAttentionBackendEnum,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def test_fastvideo_vsa_backend_is_registered():
-    assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
-
-
-def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
-    calls: dict[str, Any] = {}
-    fake_module = types.ModuleType("fastvideo_kernel")
-
-    def fake_video_sparse_attn_bshd(
-        q,
-        k,
-        v,
-        variable_block_sizes,
-        q_variable_block_sizes,
-        compress_attn_weight,
-        topk,
-        block_size,
-    ):
-        calls["q_shape"] = tuple(q.shape)
-        calls["vbs"] = variable_block_sizes.detach().cpu().tolist()
-        calls["q_vbs"] = q_variable_block_sizes.detach().cpu().tolist()
-        calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
-        calls["compress_shape"] = tuple(compress_attn_weight.shape)
-        calls["topk"] = topk
-        calls["block_size"] = block_size
-        return q + k + v
-
-    setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
-    monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
-
-    impl = FastVideoVSAImpl(
-        num_heads=2,
-        head_size=8,
-        softmax_scale=8**-0.5,
-        causal=False,
-        backend_kwargs={
-            "topk": 1,
-            "block_size": (4, 8, 8),
-            "min_seq_len": 1,
-            "disable_when_sp_active": False,
-        },
-    )
-    query = torch.randn(1, 300, 2, 8)
-    key = torch.randn_like(query)
-    value = torch.randn_like(query)
-    attn_metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10)})
-    monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
-
-    output = impl.forward_cuda(query, key, value, attn_metadata)
-
-    assert output.shape == query.shape
-    assert calls["q_shape"] == (1, 1024, 2, 8)
-    assert calls["vbs"] == [192, 48, 48, 12]
-    assert calls["q_vbs"] == [192, 48, 48, 12]
-    assert calls["compress_shape"] == (1, 1024, 2, 8)
-    assert calls["compress_sum"] == 0.0
-    assert calls["topk"] == 1
-    assert calls["block_size"] == (4, 8, 8)
-
-
-def test_fastvideo_vsa_uses_learned_gate_when_provided(monkeypatch):
-    calls: dict[str, Any] = {}
-    fake_module = types.ModuleType("fastvideo_kernel")
-
-    def fake_video_sparse_attn_bshd(
-        q,
-        k,
-        v,
-        variable_block_sizes,
-        q_variable_block_sizes,
-        compress_attn_weight,
-        topk,
-        block_size,
-    ):
-        calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
-        calls["compress_shape"] = tuple(compress_attn_weight.shape)
-        return q + k + v
-
-    setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
-    monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
-
-    impl = FastVideoVSAImpl(
-        num_heads=2,
-        head_size=8,
-        softmax_scale=8**-0.5,
-        causal=False,
-        backend_kwargs={
-            "topk": 1,
-            "block_size": (4, 8, 8),
-            "min_seq_len": 1,
-            "disable_when_sp_active": False,
-        },
-    )
-    query = torch.randn(1, 300, 2, 8)
-    gate = torch.ones_like(query)
-    metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10), "gate_compress": gate})
-    monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
-
-    output = impl.forward_cuda(query, query, query, metadata)
-
-    assert output.shape == query.shape
-    assert calls["compress_shape"] == (1, 1024, 2, 8)
-    # Only the 300 real tokens carry ones; padded gate slots stay zero.
-    assert calls["compress_sum"] == gate.numel()
-
-
-def test_fastvideo_vsa_falls_back_without_dit_shape(monkeypatch):
-    calls: dict[str, Any] = {}
-
-    impl = FastVideoVSAImpl(
-        num_heads=2,
-        head_size=8,
-        softmax_scale=8**-0.5,
-        causal=False,
-        backend_kwargs={
-            "topk": 1,
-            "block_size": (4, 8, 8),
-            "min_seq_len": 1,
-            "disable_when_sp_active": False,
-        },
-    )
-
-    def fake_fallback(query, key, value, attn_metadata, reason):
-        calls["reason"] = reason
-        return torch.zeros_like(query)
-
-    monkeypatch.setattr(impl, "_fallback", fake_fallback)
-
-    query = torch.randn(1, 512, 2, 8)
-    output = impl.forward_cuda(query, query, query, AttentionMetadata())
-
-    assert output.shape == query.shape
-    assert calls["reason"] == "vsa_dit_seq_shape metadata is required"
-
-
-def test_fastvideo_vsa_falls_back_for_mask(monkeypatch):
-    calls: dict[str, Any] = {}
-
-    impl = FastVideoVSAImpl(
-        num_heads=2,
-        head_size=8,
-        softmax_scale=8**-0.5,
-        causal=False,
-        backend_kwargs={"min_seq_len": 1},
-    )
-
-    def fake_fallback(query, key, value, attn_metadata, reason):
-        calls["reason"] = reason
-        return torch.zeros_like(query)
-
-    monkeypatch.setattr(impl, "_fallback", fake_fallback)
-
-    query = torch.randn(1, 512, 2, 8)
-    mask = torch.ones(1, 512, dtype=torch.bool)
-    output = impl.forward_cuda(query, query, query, AttentionMetadata(attn_mask=mask))
-
-    assert output.shape == query.shape
-    assert calls["reason"]
-
-
-def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
-    query = torch.randn(1, 512, 2, 8)
-    metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (4, 8, 16)})
-
-    all_blocks = FastVideoVSAImpl(
-        num_heads=2,
-        head_size=8,
-        softmax_scale=8**-0.5,
-        causal=False,
-        backend_kwargs={"topk": 2, "block_size": (4, 8, 8), "min_seq_len": 1},
-    )
-    too_many = FastVideoVSAImpl(
-        num_heads=2,
-        head_size=8,
-        softmax_scale=8**-0.5,
-        causal=False,
-        backend_kwargs={"topk": 3, "block_size": (4, 8, 8), "min_seq_len": 1},
-    )
-
-    # CPU/float32 is rejected later, but k=N itself must not trigger fallback.
-    assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
-    assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"

tests/diffusion/lora/test_lora_manager.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from __future__ import annotations
 
@@ -20,10 +20,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummyLoRALayer:
-    def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1):
+    def __init__(self, n_slices: int, output_slices: tuple[int, ...]):
         self.n_slices = n_slices
         self.output_slices = output_slices
-        self.tp_size = tp_size
         self.set_calls: list[
             tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
         ] = []
@@ -237,14 +236,6 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
 
 def test_lora_manager_activates_packed_lora_from_sublayers():
     pipeline = torch.nn.Module()
-    bound_names = None
-
-    def validate_binding(*, lora_model, bound_lora_names):
-        nonlocal bound_names
-        assert lora_model.id == 1
-        bound_names = bound_lora_names
-
-    pipeline._validate_diffusion_lora_binding = validate_binding
     pipeline.stacked_params_mapping = [
         (".to_qkv", ".to_q", "q"),
         (".to_qkv", ".to_k", "k"),
@@ -288,70 +279,6 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
     assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0)
     assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0)
     assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0)
-    assert bound_names == frozenset(loras)
-
-
-def test_lora_manager_rolls_back_all_layers_when_activation_fails():
-    class _FailingLoRALayer(_DummyLoRALayer):
-        fail = False
-
-        def set_lora(self, index: int, lora_a, lora_b):
-            if self.fail:
-                raise RuntimeError("bind failed")
-            super().set_lora(index, lora_a, lora_b)
-
-    manager = DiffusionLoRAManager(
-        pipeline=torch.nn.Module(),
-        device=torch.device("cpu"),
-        dtype=torch.bfloat16,
-        max_cached_adapters=2,
-    )
-    first = _DummyLoRALayer(n_slices=1, output_slices=(2,))
-    second = _FailingLoRALayer(n_slices=1, output_slices=(2,))
-    manager._lora_modules = {
-        "transformer.first": first,
-        "transformer.second": second,
-    }
-
-    def adapter(adapter_id: int):
-        loras = {
-            name: LoRALayerWeights(
-                module_name=name,
-                rank=2,
-                lora_alpha=2,
-                lora_a=torch.full((2, 2), float(adapter_id)),
-                lora_b=torch.full((2, 2), float(adapter_id)),
-            )
-            for name in manager._lora_modules
-        }
-        return type(
-            "LM",
-            (),
-            {
-                "id": adapter_id,
-                "loras": loras,
-                "get_lora": lambda self, key: self.loras.get(key),
-            },
-        )()
-
-    manager._registered_adapters = {1: adapter(1), 2: adapter(2)}
-    manager._activate_adapter(1, scale=1.0)
-    first_calls_after_success = len(first.set_calls)
-
-    second.fail = True
-    with pytest.raises(RuntimeError, match="bind failed"):
-        manager._activate_adapter(2, scale=1.0)
-
-    assert manager._active_adapter_id is None
-    assert first.reset_calls == 1
-    assert second.reset_calls == 1
-
-    # The failed activation must not leave the old adapter eligible for the
-    # fast path: reactivating it binds every layer again.
-    second.fail = False
-    manager._activate_adapter(1, scale=1.0)
-    assert manager._active_adapter_id == 1
-    assert len(first.set_calls) == first_calls_after_success + 2
 
 
 def _dummy_lora_request(adapter_id: int) -> LoRARequest:

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
 import sys
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
-from typing import Any, TypeVar
 from unittest.mock import Mock, patch
 
 import numpy as np
@@ -15,14 +14,6 @@ import torch.nn as nn
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
-_ItemT = TypeVar("_ItemT")
-_ResultT = TypeVar("_ResultT")
-
-
-def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ResultT:
-    items.append(item)
-    return result
-
 
 def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -194,7 +185,6 @@ def test_startup_task_rejects_unsupported_value():
 
 def test_combined_task_inference_and_transformer_routing():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.errors import OmniClientError
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
@@ -206,9 +196,6 @@ def test_combined_task_inference_and_transformer_routing():
     assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
     assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
     assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
-    assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va"
-    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
-        pipeline._resolve_task("ref2va", {}, has_turbo_lora=True)
     pipeline.partition = "ref2va"
     pipeline.supported_tasks = frozenset({"ref2va"})
     assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
@@ -273,7 +260,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
         ):
             (tmp_path / partition_name / component).mkdir()
 
-    created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
+    created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
 
     class FakeModule(torch.nn.Module):
         def __init__(self, *args, **kwargs):
@@ -291,15 +278,15 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
 
         return create
 
-    tokenizer_calls: list[Any] = []
-    processor_calls: list[Any] = []
+    tokenizer_calls = []
+    processor_calls = []
     tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast)
-    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
-        tokenizer_calls, (args, kwargs), object()
+    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
+        tokenizer_calls.append((args, kwargs)) or object()
     )
     processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor)
-    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
-        processor_calls, (args, kwargs), object()
+    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
+        processor_calls.append((args, kwargs)) or object()
     )
     monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT)
     monkeypatch.setattr(
@@ -639,7 +626,7 @@ def test_partially_constructed_pipeline_falls_back_to_the_uniform_schedule():
 
 
 def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
-    diffuse_calls: list[dict[str, Any]] = []
+    diffuse_calls = []
     base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0]
     pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": base_schedule, "ref2va": None})
 
@@ -656,7 +643,7 @@ def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
 
 
 def test_distilled_forward_accepts_the_matching_explicit_step_count():
-    diffuse_calls: list[dict[str, Any]] = []
+    diffuse_calls = []
     pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None})
 
     pipeline.forward(_t2va_batch(num_inference_steps=4))
@@ -726,10 +713,8 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
     )
 
     assert output.shape == q.shape
-    metadata = attention.attention.metadata
-    assert metadata is not None
-    assert metadata.attn_mask is None
-    assert metadata.extra["valid_kv_length"] == 5
+    assert attention.attention.metadata.attn_mask is None
+    assert attention.attention.metadata.extra["valid_kv_length"] == 5
 
 
 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
@@ -769,11 +754,9 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
         packed_total=8,
     )
 
-    metadata = attention.attention.metadata
-    assert metadata is not None
-    assert metadata.attn_mask is None
-    assert metadata.extra["valid_kv_length"] == 5
-    assert metadata.extra["npu_attn_varlen"] is True
+    assert attention.attention.metadata.attn_mask is None
+    assert attention.attention.metadata.extra["valid_kv_length"] == 5
+    assert attention.attention.metadata.extra["npu_attn_varlen"] is True
 
 
 def test_packed_attention_keeps_padding_mask_for_other_backends():
@@ -813,10 +796,8 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
         packed_total=8,
     )
 
-    metadata = attention.attention.metadata
-    assert metadata is not None
     assert torch.equal(
-        metadata.attn_mask,
+        attention.attention.metadata.attn_mask,
         torch.tensor([[True, True, True, True, True, False, False, False]]),
     )
 
@@ -1503,12 +1484,12 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
         "audio_codecs": (),
         "file_size": 1024,
     }
-    transcode_calls: list[Any] = []
+    transcode_calls = []
     monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata)
     monkeypatch.setattr(
         reference_video_module,
         "_transcode_reference_video",
-        lambda source, **kwargs: _append_and_return(transcode_calls, (source, kwargs), "prepared.mp4"),
+        lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4",
     )
 
     prepared = reference_video_module.prepare_reference_videos(
@@ -1529,7 +1510,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
     decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3)
-    decode_calls: list[Any] = []
+    decode_calls = []
     monkeypatch.setattr(
         reference_video_module,
         "_probe_video",
@@ -1542,7 +1523,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
     monkeypatch.setattr(
         reference_video_module,
         "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), decoded),
+        lambda path, **kwargs: decode_calls.append((path, kwargs)) or decoded,
     )
 
     sampled = reference_video_module.sample_reference_video_frames("prepared.mp4")
@@ -1603,7 +1584,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
     expected = np.zeros((2, 2, 3, 3), dtype=np.uint8)
-    decode_calls: list[Any] = []
+    decode_calls = []
     monkeypatch.setitem(sys.modules, "decord", None)
     monkeypatch.setattr(
         reference_video_module,
@@ -1617,7 +1598,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
     monkeypatch.setattr(
         reference_video_module,
         "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), expected),
+        lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected,
     )
 
     actual = reference_video_module.load_video_frames("prepared.mp4")
@@ -1657,7 +1638,7 @@ def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message):
 def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
-    calls: list[Any] = []
+    calls = []
     probe = {
         "streams": [
             {
@@ -1675,11 +1656,7 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
     monkeypatch.setattr(
         reference_video_module.subprocess,
         "run",
-        lambda command, **kwargs: _append_and_return(
-            calls,
-            (command, kwargs),
-            SimpleNamespace(stdout=json.dumps(probe)),
-        ),
+        lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)),
     )
 
     metadata = reference_video_module._probe_video("prepared.mp4")
@@ -1692,7 +1669,7 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
 def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
-    calls: list[Any] = []
+    calls = []
     stream = {
         "codec_type": "video",
         "width": 4,
@@ -1712,11 +1689,8 @@ def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
     monkeypatch.setattr(
         reference_video_module.subprocess,
         "run",
-        lambda command, **kwargs: _append_and_return(
-            calls,
-            (command, kwargs),
-            SimpleNamespace(stdout=json.dumps(first_probe if not calls else second_probe)),
-        ),
+        lambda command, **kwargs: calls.append((command, kwargs))
+        or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)),
     )
 
     metadata = reference_video_module._probe_video("prepared.mp4")
@@ -1776,16 +1750,12 @@ def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_p
             "file_size": 1024,
         },
     }
-    transcode_calls: list[Any] = []
+    transcode_calls = []
     monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
     monkeypatch.setattr(
         reference_video_module,
         "_transcode_reference_video",
-        lambda source, **kwargs: _append_and_return(
-            transcode_calls,
-            (source, kwargs),
-            f"{source}.prepared.mp4",
-        ),
+        lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4",
     )
 
     prepared = reference_video_module.prepare_reference_videos(

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -1,381 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from safetensors.torch import save_file
-from vllm.lora.lora_weights import PackedLoRALayerWeights
-
-from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
-from vllm_omni.diffusion.models.minimax_h3 import lora as lora_module
-from vllm_omni.diffusion.models.minimax_h3.lora import (
-    load_minimax_h3_turbo_lora,
-)
-from vllm_omni.errors import OmniClientError
-from vllm_omni.lora.request import LoRARequest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-_TINY_TARGET_DIMS = {
-    "attn.to_q": (2, 3),
-    "attn.to_k": (2, 3),
-    "attn.to_v": (2, 3),
-    "attn.to_out.0": (3, 2),
-    "ff.net.0.proj": (2, 4),
-    "ff.net.2": (2, 2),
-}
-
-
-@pytest.fixture(autouse=True)
-def _use_tiny_h3_dimensions(monkeypatch):
-    monkeypatch.setattr(lora_module, "_TURBO_TARGET_DIMS", _TINY_TARGET_DIMS)
-
-
-def _request(path) -> LoRARequest:
-    return LoRARequest(
-        lora_name="turbo",
-        lora_int_id=1,
-        lora_path=str(path),
-    )
-
-
-def _write_tiny_turbo(
-    path,
-    *,
-    alpha: str = "128",
-    key_format: str = "minimax-h3-diffusers",
-    omit_target: str | None = None,
-    shape_overrides: dict[str, tuple[int, int]] | None = None,
-) -> None:
-    rank = 128
-    tensors = {}
-    overrides = shape_overrides or {}
-    target_suffixes = (
-        "attn.to_q",
-        "attn.to_k",
-        "attn.to_v",
-        "attn.to_out.0",
-        "ff.net.0.proj",
-        "ff.net.2",
-    )
-    for prefix, block_count in (
-        ("transformer_blocks", 50),
-        ("token_refiner.refiner_blocks", 2),
-    ):
-        for block_index in range(block_count):
-            for suffix in target_suffixes:
-                target = f"{prefix}.{block_index}.{suffix}"
-                if target == omit_target:
-                    continue
-                input_dim, output_dim = _TINY_TARGET_DIMS[suffix]
-                a_name = f"{target}.lora_A.default.weight"
-                b_name = f"{target}.lora_B.default.weight"
-                tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim)))
-                if suffix == "ff.net.0.proj":
-                    if b_name in overrides:
-                        tensors[b_name] = torch.ones(overrides[b_name])
-                    else:
-                        tensors[b_name] = torch.cat(
-                            (
-                                torch.ones(output_dim // 2, rank),
-                                torch.full((output_dim // 2, rank), 2.0),
-                            ),
-                            dim=0,
-                        )
-                else:
-                    tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
-    save_file(
-        tensors,
-        str(path),
-        metadata={"alpha": alpha, "key_format": key_format},
-    )
-
-
-def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path):
-    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    _write_tiny_turbo(path)
-
-    loaded = load_minimax_h3_turbo_lora(
-        partition="fl2va",
-        lora_request=_request(path),
-        lora_path=path,
-        dtype=torch.float32,
-    )
-
-    assert loaded is not None
-    lora_model, peft_helper = loaded
-    assert peft_helper.r == 128
-    assert peft_helper.lora_alpha == 128
-    assert len(lora_model.loras) == 312
-    assert "blocks.0.attn.to_q" in lora_model.loras
-    assert "blocks.0.mlp.fc1" in lora_model.loras
-    fc1 = lora_model.get_lora("blocks.0.mlp.fc1")
-    assert isinstance(fc1, PackedLoRALayerWeights)
-    assert len(fc1.lora_a) == 2
-    assert len(fc1.lora_b) == 2
-    torch.testing.assert_close(fc1.lora_a[0], torch.ones(128, 2))
-    torch.testing.assert_close(fc1.lora_a[1], torch.ones(128, 2))
-    torch.testing.assert_close(fc1.lora_b[0], torch.full((2, 128), 2.0))
-    torch.testing.assert_close(fc1.lora_b[1], torch.ones(2, 128))
-
-
-@pytest.mark.parametrize(
-    ("tensor_suffix", "bad_shape"),
-    [
-        ("attn.to_q.lora_A.default.weight", (128, 1)),
-        ("attn.to_q.lora_A.default.weight", (128, 3)),
-        ("attn.to_q.lora_B.default.weight", (2, 128)),
-        ("attn.to_q.lora_B.default.weight", (4, 128)),
-    ],
-)
-def test_h3_turbo_rejects_under_and_oversized_global_shapes(tmp_path, tensor_suffix, bad_shape):
-    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    tensor_name = f"transformer_blocks.0.{tensor_suffix}"
-    _write_tiny_turbo(path, shape_overrides={tensor_name: bad_shape})
-
-    with pytest.raises(ValueError, match="invalid global shape"):
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(path),
-            lora_path=path,
-            dtype=torch.float32,
-        )
-
-
-def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(tmp_path):
-    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    _write_tiny_turbo(path)
-
-    class _Pipeline:
-        def _load_diffusion_lora_adapter(self, **kwargs):
-            return load_minimax_h3_turbo_lora(partition="fl2va", **kwargs)
-
-    manager = object.__new__(DiffusionLoRAManager)
-    manager.pipeline = _Pipeline()
-    manager.dtype = torch.float32
-    manager._expected_lora_modules = {"to_q", "fc1"}
-
-    lora_model, peft_helper = manager._load_adapter(_request(path))
-
-    assert lora_model.id == 1
-    assert peft_helper.lora_alpha == 128
-    assert len(lora_model.loras) == 312
-    assert "blocks.0.attn.to_q" in lora_model.loras
-    assert "blocks.0.mlp.fc1" in lora_model.loras
-
-
-def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
-    wrong_alpha = tmp_path / "wrong_alpha" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    wrong_alpha.parent.mkdir()
-    _write_tiny_turbo(wrong_alpha, alpha="8")
-    with pytest.raises(ValueError, match="requires alpha=128"):
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(wrong_alpha),
-            lora_path=wrong_alpha,
-            dtype=torch.float32,
-        )
-
-    valid = tmp_path / "valid" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    valid.parent.mkdir()
-    _write_tiny_turbo(valid)
-    with pytest.raises(ValueError, match="supports FL2VA/T2VA only"):
-        load_minimax_h3_turbo_lora(
-            partition="ref2va",
-            lora_request=_request(valid),
-            lora_path=valid,
-            dtype=torch.float32,
-        )
-
-
-@pytest.mark.parametrize(
-    "offload_mode",
-    [
-        "model-level CPU offload (--enable-cpu-offload)",
-        "layerwise offload (--enable-layerwise-offload)",
-        "distributed layerwise offload (--enable-distributed-layerwise-offload)",
-    ],
-)
-def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
-    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    _write_tiny_turbo(path)
-
-    with pytest.raises(ValueError, match="does not support"):
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(path),
-            lora_path=path,
-            dtype=torch.float32,
-            unsupported_offload_mode=offload_mode,
-        )
-
-
-def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
-    unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
-    _write_tiny_turbo(unsupported)
-
-    with pytest.raises(ValueError, match="supports only"):
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(unsupported),
-            lora_path=unsupported,
-            dtype=torch.float32,
-        )
-
-    supported = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    _write_tiny_turbo(supported)
-    # Directory resolution selects the declared artifact even when another
-    # v1.0 checkpoint is present beside it.
-    assert (
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(tmp_path),
-            lora_path=tmp_path,
-            dtype=torch.float32,
-        )
-        is not None
-    )
-
-
-def test_h3_turbo_rejects_a_truncated_declared_artifact(tmp_path):
-    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    _write_tiny_turbo(path, omit_target="transformer_blocks.49.ff.net.2")
-
-    with pytest.raises(ValueError, match="target set does not match"):
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(path),
-            lora_path=path,
-            dtype=torch.float32,
-        )
-
-
-def test_non_h3_checkpoint_falls_back_to_the_generic_peft_loader(tmp_path):
-    path = tmp_path / "other.safetensors"
-    _write_tiny_turbo(path, key_format="other")
-
-    assert (
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(path),
-            lora_path=path,
-            dtype=torch.float32,
-        )
-        is None
-    )
-
-    declared = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    _write_tiny_turbo(declared, key_format="other")
-    with pytest.raises(ValueError, match="requires safetensors metadata"):
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(declared),
-            lora_path=declared,
-            dtype=torch.float32,
-        )
-
-
-def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.diffusion.models.minimax_h3 import (
-        pipeline_minimax_h3 as pipeline_module,
-    )
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.partition = "combined"
-    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
-    pipeline._turbo_lora_adapter_ids = set()
-    request = _request("generic-peft")
-
-    def load():
-        return pipeline._load_diffusion_lora_adapter(
-            lora_request=request,
-            lora_path=request.lora_path,
-            dtype=torch.float32,
-        )
-
-    def resolve(scale):
-        sampling = SimpleNamespace(lora_request=request, lora_scale=scale)
-        return pipeline._resolve_task(
-            "ref2va",
-            {},
-            has_turbo_lora=pipeline._has_active_turbo_lora(sampling),
-        )
-
-    recognized = (object(), object())
-    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: recognized)
-    assert load() is recognized
-    assert resolve(0.0) == "ref2va"
-    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
-        resolve(1.0)
-
-    # Simulate manager eviction followed by a generic PEFT adapter reusing the
-    # same client-supplied ID. A real reload must replace the stale kind.
-    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
-    assert load() is None
-    assert resolve(1.0) == "ref2va"
-
-
-def test_h3_turbo_requires_all_loaded_targets_to_bind():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline._turbo_lora_adapter_ids = {1}
-    lora_model = SimpleNamespace(id=1, loras={"to_q": object(), "to_k": object()})
-
-    pipeline._validate_diffusion_lora_binding(
-        lora_model=lora_model,
-        bound_lora_names=frozenset(lora_model.loras),
-    )
-    with pytest.raises(ValueError, match="bound=1/2"):
-        pipeline._validate_diffusion_lora_binding(
-            lora_model=lora_model,
-            bound_lora_names=frozenset({"to_q"}),
-        )
-
-
-@pytest.mark.parametrize(
-    ("num_inference_steps", "extra_args", "error"),
-    [
-        (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
-        (4, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
-        ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
-        (5, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"),
-        (5, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"),
-        (5, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"),
-        (5, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"),
-    ],
-)
-def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, error):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    pipeline.default_video_shift = 12.0
-    pipeline.default_audio_shift = 3.0
-    sampling = SimpleNamespace(
-        num_inference_steps=num_inference_steps,
-        extra_args=extra_args,
-    )
-
-    with pytest.raises(OmniClientError, match=error):
-        pipeline._validate_turbo_sampling(sampling)
-
-
-def test_h3_turbo_accepts_five_sigma_points_for_four_nfe():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    pipeline.default_video_shift = 12.0
-    pipeline.default_audio_shift = 3.0
-    pipeline._validate_turbo_sampling(
-        SimpleNamespace(
-            num_inference_steps=5,
-            extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0},
-        )
-    )

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import importlib
 from contextlib import contextmanager
 from types import SimpleNamespace
 
@@ -10,7 +9,6 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
-from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -84,7 +82,6 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline.vae_scale_factor_spatial = 8
     pipeline.boundary_ratio = 0.875
     pipeline.expand_timesteps = False
-    pipeline.is_dmd = False
     pipeline._guidance_scale = None
     pipeline._guidance_scale_2 = None
     pipeline._num_timesteps = None
@@ -97,7 +94,7 @@ def _make_pipeline() -> Wan22Pipeline:
 
 
 def _make_sampling(**overrides):
-    values: dict[str, object] = {
+    values = {
         "height": None,
         "width": None,
         "num_frames": 1,
@@ -341,125 +338,3 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
         (3.0, 3, 28.0, False),
     ]
     assert torch.equal(result, torch.full_like(latents, 10.0))
-
-
-class _StubDMDScheduler:
-    def __init__(self) -> None:
-        self.predict_clean_calls: list[tuple[float, float, float]] = []
-        self.add_noise_calls: list[tuple[float, float, float]] = []
-
-    def predict_clean(self, model_output, sample, timestep):
-        self.predict_clean_calls.append((float(model_output.mean()), float(sample.mean()), float(timestep)))
-        return sample - model_output
-
-    def add_noise(self, clean_sample, noise, timestep):
-        self.add_noise_calls.append((float(clean_sample.mean()), float(noise.mean()), float(timestep)))
-        return clean_sample + 10.0
-
-
-def test_diffuse_dmd_predicts_clean_and_renoises_between_steps(monkeypatch) -> None:
-    pipeline = _make_pipeline()
-    pipeline.is_dmd = True
-    pipeline.scheduler = _StubDMDScheduler()
-    latents = torch.zeros((1, 1, 1, 1, 1), dtype=torch.float32)
-    timesteps = torch.tensor([1000.0, 757.0, 522.0])
-
-    pipeline.predict_noise_maybe_with_cfg = lambda **kwargs: torch.ones_like(latents)  # type: ignore[method-assign]
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.randn_tensor",
-        lambda *args, **kwargs: torch.full(args[0], 2.0, dtype=kwargs["dtype"]),
-    )
-
-    result = pipeline.diffuse(
-        latents=latents,
-        timesteps=timesteps,
-        prompt_embeds=torch.zeros(1, 8),
-        negative_prompt_embeds=None,
-        guidance_low=1.0,
-        guidance_high=1.0,
-        boundary_timestep=None,
-        dtype=torch.float32,
-        attention_kwargs={},
-        generator=torch.Generator(device="cpu").manual_seed(1),
-    )
-
-    assert pipeline.scheduler.predict_clean_calls == [
-        (1.0, 0.0, 1000.0),
-        (1.0, 9.0, 757.0),
-        (1.0, 18.0, 522.0),
-    ]
-    assert pipeline.scheduler.add_noise_calls == [
-        (-1.0, 2.0, 757.0),
-        (8.0, 2.0, 522.0),
-    ]
-    torch.testing.assert_close(result, torch.tensor([[[[[17.0]]]]]))
-
-
-def _make_gate_loading_pipeline():
-    pipeline = Wan22Pipeline.__new__(Wan22Pipeline)
-    nn.Module.__init__(pipeline)
-    gate = WanSelfAttention.__new__(WanSelfAttention)
-    nn.Module.__init__(gate)
-    gate.to_gate_compress = nn.Linear(1, 1)
-    pipeline.gate_holder = gate
-    return pipeline, gate
-
-
-@pytest.mark.parametrize(
-    ("module_name", "class_name"),
-    [
-        ("pipeline_wan2_2", "Wan22Pipeline"),
-        ("pipeline_wan2_2_i2v", "Wan22I2VPipeline"),
-        ("pipeline_wan2_2_s2v", "Wan22S2VPipeline"),
-        ("pipeline_wan2_2_vace", "Wan22VACEPipeline"),
-    ],
-)
-def test_wan_pipeline_loaders_share_optional_gate_cleanup(monkeypatch, module_name, class_name) -> None:
-    module = importlib.import_module(f"vllm_omni.diffusion.models.wan2_2.{module_name}")
-    pipeline_cls = getattr(module, class_name)
-    pipeline = pipeline_cls.__new__(pipeline_cls)
-    expected = {"loaded"}
-
-    def fake_loader(model, weights):
-        assert model is pipeline
-        assert list(weights) == [("weight", torch.ones(1))]
-        return expected
-
-    monkeypatch.setattr(module, "load_wan_weights_with_optional_gate", fake_loader)
-
-    assert pipeline_cls.load_weights(pipeline, iter((("weight", torch.ones(1)),))) is expected
-
-
-def test_load_weights_removes_unloaded_vsa_gate(monkeypatch) -> None:
-    pipeline, gate = _make_gate_loading_pipeline()
-
-    class _Loader:
-        def __init__(self, model):
-            del model
-
-        def load_weights(self, weights):
-            return {name for name, _ in weights}
-
-    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
-    pipeline.load_weights(iter((("other.weight", torch.ones(1)),)))
-
-    assert pipeline.has_gate_compress_weights is False
-    assert gate.to_gate_compress is None
-
-
-def test_load_weights_keeps_trained_vsa_gate(monkeypatch) -> None:
-    pipeline, gate = _make_gate_loading_pipeline()
-    original_gate = gate.to_gate_compress
-
-    class _Loader:
-        def __init__(self, model):
-            del model
-
-        def load_weights(self, weights):
-            return {name for name, _ in weights}
-
-    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
-    pipeline.load_weights(iter((("gate_holder.to_gate_compress.weight", torch.ones(1)),)))
-
-    assert pipeline.has_gate_compress_weights is True
-    assert gate.to_gate_compress is original_gate

tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py

@@ -1,32 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def test_dmd_predict_clean_and_add_noise_match_flow_formulas() -> None:
-    scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
-    timestep = torch.tensor(757.0)
-    sigma = scheduler.sigma_for_timestep(timestep)
-    sample = torch.tensor([2.0, -1.0])
-    model_output = torch.tensor([0.5, 0.25])
-    noise = torch.tensor([-0.5, 1.5])
-
-    clean = scheduler.predict_clean(model_output, sample, timestep)
-    expected_clean = sample - sigma * model_output
-    torch.testing.assert_close(clean, expected_clean)
-
-    renoised = scheduler.add_noise(clean, noise, timestep)
-    expected_renoised = (1.0 - sigma) * expected_clean + sigma * noise
-    torch.testing.assert_close(renoised, expected_renoised)
-
-
-def test_dmd_sigma_requires_scalar_timestep() -> None:
-    scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
-    with pytest.raises(ValueError, match="scalar timestep"):
-        scheduler.sigma_for_timestep(torch.tensor([757.0, 522.0]))

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -546,7 +546,7 @@ def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
     assert second.codes is None
 
 
-def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
+def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
     request.resumable = True
@@ -565,6 +565,51 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt
     assert adapter.requests_num_chunks_sent["ext-stream"] == 0
 
 
+def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
+    request.resumable = True
+    request._omni_segment_generation = 0
+    adapter.requests_num_chunks_sent["ext-stream"] = 111
+
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+
+    late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream")
+    late.resumable = True
+    late.num_computed_tokens = 111
+    late._omni_segment_generation = 0
+    adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False)
+
+    next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream")
+    next_segment.resumable = True
+    next_segment.num_computed_tokens = 3
+    next_segment._omni_segment_generation = 1
+    adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False)
+
+    assert len(adapter._pending_save_reqs) == 2
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 3
+    assert adapter._pending_save_reqs[-1]["request"] is next_segment
+
+
+def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter):
+    adapter, connector = build_adapter(stage_id=1)
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct())
+    request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old")
+    request.resumable = True
+    request._omni_segment_generation = 0
+
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+
+    old_task = adapter._pending_save_reqs.popleft()
+    adapter._send_single_request(old_task)
+    connector.put.assert_called_once()
+
+    boundary_task = adapter._pending_save_reqs.popleft()
+    adapter._send_single_request(boundary_task)
+    assert connector.put.call_count == 2
+
+
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")

tests/engine/test_orchestrator.py

@@ -30,6 +30,7 @@ from vllm_omni.engine.messages import (
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
+    StreamingSegmentState,
     _build_terminal_empty_output,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -1951,7 +1952,7 @@ async def test_resumable_segment_boundary_builds_stage_metrics() -> None:
         final_stage_id=0,
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segment_finished = True
+    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
     req_state.stage_submit_ts[0] = time.time()
     orchestrator.request_states = {"req-stream": req_state}
     orchestrator.stage_pools = [pool]

tests/engine/test_orchestrator_segment_stage_keying.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Streaming segment-boundary state must be keyed by stage.
+
+``Orchestrator._orchestration_loop`` polls every stage into the same
+``OrchestratorRequestState``. While the segment-boundary fields lived in a single
+flat slot per request, whichever stage polled most recently overwrote the others,
+and the per-stage consumers then read back a different stage's boundary.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import (
+    Orchestrator,
+    OrchestratorRequestState,
+    StreamingSegmentState,
+)
+from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+THINKER_STAGE = 0
+TALKER_STAGE = 1
+
+
+@dataclass
+class _StageMetrics:
+    pipeline_timings: dict[str, float] = field(default_factory=dict)
+
+
+@dataclass
+class _StageOutput:
+    """Stand-in for a processed stage output, with the fields the router reads."""
+
+    request_id: str
+    finished: bool = False
+    error: object | None = None
+
+
+class _RecordingPool:
+    """Stage pool that records the requests it was asked to build metrics for."""
+
+    def __init__(self) -> None:
+        self.built_for: list[str] = []
+
+    def build_stage_metrics(self, outputs: list[_StageOutput], **_kwargs: object) -> _StageMetrics:
+        self.built_for.extend(output.request_id for output in outputs)
+        return _StageMetrics()
+
+
+def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState:
+    req_state = OrchestratorRequestState(
+        request_id=request_id,
+        sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)],
+        final_stage_id=TALKER_STAGE,
+    )
+    req_state.streaming.enabled = True
+    req_state.duplex_identity = DuplexRequestIdentity(
+        session_id="session-1",
+        fence=DuplexFence(session_id="session-1"),
+    )
+    return req_state
+
+
+def test_duplex_output_context_reads_only_its_own_stage_segment() -> None:
+    req_state = _duplex_req_state()
+    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(
+        finished=True,
+        token_ids=[11, 22],
+        output_metadata={"stage": "thinker"},
+    )
+    # The talker polls after the thinker and is still mid-segment.
+    req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False)
+
+    thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE)
+    talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE)
+
+    assert thinker is not None
+    assert talker is not None
+
+    # These are the assertions that fail on the flat field: the talker is written
+    # second, so the thinker loses its boundary to whichever stage polled last.
+    assert thinker.segment_finished is True
+    assert thinker.segment_token_ids == (11, 22)
+    assert thinker.segment_output_metadata == {"stage": "thinker"}
+
+    assert talker.segment_finished is False
+    assert talker.segment_token_ids == ()
+    assert talker.segment_output_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None:
+    req_state = _duplex_req_state()
+    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True)
+
+    talker_pool = _RecordingPool()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.request_states = {req_state.request_id: req_state}
+    orchestrator.stage_pools = [_RecordingPool(), talker_pool]
+
+    routed: list[_StageMetrics | None] = []
+
+    async def record_route(
+        _stage: int,
+        _replica: int,
+        _output: _StageOutput,
+        _state: OrchestratorRequestState,
+        metrics: _StageMetrics | None,
+    ) -> None:
+        routed.append(metrics)
+
+    orchestrator._route_output = record_route
+
+    # The talker's own output is neither finished nor at a segment boundary, so
+    # it must not be treated as one just because the thinker reached one.
+    await orchestrator._handle_processed_outputs(
+        TALKER_STAGE,
+        0,
+        [_StageOutput(request_id=req_state.request_id)],
+    )
+
+    assert talker_pool.built_for == []
+    assert routed == [None]

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
+    StreamingSegmentState,
     _OrchestratorDuplexStagePort,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -358,7 +359,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None:
         final_output_stage_ids={0},
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segment_finished = True
+    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
     output = SimpleNamespace(
         request_id=req_state.request_id,
         finished=True,

tests/entrypoints/openai/test_duplex_protocol.py

@@ -462,6 +462,16 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
     assert caps["stage_handoff_transport"] == "scheduler_data_plane"
 
 
+def test_duplex_overlap_policy_auto_remains_a_compatibility_value():
+    from vllm_omni.experimental.fullduplex.openai.protocol import (
+        DuplexOverlapPolicy,
+        DuplexSessionConfig,
+    )
+
+    assert DuplexOverlapPolicy.AUTO.value == "auto"
+    assert DuplexSessionConfig._normalize_overlap_policy("auto") == "auto"
+
+
 def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
     caps = DuplexCapabilities.minicpmo45_native(max_sessions=1).as_dict()
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from types import SimpleNamespace
 
@@ -400,9 +400,7 @@ def test_serve_cli_accepts_diffusion_attention_backend():
             "Qwen/Qwen-Image",
             "--omni",
             "--diffusion-attention-backend",
-            "FASTVIDEO_VSA",
-            "--fastvideo-vsa-topk",
-            "96",
+            "FLASH_ATTN",
         ]
     )
 
@@ -410,12 +408,10 @@ def test_serve_cli_accepts_diffusion_attention_backend():
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
     diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
 
-    assert args.diffusion_attention_backend == "FASTVIDEO_VSA"
-    assert args.fastvideo_vsa_topk == 96
+    assert args.diffusion_attention_backend == "FLASH_ATTN"
     assert isinstance(diffusion_attention_config, AttentionConfig)
     assert diffusion_attention_config.default is not None
-    assert diffusion_attention_config.default.backend == "FASTVIDEO_VSA"
-    assert diffusion_attention_config.default.backend_kwargs() == {"topk": 96}
+    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
 
 
 def test_serve_cli_accepts_request_batch_max_wait_ms():

vllm_omni/config/config_factory.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Config factories for vllm-omni, e.g., StageConfigFactory."""
 
 from __future__ import annotations
@@ -226,10 +226,7 @@ class StageConfigFactory:
                     # If we have a resolver, call it with the optional hf_config
                     # to get the default pipeline config for this key
                     pipeline_cfg = obj(hf_config) if callable(obj) else obj
-                    if pipeline_cfg is not None and class_name in (
-                        pipeline_cfg.diffusers_class_name,
-                        *pipeline_cfg.diffusers_class_aliases,
-                    ):
+                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
                         logger.info(
                             "Detected pipeline %r from model_index.json (_class_name=%r)",
                             pipeline_cfg.model_type,
@@ -436,7 +433,6 @@ class StageConfigFactory:
         load-balancer policy (``None`` when no strategy set one) travels with the
         stages instead of through a mutable out-param.
         """
-        deploy_cfg: DeployConfig | None
         if user_deploy_config is not None:
             deploy_cfg = user_deploy_config
         elif deploy_config_path is not None:
@@ -444,12 +440,9 @@ class StageConfigFactory:
             assert deploy_cfg is not None
         elif pipeline_cfg.default_deploy_config_name is not None:
             deploy_cfg = load_deploy_config(_DEPLOY_DIR / pipeline_cfg.default_deploy_config_name)
-            assert deploy_cfg is not None
         else:
             deploy_cfg = DeployConfig()
 
-        assert deploy_cfg is not None
-
         cli_async_chunk = cli_overrides.get("async_chunk")
         if cli_async_chunk is not None:
             deploy_cfg.async_chunk = bool(cli_async_chunk)

vllm_omni/config/omni_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Structured vLLM-Omni configuration classes.
 
 This module is additive for Phase 2 of RFC #4021.
@@ -126,7 +126,7 @@ def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]:
 
     cls.__init__ = keyword_only_init
     cls.__signature__ = keyword_only_signature
-    setattr(cls, "__match_args__", ())
+    cls.__match_args__ = ()
     return cls
 
 
@@ -447,7 +447,7 @@ class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig):
     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  # type: ignore[assignment]
+    is_encoder_decoder: InitVar[bool] = False
     enable_chunked_prefill: bool | None = None
     async_scheduling: bool | None = None
 
@@ -519,8 +519,8 @@ class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig):
             self.data_parallel_rank,
             self.data_parallel_rpc_port,
             self.worker_cls,
-            getattr(self, "all2all_backend", None),
-            getattr(self, "disable_custom_all_reduce", None),
+            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(
@@ -741,7 +741,7 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
 
         normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs)
-        valid_fields = {f.name for f in fields(cast(Any, cls))}
+        valid_fields = {f.name for f in fields(cls)}
         return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields})
 
     def __post_init__(self) -> None:
@@ -893,7 +893,7 @@ class _DiffusionConfigProjection:
                 setattr(self, name, _copy_value(getattr(omni_diffusion_config, name)))
 
 
-_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(cast(Any, _DiffusionConfigProjection)))
+_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(_DiffusionConfigProjection))
 
 # Current OmniDiffusionConfig still contains a flat mix of shared engine,
 # runtime, parallel, and diffusion-specific knobs. Keep this classification
@@ -956,7 +956,6 @@ _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS)
 _DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset(
     {
         "diffusion_attention_backend",
-        "fastvideo_vsa_topk",
         "kv_cache_dtype",
         "kv_cache_skip_layers",
         "kv_cache_skip_steps",
@@ -1141,7 +1140,7 @@ def _global_stage_cli_fields() -> frozenset[str]:
     )
     externally_consumed = (
         _NON_STAGE_ENGINE_CLI_FIELDS
-        | frozenset(f.name for f in fields(cast(Any, VllmOmniOrchestratorConfig)))
+        | frozenset(f.name for f in fields(VllmOmniOrchestratorConfig))
         | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS)
     )
     return candidates - externally_consumed
@@ -1239,7 +1238,7 @@ def _stage_sampling_params(
 
 def _orchestrator_cli_overrides(cli_overrides: Mapping[str, Any]) -> dict[str, Any]:
     overrides: dict[str, Any] = {}
-    for config_field in fields(cast(Any, VllmOmniOrchestratorConfig)):
+    for config_field in fields(VllmOmniOrchestratorConfig):
         name = config_field.name
         if name == "deploy_config_path":
             continue
@@ -1444,8 +1443,8 @@ def _with_resolved_processors(
     input_proc: str | None,
     next_stage_proc: str | None,
 ) -> StageConfigType:
-    setattr(stage_config, "_resolved_custom_process_input_func", input_proc)
-    setattr(stage_config, "_resolved_custom_process_next_stage_input_func", next_stage_proc)
+    stage_config._resolved_custom_process_input_func = input_proc
+    stage_config._resolved_custom_process_next_stage_input_func = next_stage_proc
     return stage_config
 
 
@@ -1559,16 +1558,13 @@ def _build_stage_config(
         builder = _STAGE_CONFIG_BUILDERS[topology.execution_type]
     except KeyError as exc:
         raise ValueError(f"Unsupported stage execution type: {topology.execution_type!r}") from exc
-    return cast(
-        StageConfigType,
-        builder(
-            pipeline,
-            deploy,
-            topology,
-            stage_deploy,
-            engine,
-            model=model,
-        ),
+    return builder(
+        pipeline,
+        deploy,
+        topology,
+        stage_deploy,
+        engine,
+        model=model,
     )
 
 
@@ -1612,7 +1608,7 @@ def _build_model_config(
         kwargs["model_subdir"] = topology.model_subdir
     if "tokenizer_subdir" not in kwargs and topology.tokenizer_subdir is not None:
         kwargs["tokenizer_subdir"] = topology.tokenizer_subdir
-    return cast(Any, OmniStageModelConfig)(
+    return OmniStageModelConfig(
         default_sampling_params=default_sampling_params,
         duplex_max_sessions=duplex_max_sessions,
         **kwargs,
@@ -1703,7 +1699,7 @@ def _build_connector_config(
 ) -> OmniStageConnectorConfig:
     output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None
     input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None
-    return cast(Any, OmniStageConnectorConfig)(
+    return OmniStageConnectorConfig(
         async_chunk=bool(deploy.async_chunk),
         omni_kv_config=_copy_value(engine.get("omni_kv_config")),
         output_connectors=_copy_value(output_connectors) if output_connectors else None,
@@ -1852,11 +1848,11 @@ class VllmOmniConfig:
             for topology in pipeline_cfg.stages
         )
 
-        orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)(
+        orchestrator_config = VllmOmniOrchestratorConfig(
             deploy_config_path=loaded_deploy_config_path,
             **_orchestrator_cli_overrides(cli_overrides),
         )
-        return cast(Any, cls)(
+        return cls(
             pipeline_config=pipeline_cfg,
             stage_configs=stage_configs,
             orchestrator_config=orchestrator_config,

vllm_omni/config/stage_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Stage configuration system for vLLM-Omni."""
 
 from __future__ import annotations
@@ -234,7 +234,6 @@ class PipelineConfig:
     # this value to auto-detect the pipeline.  Only needed for diffusers-style
     # multi-component repos (e.g. GLM-Image).  ``None`` = not a diffusers model.
     diffusers_class_name: str | None = None
-    diffusers_class_aliases: tuple[str, ...] = ()
     endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
     # Optional model-owned duplex planner loaded by the stable engine runtime.
     duplex_runtime_extension: str | None = None
@@ -357,7 +356,6 @@ class StageDeployConfig:
     diffusers_call_kwargs: dict[str, Any] | None = None
     diffusion_quantization_config: str | None = None
     diffusion_attention_backend: str | None = None
-    fastvideo_vsa_topk: int | None = None
     diffusion_attention_config: dict[str, Any] | None = None
 
     # Diffusion execution, cache, and VAE behavior.

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -218,16 +218,21 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if stop_after_transfer and req_id in self.requests_needing_kv_transfer:
             self.pending_stop_after_extraction.add(req_id)
 
+    def _drop_aborted_queued_requests(self) -> None:
+        for queue in (self.waiting, self.skipped_waiting):
+            aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED]
+            if aborted:
+                queue.remove_requests(aborted)
+        self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED]
+
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         # Remove FINISHED_ABORTED requests before the upstream scheduler sees
         # them. Upstream vllm raises RuntimeError on this status; omni allows
         # async abort (e.g. client disconnect during TTS streaming) to leave
         # requests in the waiting/running queues temporarily.
-        for queue in (self.waiting, self.running):
-            for req in list(queue):
-                if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
-                    queue.remove(req)
+        self._drop_aborted_queued_requests()
         self._process_pending_omni_inputs(model_mode="ar")
+        self._resync_streaming_input_counter()
 
         original_waiting = None
         if self._should_defer_waiting_admission():
@@ -354,6 +359,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+
+            if generated_token_ids and getattr(request, "async_tokens_to_discard", 0) > 0:
+                self._update_request_with_output(request, generated_token_ids)
+                continue
+
             status_before_stop = request.status
             new_logprobs = None
             logprob_validation_failed = False
@@ -651,6 +661,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter and self.chunk_transfer_adapter.receives_chunks:
             self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None)
             if stage_id != 0:
+                session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
                 # Downstream async-chunk stages receive real payloads from the
                 # connector. This update only resumes polling for the next segment.
                 self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
@@ -682,6 +693,30 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             self._release_replaced_streaming_prompt_cache(session)
             self._replace_streaming_session(session, update)
             return
+        native_duplex_update = any(
+            isinstance(info, dict) and info.get("native_duplex") is True for info in update_infos
+        )
+        if native_duplex_update:
+            max_model_len = int(getattr(self.vllm_config.model_config, "max_model_len", 0) or 0)
+            current_prompt_len = len(getattr(session, "prompt_token_ids", None) or ())
+            computed_tokens = int(getattr(session, "num_computed_tokens", 0) or 0)
+            kept_output_len = max(0, computed_tokens - current_prompt_len)
+            next_prompt_len = len(update.prompt_token_ids or ())
+            candidate_len = current_prompt_len + kept_output_len + next_prompt_len
+            # Keep one context position available for the next sampled codec token.
+            if max_model_len > 0 and candidate_len >= max_model_len:
+                logger.info(
+                    "Resetting native duplex Stage-1 prompt before context overflow: "
+                    "req=%s current=%s kept_output=%s next=%s max_model_len=%s",
+                    session.request_id,
+                    current_prompt_len,
+                    kept_output_len,
+                    next_prompt_len,
+                    max_model_len,
+                )
+                self._replace_streaming_session(session, update)
+                return
+        session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
         super()._update_request_as_session(session, update)
         if hasattr(update, "model_intermediate_buffer"):
             session.model_intermediate_buffer = update.model_intermediate_buffer

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -146,6 +146,7 @@ class OmniSchedulerMixin:
 
     def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None:
         """Replace a downstream stage's placeholder with its next payload."""
+        session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
         adapter = getattr(self, "chunk_transfer_adapter", None)
         if adapter is not None:
             adapter.segment_finished_requests.discard(session.request_id)
@@ -645,6 +646,7 @@ class OmniSchedulerMixin:
         )
         finished = super().finish_requests(request_ids, finished_status)
         self._purge_finished_from_running(target_request_ids)
+        self._resync_streaming_input_counter()
 
         for request in finished:
             self._free_input_coordinator_request(request.request_id)
@@ -810,3 +812,18 @@ class OmniSchedulerMixin:
             return resumable_segment_stop and req.request_id not in target_request_ids
 
         self.running[:] = [req for req in self.running if keep_running(req)]
+
+    def _resync_streaming_input_counter(self) -> None:
+        counter = getattr(self, "num_waiting_for_streaming_input", None)
+        if counter is None:
+            return
+        parked = sum(
+            1
+            for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None))
+            if queue
+            for request in queue
+            if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+        )
+        if parked != counter:
+            logger.debug("[Omni] resynced streaming-input counter %s -> %s", counter, parked)
+            self.num_waiting_for_streaming_input = parked

vllm_omni/deploy/qwen3_tts.yaml

@@ -26,14 +26,13 @@ connectors:
       codec_left_context_frames: 72
       # Emit only the first audio chunk early, then return to codec_chunk_frames.
       initial_codec_chunk_frames: 1
-      # Gradual size increase for the first N chunks eliminates the
-      # IC-to-steady cliff (1->25) and improves streaming continuity.
+      # Chunk ramp (opt-in): gradual size increase for first N chunks to
+      # eliminate the IC→steady cliff (1→25). Uncomment to enable.
       # Each entry is the size of that chunk index; indices past the table
-      # fall back to codec_chunk_frames. Ramp overrides
-      # initial_codec_chunk_frames, including the per-request API field.
-      # Set this to null in a deployment overlay to restore the legacy
-      # initial_codec_chunk_frames behavior.
-      codec_chunk_ramp: [4, 4, 8, 16, 25]
+      # fall back to codec_chunk_frames. When enabled, ramp overrides
+      # initial_codec_chunk_frames (including the per-request API field).
+      # Example:
+      #   codec_chunk_ramp: [4, 4, 8, 16, 25]
       # Decoder CUDA Graph sizing:
       #   * async_chunk=true derives all stateful ICL/xvec shapes from
       #     codec_chunk_ramp when enabled, otherwise from

vllm_omni/deploy/qwen3_tts_high_concurrency.yaml

@@ -31,14 +31,13 @@ connectors:
       ref_code_context_frames: 72
       # Emit only the first audio chunk early, then return to codec_chunk_frames.
       initial_codec_chunk_frames: 1
-      # Gradual size increase for the first N chunks eliminates the
-      # IC-to-steady cliff (1->25) and improves streaming continuity.
+      # Chunk ramp (opt-in): gradual size increase for first N chunks to
+      # eliminate the IC→steady cliff (1→25). Uncomment to enable.
       # Each entry is the size of that chunk index; indices past the table
-      # fall back to codec_chunk_frames. Ramp overrides
-      # initial_codec_chunk_frames, including the per-request API field.
-      # Set this to null in a deployment overlay to restore the legacy
-      # initial_codec_chunk_frames behavior.
-      codec_chunk_ramp: [4, 4, 8, 16, 25]
+      # fall back to codec_chunk_frames. When enabled, ramp overrides
+      # initial_codec_chunk_frames (including the per-request API field).
+      # Example:
+      #   codec_chunk_ramp: [4, 4, 8, 16, 25]
       # async_chunk derives stateful ICL/xvec graph shapes from the ramp when
       # enabled, otherwise from the fixed chunk settings above, and ignores
       # decode_cudagraph_capture_sizes. Keep B>1

vllm_omni/diffusion/attention/backends/fastvideo_vsa.py

@@ -1,431 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import functools
-import math
-from collections.abc import Mapping
-from typing import Any
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
-from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
-from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
-
-logger = init_logger(__name__)
-
-
-# Keep the external FastVideo pybind/CUDA kernel opaque to torch.compile.
-# This mirrors the SageAttention3 backend pattern: tracing the raw extension
-# through Dynamo can reach Inductor scheduling with unstable internal op names
-# (e.g. KeyError: "op12").  The custom op gives Dynamo a single Tensor->Tensor
-# boundary and lets Inductor schedule the surrounding Wan block normally.
-if not hasattr(torch.ops.vllm_omni, "fastvideo_vsa_bshd"):
-
-    @torch.library.custom_op("vllm_omni::fastvideo_vsa_bshd", mutates_args=())
-    def _fastvideo_vsa_bshd_op(
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        variable_block_sizes: torch.Tensor,
-        q_variable_block_sizes: torch.Tensor,
-        compress_attn_weight: torch.Tensor,
-        topk: int,
-        block_t: int,
-        block_h: int,
-        block_w: int,
-    ) -> torch.Tensor:
-        from fastvideo_kernel import video_sparse_attn_bshd
-
-        return video_sparse_attn_bshd(
-            query,
-            key,
-            value,
-            variable_block_sizes=variable_block_sizes,
-            q_variable_block_sizes=q_variable_block_sizes,
-            topk=topk,
-            block_size=(block_t, block_h, block_w),
-            compress_attn_weight=compress_attn_weight if compress_attn_weight.numel() else None,
-        )
-
-    @_fastvideo_vsa_bshd_op.register_fake
-    def _(
-        query,
-        key,
-        value,
-        variable_block_sizes,
-        q_variable_block_sizes,
-        compress_attn_weight,
-        topk,
-        block_t,
-        block_h,
-        block_w,
-    ):
-        del (
-            key,
-            value,
-            variable_block_sizes,
-            q_variable_block_sizes,
-            compress_attn_weight,
-            topk,
-            block_t,
-            block_h,
-            block_w,
-        )
-        return torch.empty_like(query)
-
-
-_fastvideo_vsa_bshd_op = torch.ops.vllm_omni.fastvideo_vsa_bshd
-
-
-@functools.lru_cache(maxsize=32)
-def _get_tile_partition_indices(
-    dit_seq_shape: tuple[int, int, int],
-    tile_size: tuple[int, int, int],
-    device: torch.device,
-) -> torch.Tensor:
-    t_size, h_size, w_size = dit_seq_shape
-    tile_t, tile_h, tile_w = tile_size
-    indices = torch.arange(t_size * h_size * w_size, device=device, dtype=torch.long).reshape(t_size, h_size, w_size)
-    tiles = []
-    for tile_t_idx in range(math.ceil(t_size / tile_t)):
-        for tile_h_idx in range(math.ceil(h_size / tile_h)):
-            for tile_w_idx in range(math.ceil(w_size / tile_w)):
-                tiles.append(
-                    indices[
-                        tile_t_idx * tile_t : min((tile_t_idx + 1) * tile_t, t_size),
-                        tile_h_idx * tile_h : min((tile_h_idx + 1) * tile_h, h_size),
-                        tile_w_idx * tile_w : min((tile_w_idx + 1) * tile_w, w_size),
-                    ].flatten()
-                )
-    return torch.cat(tiles, dim=0)
-
-
-@functools.lru_cache(maxsize=32)
-def _construct_variable_block_sizes(
-    dit_seq_shape: tuple[int, int, int],
-    tile_size: tuple[int, int, int],
-    device: torch.device,
-) -> torch.Tensor:
-    num_tiles = tuple(math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, tile_size))
-
-    def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.Tensor:
-        sizes = torch.full((n_tiles,), tile, dtype=torch.int32, device=device)
-        remainder = dim_len - (n_tiles - 1) * tile
-        sizes[-1] = remainder if remainder > 0 else tile
-        return sizes
-
-    t_sizes = _sizes(dit_seq_shape[0], tile_size[0], num_tiles[0])
-    h_sizes = _sizes(dit_seq_shape[1], tile_size[1], num_tiles[1])
-    w_sizes = _sizes(dit_seq_shape[2], tile_size[2], num_tiles[2])
-    return (t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :]).reshape(-1)
-
-
-@functools.lru_cache(maxsize=32)
-def _get_non_pad_index(variable_block_sizes: torch.Tensor, max_block_size: int) -> torch.Tensor:
-    num_blocks = variable_block_sizes.shape[0]
-    device = variable_block_sizes.device
-    starts = torch.arange(num_blocks, device=device) * max_block_size
-    padded_index = starts[:, None] + torch.arange(max_block_size, device=device)[None, :]
-    valid = torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None]
-    return padded_index[valid]
-
-
-@torch.compiler.disable
-def _get_tile_metadata(
-    dit_seq_shape: tuple[int, int, int],
-    tile_size: tuple[int, int, int],
-    block_elements: int,
-    device: torch.device,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, tile_size, device)
-    variable_block_sizes = _construct_variable_block_sizes(dit_seq_shape, tile_size, device)
-    non_pad_index = _get_non_pad_index(variable_block_sizes, block_elements)
-    untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)]
-    return tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index
-
-
-def _get_vsa_dit_seq_shape(attn_metadata: AttentionMetadata | None) -> tuple[int, int, int] | None:
-    if attn_metadata is None:
-        return None
-    value = attn_metadata.extra.get("vsa_dit_seq_shape")
-    if value is None:
-        return None
-    if not isinstance(value, (list, tuple)) or len(value) != 3:
-        return None
-    return (int(value[0]), int(value[1]), int(value[2]))
-
-
-def _get_gate_compress(attn_metadata: AttentionMetadata | None) -> torch.Tensor | None:
-    if attn_metadata is None:
-        return None
-    value = attn_metadata.extra.get("gate_compress")
-    return value if isinstance(value, torch.Tensor) else None
-
-
-def _preserve_vsa_all_blocks(attn_metadata: AttentionMetadata | None) -> bool:
-    if attn_metadata is None:
-        return False
-    return attn_metadata.extra.get("preserve_vsa_all_blocks") is True
-
-
-class FastVideoVSABackend(AttentionBackend):
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        # FastVideo VSA is intended for video DiT head sizes such as 64/128.
-        # Keep this permissive and let the runtime fallback handle unsupported
-        # cases from the installed fastvideo-kernel build.
-        return []
-
-    @staticmethod
-    def get_name() -> str:
-        return "FASTVIDEO_VSA"
-
-    @staticmethod
-    def get_impl_cls() -> type[FastVideoVSAImpl]:
-        return FastVideoVSAImpl
-
-
-class FastVideoVSAImpl(AttentionImpl):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        qkv_layout: str | None = None,
-        backend_kwargs: Mapping[str, Any] | None = None,
-        **extra_impl_args,
-    ) -> None:
-        backend_kwargs = backend_kwargs or {}
-        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
-        self.softmax_scale = softmax_scale
-        self.causal = causal
-        self.qkv_layout = qkv_layout
-
-        self.topk = int(backend_kwargs.get("topk", 64))
-        self.block_size = self._parse_block_size(backend_kwargs.get("block_size", (4, 8, 8)))
-        self.block_elements = self.block_size[0] * self.block_size[1] * self.block_size[2]
-        self.min_seq_len = int(backend_kwargs.get("min_seq_len", self.block_elements * 2))
-        self.fallback_on_error = bool(backend_kwargs.get("fallback_on_error", True))
-        self.disable_when_sp_active = bool(backend_kwargs.get("disable_when_sp_active", True))
-
-        self.sdpa_fallback = SDPAImpl(
-            num_heads=num_heads,
-            head_size=head_size,
-            softmax_scale=softmax_scale,
-            causal=causal,
-            num_kv_heads=num_kv_heads,
-            qkv_layout=qkv_layout,
-        )
-
-        if self.block_elements != 256:
-            logger.warning(
-                "FASTVIDEO_VSA currently uses fastvideo_kernel.video_sparse_attn_bshd, "
-                "which supports only 256-token blocks. Configured block_size=%s "
-                "(product=%d) will fall back to SDPA.",
-                self.block_size,
-                self.block_elements,
-            )
-
-    @staticmethod
-    def _parse_block_size(value: Any) -> tuple[int, int, int]:
-        if isinstance(value, int):
-            return (value, value, value)
-        if isinstance(value, (list, tuple)) and len(value) == 3:
-            return (int(value[0]), int(value[1]), int(value[2]))
-        raise ValueError(f"FASTVIDEO_VSA block_size must be an int or length-3 tuple/list, got {value!r}")
-
-    def _fallback(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata | None,
-        reason: str,
-    ) -> torch.Tensor:
-        logger.warning_once("FASTVIDEO_VSA falling back to SDPA: %s", reason)
-        return self.sdpa_fallback.forward(query, key, value, attn_metadata)
-
-    def _fallback_reason(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata | None,
-    ) -> str | None:
-        if self.causal:
-            return "causal attention is not supported"
-        if self.block_elements != 256:
-            return f"block_elements must be 256, got {self.block_elements}"
-        if self.topk <= 0:
-            return f"topk must be positive, got {self.topk}"
-        if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
-            return f"expected [B, S, H, D] tensors, got {query.shape}, {key.shape}, {value.shape}"
-        if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
-            return "batch dimensions must match"
-        if query.shape[2:] != key.shape[2:] or query.shape[2:] != value.shape[2:]:
-            return "head/head_dim dimensions must match"
-        if query.shape[1] != key.shape[1] or query.shape[1] != value.shape[1]:
-            return "initial VSA backend supports self-attention with Sq == Skv only"
-        if query.shape[1] < self.min_seq_len:
-            return f"sequence length {query.shape[1]} is below min_seq_len {self.min_seq_len}"
-        dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata)
-        if dit_seq_shape is None:
-            return "vsa_dit_seq_shape metadata is required"
-        if math.prod(dit_seq_shape) != query.shape[1]:
-            return f"vsa_dit_seq_shape product {math.prod(dit_seq_shape)} != sequence length {query.shape[1]}"
-        num_blocks = math.prod(
-            math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size)
-        )
-        if self.topk > num_blocks:
-            return f"topk {self.topk} > num_blocks {num_blocks}"
-        if query.dtype not in (torch.float16, torch.bfloat16):
-            return f"dtype {query.dtype} is not supported"
-        if key.dtype != query.dtype or value.dtype != query.dtype:
-            return "q/k/v dtypes must match"
-        if query.device.type != "cuda" or key.device.type != "cuda" or value.device.type != "cuda":
-            return "q/k/v must be CUDA tensors"
-        expected_scale = self.head_size**-0.5
-        if abs(float(self.softmax_scale) - float(expected_scale)) > 1e-6:
-            return f"softmax_scale {self.softmax_scale} differs from FastVideo VSA scale {expected_scale}"
-        if attn_metadata is not None and attn_metadata.attn_mask is not None:
-            return "attention masks are not supported"
-        if attn_metadata is not None and attn_metadata.full_attn_spans is not None:
-            return "piecewise/full attention spans are not supported"
-        if self.num_heads != self.num_kv_heads:
-            return "GQA/MQA is not supported"
-        if self.disable_when_sp_active and is_forward_context_available():
-            ctx = get_forward_context()
-            if getattr(ctx, "sp_active", False):
-                return "sequence parallel context is active"
-        return None
-
-    def forward_cuda(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata | None = None,
-    ) -> torch.Tensor:
-        reason = self._fallback_reason(query, key, value, attn_metadata)
-        if reason is not None:
-            return self._fallback(query, key, value, attn_metadata, reason)
-
-        seq_len = query.shape[1]
-        dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata)
-        assert dit_seq_shape is not None
-        num_blocks = math.prod(
-            math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size)
-        )
-        preserve_all_blocks = _preserve_vsa_all_blocks(attn_metadata)
-        use_native_sdpa = self.topk == num_blocks and not preserve_all_blocks
-        route = "SDPA" if use_native_sdpa else "VSA_ALL_BLOCKS" if self.topk == num_blocks else "VSA"
-        checkpoint_mode = "fastvideo_dmd" if preserve_all_blocks else "native"
-        logger.info_once(
-            "FASTVIDEO_VSA routing: seq_len=%d, dit_seq_shape=%s, block_size=%s, num_blocks=%d, "
-            "topk=%d, keep_ratio=%.1f%%, checkpoint_mode=%s, route=%s",
-            seq_len,
-            dit_seq_shape,
-            self.block_size,
-            num_blocks,
-            self.topk,
-            100.0 * self.topk / num_blocks,
-            checkpoint_mode,
-            route,
-        )
-        if use_native_sdpa:
-            return self._fallback(
-                query,
-                key,
-                value,
-                attn_metadata,
-                f"topk {self.topk} selects all blocks for a native checkpoint",
-            )
-
-        try:
-            tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index = _get_tile_metadata(
-                dit_seq_shape,
-                self.block_size,
-                self.block_elements,
-                query.device,
-            )
-
-            padded_len = variable_block_sizes.numel() * self.block_elements
-            target_shape = (query.shape[0], padded_len, query.shape[2], query.shape[3])
-            query_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype)
-            key_tiled = torch.zeros_like(query_tiled)
-            value_tiled = torch.zeros_like(query_tiled)
-            query_tiled[:, non_pad_index] = query[:, tile_partition_indices]
-            key_tiled[:, non_pad_index] = key[:, tile_partition_indices]
-            value_tiled[:, non_pad_index] = value[:, tile_partition_indices]
-            # Gate behavior is checkpoint-driven, not user-configured.
-            # Wan VSA layers always provide a gate projection. Its zero
-            # initialization makes checkpoints without gate weights sparse-only;
-            # checkpoints containing to_gate_compress weights use the learned gate.
-            gate_compress = _get_gate_compress(attn_metadata)
-            if gate_compress is None:
-                gate_compress = torch.zeros_like(query)
-            elif gate_compress.shape != query.shape:
-                raise ValueError(f"gate_compress shape {gate_compress.shape} must match query shape {query.shape}")
-            gate_tiled = torch.zeros_like(query_tiled)
-            gate_tiled[:, non_pad_index] = gate_compress[:, tile_partition_indices]
-            compress_attn_weight = gate_tiled
-
-            output = _fastvideo_vsa_bshd_op(
-                query_tiled.contiguous(),
-                key_tiled.contiguous(),
-                value_tiled.contiguous(),
-                variable_block_sizes,
-                variable_block_sizes,
-                compress_attn_weight,
-                self.topk,
-                self.block_size[0],
-                self.block_size[1],
-                self.block_size[2],
-            )
-            return output[:, untile_combined_index].contiguous()
-        except Exception as exc:
-            if not self.fallback_on_error:
-                raise
-            return self._fallback(query, key, value, attn_metadata, f"VSA kernel failed: {exc}")
-
-    def forward_npu(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata | None = None,
-    ) -> torch.Tensor:
-        return self.sdpa_fallback.forward_npu(query, key, value, attn_metadata)
-
-    def forward_xpu(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata | None = None,
-    ) -> torch.Tensor:
-        return self.sdpa_fallback.forward_xpu(query, key, value, attn_metadata)
-
-    def forward_musa(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata | None = None,
-    ) -> torch.Tensor:
-        return self.sdpa_fallback.forward_musa(query, key, value, attn_metadata)

vllm_omni/diffusion/attention/backends/registry.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Diffusion attention backend registry.
 
 This module provides an enum-based registry for diffusion attention backends,
@@ -23,7 +23,7 @@ logger = init_logger(__name__)
 class _DiffusionBackendEnumMeta(EnumMeta):
     """Metaclass for DiffusionAttentionBackendEnum to provide better error messages."""
 
-    def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":  # type: ignore[override]
+    def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":
         """Get backend by name with helpful error messages."""
         try:
             return super().__getitem__(name)  # type: ignore[return-value]
@@ -62,7 +62,6 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
     CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
-    FASTVIDEO_VSA = "vllm_omni.diffusion.attention.backends.fastvideo_vsa.FastVideoVSABackend"
     FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
     FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
     TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend"

vllm_omni/diffusion/data.py

@@ -1,6 +1,6 @@
 # adapted from sglang and fastvideo
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import copy
 import math
 import os
@@ -77,13 +77,11 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
     # Handle "diffusion_attention_backend" shorthand: merge into
     # diffusion_attention_config before field filtering.
     diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
-    fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None)
-    if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None:
+    if diffusion_attn_backend is not None:
         existing = normalized.get("diffusion_attention_config")
         normalized["diffusion_attention_config"] = parse_attention_config(
             existing,
             attention_backend=diffusion_attn_backend,
-            fastvideo_vsa_topk=fastvideo_vsa_topk,
         )
 
     # Check environment variable as fallback for cache_backend.
@@ -106,7 +104,6 @@ def parse_kv_cache_skip_selector(
     """Parse a non-negative index selector such as "0-9,20,25-30"."""
     if selector is None:
         return None
-    values: set[int]
     if isinstance(selector, set):
         values = selector
     elif isinstance(selector, (list, tuple)):
@@ -115,7 +112,7 @@ def parse_kv_cache_skip_selector(
         text = selector.strip()
         if not text:
             return None
-        values = set()
+        values: set[int] = set()
         for chunk in text.split(","):
             token = chunk.strip()
             if not token:
@@ -248,7 +245,6 @@ class DiffusionParallelConfig:
         assert self.pipeline_parallel_size > 0, "Pipeline 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 is not None
         assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
         assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
         assert self.ring_degree > 0, "Ring degree must be > 0"
@@ -366,7 +362,6 @@ class DiffusionParallelConfig:
             self.world_size = world_size
             return 1
 
-        assert self.sequence_parallel_size is not None
         non_dp_size = (
             self.pipeline_parallel_size
             * self.tensor_parallel_size
@@ -1068,7 +1063,6 @@ class OmniDiffusionConfig:
 
         if self.diffusion_compile_granularity == "full":
             incompatible_features = []
-            assert self.parallel_config.sequence_parallel_size is not None
             if self.parallel_config.use_hsdp:
                 incompatible_features.append("HSDP")
             if self.parallel_config.sequence_parallel_size > 1:
@@ -1268,7 +1262,6 @@ class OmniDiffusionConfig:
         if self.model_class_name is None and self.diffusion_load_format == "diffusers":
             self.model_class_name = "DiffusersAdapterPipeline"
 
-        assert self.model is not None
         try:
             config_dict = get_diffusion_model_index(
                 self.model,
@@ -1345,7 +1338,6 @@ class OmniDiffusionConfig:
 
                 from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1
 
-                assert self.model is not None
                 is_hidream_o1 = _looks_like_hidream_o1(self.model, cfg)
                 if self.model_class_name == "HiDreamO1ImagePipeline" and not is_hidream_o1:
                     raise ValueError(f"Checkpoint {self.model} does not have the HiDream-O1 signature")
@@ -1403,7 +1395,6 @@ class OmniDiffusionConfig:
                 elif model_type == "vla":
                     from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 
-                    assert self.model is not None
                     if _looks_like_dreamzero(self.model):
                         self.model_class_name = "DreamZeroPipeline"
                         self.set_tf_model_config(TransformerConfig())
@@ -1602,11 +1593,10 @@ class AttnQuantSpec:
         for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)):
             if v is not None and v not in self._VALID_DTYPES:
                 raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.")
-        for block_name, block_value in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
-            if block_value not in self._VALID_BLOCK_SIZES:
+        for name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
+            if v not in self._VALID_BLOCK_SIZES:
                 raise ValueError(
-                    f"quant.{block_name}={block_value!r} unsupported; "
-                    f"kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
+                    f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
                 )
 
     @property
@@ -1648,7 +1638,6 @@ class AttentionSpec:
     backend: str
     skip_softmax: SkipSoftmaxSpec | None = None
     quant: AttnQuantSpec | None = None
-    fastvideo_vsa_topk: int | None = None
     block_sparse: BlockSparseSpec | None = None
     skip_calibration: dict | None = field(default=None, repr=False)
 
@@ -1668,11 +1657,6 @@ class AttentionSpec:
                 f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
                 f"backend={self.backend!r}. Remove quant or set a supported backend."
             )
-        if self.fastvideo_vsa_topk is not None:
-            if self.backend.upper() != "FASTVIDEO_VSA":
-                raise ValueError("fastvideo_vsa_topk is only supported by the FASTVIDEO_VSA backend.")
-            if self.fastvideo_vsa_topk <= 0:
-                raise ValueError("fastvideo_vsa_topk must be positive.")
         if self.backend.upper() in BLOCK_SPARSE_BACKENDS:
             # Selecting the backend is the opt-in; without an explicit block the
             # defaults apply rather than silently running dense.
@@ -1714,15 +1698,12 @@ class AttentionSpec:
             if q.flashinfer_backend is not None:
                 quant_kw["flashinfer_backend"] = q.flashinfer_backend
             kw["quant"] = quant_kw
-        if self.fastvideo_vsa_topk is not None:
-            kw["topk"] = self.fastvideo_vsa_topk
         if self.block_sparse is not None:
             bs = self.block_sparse
             kw["sparsity"] = bs.sparsity
             kw["start_step"] = bs.start_step
             if bs.skip_layer_indices:
                 kw["skip_layers"] = sorted(bs.skip_layer_indices)
-
         return kw or None
 
 
@@ -1792,7 +1773,7 @@ class AttentionConfig:
             normalized[role] = node
             return
 
-        spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse"}
+        spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"}
         node_dict = dict(node)
         node_keys = set(node_dict)
         if node_keys & spec_keys:
@@ -1831,7 +1812,6 @@ def parse_attention_config(
     attention_config: AttentionConfig | Mapping[str, Any] | None = None,
     *,
     attention_backend: str | None = None,
-    fastvideo_vsa_topk: int | None = None,
 ) -> AttentionConfig:
     """Pure type-conversion: coerce *attention_config* to an AttentionConfig.
 
@@ -1859,14 +1839,6 @@ def parse_attention_config(
         if attention_backend.lower() != "auto":
             normalized.default = AttentionSpec(backend=attention_backend)
 
-    if fastvideo_vsa_topk is not None:
-        if normalized.default is None:
-            raise ValueError("--fastvideo-vsa-topk requires --diffusion-attention-backend FASTVIDEO_VSA.")
-        if normalized.default.backend.upper() != "FASTVIDEO_VSA":
-            raise ValueError("--fastvideo-vsa-topk is only valid with the FASTVIDEO_VSA backend.")
-        normalized.default.fastvideo_vsa_topk = fastvideo_vsa_topk
-        normalized.default.__post_init__()
-
     return normalized
 
 

vllm_omni/diffusion/lora/manager.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import math
 import time
@@ -289,35 +289,11 @@ class DiffusionLoRAManager:
         lora_path = get_adapter_absolute_path(lora_request.lora_path)
         logger.debug("Resolved LoRA path: %s", lora_path)
 
-        model_loader = getattr(self.pipeline, "_load_diffusion_lora_adapter", None)
-        loaded = None
-        if callable(model_loader):
-            loaded = model_loader(
-                lora_request=lora_request,
-                lora_path=lora_path,
-                dtype=self.dtype,
-            )
-
-        if loaded is None:
-            peft_helper = PEFTHelper.from_local_dir(
-                lora_path,
-                max_position_embeddings=None,  # no need in diffusion
-                tensorizer_config_dict=lora_request.tensorizer_config_dict,
-            )
-
-            lora_model = LoRAModel.from_local_checkpoint(
-                lora_path,
-                expected_lora_modules=self._expected_lora_modules,
-                peft_helper=peft_helper,
-                lora_model_id=lora_request.lora_int_id,
-                device="cpu",  # consistent w/ vllm's behavior
-                dtype=self.dtype,
-                model_vocab_size=None,
-                tensorizer_config_dict=lora_request.tensorizer_config_dict,
-                weights_mapper=None,
-            )
-        else:
-            lora_model, peft_helper = loaded
+        peft_helper = PEFTHelper.from_local_dir(
+            lora_path,
+            max_position_embeddings=None,  # no need in diffusion
+            tensorizer_config_dict=lora_request.tensorizer_config_dict,
+        )
 
         logger.info(
             "Loaded PEFT config: r=%d, lora_alpha=%d, target_modules=%s",
@@ -326,6 +302,18 @@ class DiffusionLoRAManager:
             peft_helper.target_modules,
         )
 
+        lora_model = LoRAModel.from_local_checkpoint(
+            lora_path,
+            expected_lora_modules=self._expected_lora_modules,
+            peft_helper=peft_helper,
+            lora_model_id=lora_request.lora_int_id,
+            device="cpu",  # consistent w/ vllm's behavior
+            dtype=self.dtype,
+            model_vocab_size=None,
+            tensorizer_config_dict=lora_request.tensorizer_config_dict,
+            weights_mapper=None,
+        )
+
         logger.info(
             "Loaded LoRA model: id=%d, num_modules=%d, modules=%s",
             lora_model.id,
@@ -534,17 +522,13 @@ class DiffusionLoRAManager:
         matches_scale = self._adapter_scales.get(adapter_id) == rounded_scale
         return is_active and matches_scale
 
-    def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None:
-        binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None)
-        lora_names_by_id = (
-            {id(weights): name for name, weights in lora_model.loras.items()} if callable(binding_validator) else {}
-        )
-        bound_lora_names: set[str] = set()
+    def _activate_adapter(self, adapter_id: int, scale: float) -> None:
+        if self._is_active_at_scale(adapter_id, scale):
+            logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale)
+            return
 
-        def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None:
-            name = lora_names_by_id.get(id(weights))
-            if name is not None:
-                bound_lora_names.add(name)
+        logger.info("Activating adapter: id=%d", adapter_id)
+        lora_model = self._registered_adapters[adapter_id]
 
         # activate weights in each LoRA layer
         for full_module_name, lora_layer in self._lora_modules.items():
@@ -586,9 +570,6 @@ class DiffusionLoRAManager:
                         lora_b_list.append(sub_lora.lora_b * scale)
 
                     lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
-                    for sub_lora in sub_loras:
-                        if sub_lora is not None:
-                            _record_bound(sub_lora)
                     logger.debug(
                         "Activated packed LoRA for %s via submodules=%s (scale=%.2f)",
                         full_module_name,
@@ -607,7 +588,6 @@ class DiffusionLoRAManager:
                     for b in lora_weights.lora_b
                 ]
                 lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
-                _record_bound(lora_weights)
                 logger.debug(
                     "Activated packed LoRA for %s (scale=%.2f)",
                     full_module_name,
@@ -638,7 +618,6 @@ class DiffusionLoRAManager:
                 lora_a_list = [lora_weights.lora_a] * n_slices
                 lora_b_list = [b * scale for b in b_splits]
                 lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
-                _record_bound(lora_weights)
                 logger.debug(
                     "Activated fused LoRA for packed layer %s (scale=%.2f)",
                     full_module_name,
@@ -648,7 +627,6 @@ class DiffusionLoRAManager:
 
             scaled_lora_b = lora_weights.lora_b * scale
             lora_layer.set_lora(index=0, lora_a=lora_weights.lora_a, lora_b=scaled_lora_b)
-            _record_bound(lora_weights)
             logger.debug(
                 "Activated LoRA for %s: lora_a shape=%s, lora_b shape=%s, scale=%.2f",
                 full_module_name,
@@ -657,33 +635,6 @@ class DiffusionLoRAManager:
                 scale,
             )
 
-        if callable(binding_validator):
-            binding_validator(
-                lora_model=lora_model,
-                bound_lora_names=frozenset(bound_lora_names),
-            )
-
-    def _reset_lora_layers(self) -> None:
-        for lora_layer in self._lora_modules.values():
-            lora_layer.reset_lora(0)
-
-    def _activate_adapter(self, adapter_id: int, scale: float) -> None:
-        if self._is_active_at_scale(adapter_id, scale):
-            logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale)
-            return
-
-        logger.info("Activating adapter: id=%d", adapter_id)
-        lora_model = self._registered_adapters[adapter_id]
-        # Binding overwrites slot 0 incrementally. Invalidate the fast-path
-        # state before the first mutation and leave every wrapper inactive if
-        # any set_lora() call or model validator fails.
-        self._active_adapter_id = None
-        try:
-            self._bind_adapter_weights(lora_model, scale)
-        except Exception:
-            self._reset_lora_layers()
-            raise
-
         self._active_adapter_id = adapter_id
         self._update_adapter_scale(adapter_id, scale)
 
@@ -692,7 +643,8 @@ class DiffusionLoRAManager:
             logger.debug("All adapters already inactive")
             return
         logger.info("Deactivating all adapters: %d layers", len(self._lora_modules))
-        self._reset_lora_layers()
+        for lora_layer in self._lora_modules.values():
+            lora_layer.reset_lora(0)
         self._active_adapter_id = None
         logger.debug("All adapters deactivated")
 

vllm_omni/diffusion/models/minimax_h3/lora.py

@@ -1,205 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import math
-from pathlib import Path
-
-import torch
-from safetensors import safe_open
-from vllm.lora.lora_model import LoRAModel
-from vllm.lora.lora_weights import PackedLoRALayerWeights
-from vllm.lora.peft_helper import PEFTHelper
-from vllm.model_executor.models.utils import WeightsMapper
-
-from vllm_omni.lora.request import LoRARequest
-
-_TURBO_RANK = 128
-_TURBO_ALPHA = 128
-_TURBO_HIDDEN_SIZE = 5376
-_TURBO_ATTENTION_INNER_SIZE = 7168
-_TURBO_FFN_HIDDEN_SIZE = 14336
-_TURBO_FILENAME = "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-_LORA_A_SUFFIX = ".lora_A.default.weight"
-_LORA_B_SUFFIX = ".lora_B.default.weight"
-_TURBO_TARGETS = frozenset({"to_q", "to_k", "to_v", "out_proj", "fc1", "fc2"})
-_TURBO_RAW_TARGET_SUFFIXES = (
-    "attn.to_q",
-    "attn.to_k",
-    "attn.to_v",
-    "attn.to_out.0",
-    "ff.net.0.proj",
-    "ff.net.2",
-)
-_TURBO_TARGET_DIMS = {
-    "attn.to_q": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE),
-    "attn.to_k": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE),
-    "attn.to_v": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE),
-    "attn.to_out.0": (_TURBO_ATTENTION_INNER_SIZE, _TURBO_HIDDEN_SIZE),
-    "ff.net.0.proj": (_TURBO_HIDDEN_SIZE, 2 * _TURBO_FFN_HIDDEN_SIZE),
-    "ff.net.2": (_TURBO_FFN_HIDDEN_SIZE, _TURBO_HIDDEN_SIZE),
-}
-_TURBO_EXPECTED_RAW_TARGETS = frozenset(
-    f"{prefix}.{block_index}.{suffix}"
-    for prefix, block_count in (
-        ("transformer_blocks", 50),
-        ("token_refiner.refiner_blocks", 2),
-    )
-    for block_index in range(block_count)
-    for suffix in _TURBO_RAW_TARGET_SUFFIXES
-)
-_TURBO_TARGET_PATTERN = (
-    r"^transformer\.(?:token_refiner\.blocks|blocks)\.\d+\."
-    r"(?:attn\.(?:to_q|to_k|to_v|out_proj)|mlp\.(?:fc1|fc2))$"
-)
-
-_TURBO_WEIGHTS_MAPPER = WeightsMapper(
-    orig_to_new_substr={
-        "token_refiner.refiner_blocks.": "token_refiner.blocks.",
-        "transformer_blocks.": "blocks.",
-        ".attn.to_out.0.": ".attn.out_proj.",
-        ".ff.net.0.proj.": ".mlp.fc1.",
-        ".ff.net.2.": ".mlp.fc2.",
-        ".lora_A.default.": ".lora_A.",
-        ".lora_B.default.": ".lora_B.",
-    }
-)
-
-
-def _select_turbo_file(artifact_path: str | Path) -> Path | None:
-    path = Path(artifact_path)
-    if path.is_file():
-        return path if path.suffix == ".safetensors" else None
-    if not path.is_dir():
-        return None
-
-    candidate = path / _TURBO_FILENAME
-    return candidate if candidate.is_file() else None
-
-
-def _validate_and_convert_tensors(checkpoint) -> dict[str, torch.Tensor]:
-    tensors: dict[str, torch.Tensor] = {}
-    pairs: dict[str, set[str]] = {}
-    raw_targets: set[str] = set()
-    for name in checkpoint.keys():
-        if name.endswith(_LORA_A_SUFFIX):
-            raw_target = name[: -len(_LORA_A_SUFFIX)]
-            side = "a"
-        elif name.endswith(_LORA_B_SUFFIX):
-            raw_target = name[: -len(_LORA_B_SUFFIX)]
-            side = "b"
-        else:
-            raise ValueError(f"Unconsumed MiniMax-H3 Turbo tensor: {name!r}")
-        raw_targets.add(raw_target)
-
-        mapped_name = _TURBO_WEIGHTS_MAPPER.apply_list([name])[0]
-        mapped_target = mapped_name.rsplit(".lora_", 1)[0]
-        if mapped_target.rsplit(".", 1)[-1] not in _TURBO_TARGETS:
-            raise ValueError(f"Unsupported MiniMax-H3 Turbo target: {raw_target!r}")
-        target_sides = pairs.setdefault(mapped_target, set())
-        if side in target_sides:
-            raise ValueError(f"Duplicate MiniMax-H3 Turbo tensor for {mapped_target}.{side}")
-        target_sides.add(side)
-
-        tensor = checkpoint.get_tensor(name)
-        if tensor.ndim != 2:
-            raise ValueError(f"MiniMax-H3 Turbo LoRA tensors must be matrices, got {name}={tuple(tensor.shape)}")
-        suffix = next((suffix for suffix in _TURBO_RAW_TARGET_SUFFIXES if raw_target.endswith(suffix)), None)
-        if suffix is None:
-            raise ValueError(f"MiniMax-H3 Turbo LoRA contains unsupported target: {raw_target}")
-        input_dim, output_dim = _TURBO_TARGET_DIMS[suffix]
-        expected_shape = (_TURBO_RANK, input_dim) if side == "a" else (output_dim, _TURBO_RANK)
-        if tuple(tensor.shape) != expected_shape:
-            raise ValueError(
-                f"MiniMax-H3 Turbo tensor has invalid global shape: {name}={tuple(tensor.shape)}, "
-                f"expected={expected_shape}"
-            )
-        if side == "b" and ".ff.net.0.proj." in name:
-            value, gate = tensor.chunk(2, dim=0)
-            tensor = torch.cat((gate, value), dim=0).contiguous()
-        tensors[name] = tensor
-
-    incomplete = sorted(target for target, sides in pairs.items() if sides != {"a", "b"})
-    if incomplete:
-        raise ValueError(f"Incomplete MiniMax-H3 Turbo LoRA pairs: {incomplete}")
-    missing = sorted(_TURBO_EXPECTED_RAW_TARGETS - raw_targets)
-    unexpected = sorted(raw_targets - _TURBO_EXPECTED_RAW_TARGETS)
-    if missing or unexpected:
-        raise ValueError(
-            "MiniMax-H3 Turbo target set does not match the supported v1.0 artifact: "
-            f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}"
-        )
-    return tensors
-
-
-def _pack_h3_turbo_fc1(lora_model: LoRAModel) -> None:
-    """Represent H3's fused gate/up projection without generic layout guesses."""
-
-    for module_name, weights in tuple(lora_model.loras.items()):
-        if not module_name.endswith(".mlp.fc1"):
-            continue
-        gate_b, up_b = weights.lora_b.chunk(2, dim=0)
-        lora_model.loras[module_name] = PackedLoRALayerWeights(
-            module_name=module_name,
-            rank=weights.rank,
-            lora_alphas=[weights.lora_alpha, weights.lora_alpha],
-            lora_a=[weights.lora_a, weights.lora_a],
-            lora_b=[gate_b.contiguous(), up_b.contiguous()],
-            scaling=[weights.scaling, weights.scaling],
-        )
-
-
-def load_minimax_h3_turbo_lora(
-    *,
-    partition: str,
-    lora_request: LoRARequest,
-    lora_path: str | Path,
-    dtype: torch.dtype,
-    unsupported_offload_mode: str | None = None,
-) -> tuple[LoRAModel, PEFTHelper] | None:
-    """Load the published LightX2V Turbo v1.0 through the legacy manager."""
-
-    lora_file = _select_turbo_file(lora_path)
-    if lora_file is None:
-        return None
-    with safe_open(lora_file, framework="pt", device="cpu") as checkpoint:
-        metadata = checkpoint.metadata() or {}
-        if metadata.get("key_format") != "minimax-h3-diffusers":
-            if lora_file.name == _TURBO_FILENAME:
-                raise ValueError(
-                    "MiniMax-H3 Turbo v1.0 requires safetensors metadata key_format='minimax-h3-diffusers'"
-                )
-            return None
-        if lora_file.name != _TURBO_FILENAME:
-            raise ValueError(f"MiniMax-H3 Turbo supports only {_TURBO_FILENAME!r}, got {lora_file.name!r}")
-        raw_alpha = metadata.get("alpha")
-        try:
-            alpha = float(raw_alpha) if raw_alpha is not None else math.nan
-        except ValueError as exc:
-            raise ValueError(f"MiniMax-H3 Turbo alpha must be numeric, got {raw_alpha!r}") from exc
-        if alpha != _TURBO_ALPHA:
-            raise ValueError(f"MiniMax-H3 Turbo v1.0 requires alpha={_TURBO_ALPHA}, got {raw_alpha!r}")
-        if partition == "ref2va":
-            raise ValueError("MiniMax-H3 Turbo LoRA supports FL2VA/T2VA only")
-        if unsupported_offload_mode is not None:
-            raise ValueError(f"MiniMax-H3 Turbo dynamic LoRA does not support {unsupported_offload_mode}")
-        tensors = _validate_and_convert_tensors(checkpoint)
-
-    peft_helper = PEFTHelper.from_dict(
-        {
-            "r": _TURBO_RANK,
-            "lora_alpha": _TURBO_ALPHA,
-            "target_modules": _TURBO_TARGET_PATTERN,
-        }
-    )
-    lora_model = LoRAModel.from_lora_tensors(
-        lora_model_id=lora_request.lora_int_id,
-        tensors=tensors,
-        peft_helper=peft_helper,
-        device="cpu",
-        dtype=dtype,
-        weights_mapper=_TURBO_WEIGHTS_MAPPER,
-    )
-    _pack_h3_turbo_fc1(lora_model)
-    return lora_model, peft_helper

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -954,14 +954,6 @@ class MiniMaxH3DiTModel(nn.Module):
     # quantization or LoRA to map onto. Address the fused layers directly, e.g.
     # ignored_layers=["blocks.0.attn.qkv_proj"].
     packed_modules_mapping = {}
-    # Turbo LoRA checkpoints publish separate Q/K/V adapters. This declaration
-    # lets the legacy diffusion LoRA manager bind them to the packed QKV layer;
-    # it does not change the fused base-checkpoint loading path above.
-    stacked_params_mapping = (
-        (".attn.qkv_proj", ".attn.to_q", "q"),
-        (".attn.qkv_proj", ".attn.to_k", "k"),
-        (".attn.qkv_proj", ".attn.to_v", "v"),
-    )
 
     def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None:
         if tp_size < 1:

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """vLLM-Omni pipeline for MiniMax H3 FL2VA and Ref2VA partitions."""
 
 from __future__ import annotations
@@ -12,7 +11,7 @@ from collections.abc import Iterable, Mapping, Sequence
 from contextlib import contextmanager, nullcontext
 from itertools import groupby
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -75,7 +74,6 @@ from .condition_noise import (
 )
 from .denoise_loop import MiniMaxH3DenoiseBranch, minimax_h3_denoise_loop
 from .encoder import MiniMaxH3Qwen3VLEncoder
-from .lora import load_minimax_h3_turbo_lora
 from .minimax_h3_transformer import MiniMaxH3DiTModel
 from .packed_sequence import (
     minimax_h3_packed_sequence,
@@ -112,12 +110,6 @@ from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE
 
 logger = init_logger(__name__)
 
-if TYPE_CHECKING:
-    from vllm.lora.lora_model import LoRAModel
-    from vllm.lora.peft_helper import PEFTHelper
-
-    from vllm_omni.lora.request import LoRARequest
-
 MINIMAX_H3_FPS = 24
 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
 MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
@@ -138,9 +130,6 @@ MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024
 MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"})
 MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0
 MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0
-MINIMAX_H3_TURBO_SIGMA_POINTS = 5
-MINIMAX_H3_TURBO_VIDEO_SHIFT = 6.0
-MINIMAX_H3_TURBO_AUDIO_SHIFT = 3.0
 MINIMAX_H3_DOWNLOAD_PATTERNS = [
     "FL2VA/**",
     "Ref2VA/model_index.json",
@@ -570,80 +559,6 @@ class MiniMaxH3Pipeline(
     # uniform path available to partially constructed pipelines.
     _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}
 
-    def _load_diffusion_lora_adapter(
-        self,
-        *,
-        lora_request: LoRARequest,
-        lora_path: str | Path,
-        dtype: torch.dtype,
-    ) -> tuple[LoRAModel, PEFTHelper] | None:
-        # A cache eviction may be followed by a different adapter reusing the
-        # same client-supplied ID. Every real load replaces the classification.
-        self._turbo_lora_adapter_ids.discard(lora_request.lora_int_id)
-        od_config = getattr(self, "od_config", None)
-        offload_modes = []
-        if getattr(od_config, "enable_cpu_offload", False):
-            offload_modes.append("model-level CPU offload (--enable-cpu-offload)")
-        if getattr(od_config, "enable_layerwise_offload", False):
-            offload_modes.append("layerwise offload (--enable-layerwise-offload)")
-        if getattr(od_config, "enable_distributed_layerwise_offload", False):
-            offload_modes.append("distributed layerwise offload (--enable-distributed-layerwise-offload)")
-        loaded = load_minimax_h3_turbo_lora(
-            partition=self.partition,
-            lora_request=lora_request,
-            lora_path=lora_path,
-            dtype=dtype,
-            unsupported_offload_mode=" or ".join(offload_modes) or None,
-        )
-        if loaded is not None:
-            self._turbo_lora_adapter_ids.add(lora_request.lora_int_id)
-        return loaded
-
-    def _validate_diffusion_lora_binding(
-        self,
-        *,
-        lora_model: LoRAModel,
-        bound_lora_names: frozenset[str],
-    ) -> None:
-        if lora_model.id not in self._turbo_lora_adapter_ids:
-            return
-        missing = sorted(set(lora_model.loras) - bound_lora_names)
-        if missing:
-            raise ValueError(
-                "MiniMax-H3 Turbo LoRA binding is incomplete: "
-                f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}"
-            )
-
-    def _has_active_turbo_lora(self, sampling: Any) -> bool:
-        lora_request = sampling.lora_request
-        return (
-            lora_request is not None
-            and not math.isclose(0.0, float(sampling.lora_scale))
-            and lora_request.lora_int_id in self._turbo_lora_adapter_ids
-        )
-
-    def _validate_turbo_sampling(self, sampling: Any) -> None:
-        extra = sampling.extra_args or {}
-        sigma_points = sampling.num_inference_steps
-        if sigma_points != MINIMAX_H3_TURBO_SIGMA_POINTS:
-            raise OmniClientError(
-                "MiniMax-H3 Turbo requires num_inference_steps=5 (five sigma points produce four denoiser evaluations)"
-            )
-        try:
-            video_shift = float(extra.get("flow_shift", self.default_video_shift))
-        except (TypeError, ValueError) as exc:
-            raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}") from exc
-        if not math.isclose(video_shift, MINIMAX_H3_TURBO_VIDEO_SHIFT):
-            raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}")
-        try:
-            audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift))
-        except (TypeError, ValueError) as exc:
-            raise OmniClientError(
-                f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}"
-            ) from exc
-        if not math.isclose(audio_shift, MINIMAX_H3_TURBO_AUDIO_SHIFT):
-            raise OmniClientError(f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}")
-
     def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
         """Adopt runner-installed generic Cache-DiT for request transitions."""
 
@@ -674,7 +589,6 @@ class MiniMaxH3Pipeline(
             getattr(od_config, "task_type", None),
             str(od_config.model),
         )
-        self._turbo_lora_adapter_ids: set[int] = set()
         model_root = _resolve_minimax_h3_model_root(
             str(od_config.model),
             od_config.revision,
@@ -867,8 +781,6 @@ class MiniMaxH3Pipeline(
         self,
         requested: str | None,
         multi_modal_data: dict[str, Any],
-        *,
-        has_turbo_lora: bool = False,
     ) -> str:
         if requested is None:
             # A Ref2VA-only startup has no FL2VA transformer; preserve its
@@ -886,8 +798,6 @@ class MiniMaxH3Pipeline(
             raise OmniClientError(
                 f"checkpoint partition {self.partition!r} supports {sorted(self.supported_tasks)}, got task={task!r}"
             )
-        if task == "ref2va" and has_turbo_lora:
-            raise OmniClientError("MiniMax-H3 Turbo LoRA supports T2VA/FL2VA requests only")
         return task
 
     def _resolve_shape(
@@ -1745,14 +1655,7 @@ class MiniMaxH3Pipeline(
         quality = sampling.quality
         logger.debug("MiniMax H3 request quality=%s", quality)
         extra = sampling.extra_args or {}
-        has_turbo_lora = self._has_active_turbo_lora(sampling)
-        task = self._resolve_task(
-            extra.get("task"),
-            multi_modal_data,
-            has_turbo_lora=has_turbo_lora,
-        )
-        if has_turbo_lora:
-            self._validate_turbo_sampling(sampling)
+        task = self._resolve_task(extra.get("task"), multi_modal_data)
 
         raw_image = multi_modal_data.get("image")
         raw_videos = multi_modal_data.get("video")

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from __future__ import annotations
 
@@ -33,7 +33,7 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
-from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, WanTransformer3DModel
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
 from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -44,8 +44,6 @@ from vllm_omni.platforms import current_omni_platform
 logger = logging.getLogger(__name__)
 DEBUG_PERF = False
 WAN_SAMPLE_SOLVER_CHOICES = {"unipc", "euler"}
-FASTWAN_DMD_TIMESTEPS = (1000.0, 757.0, 522.0)
-FASTWAN_DMD_SCHEDULER_SHIFT = 8.0
 
 
 def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any:
@@ -66,30 +64,6 @@ def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any:
     )
 
 
-def load_wan_weights_with_optional_gate(model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-    """Load Wan weights and discard optional VSA gates absent from the checkpoint."""
-    gate_param_names = {name for name, _ in model.named_parameters() if ".to_gate_compress." in name}
-    has_gate_compress_weights = False
-
-    def tracked_weights():
-        nonlocal has_gate_compress_weights
-        for name, weight in weights:
-            if ".to_gate_compress." in name:
-                has_gate_compress_weights = True
-            yield name, weight
-
-    loaded_weights = AutoWeightsLoader(model).load_weights(tracked_weights())
-    setattr(model, "has_gate_compress_weights", has_gate_compress_weights)
-    if not has_gate_compress_weights:
-        for module in model.modules():
-            if isinstance(module, WanSelfAttention):
-                module.to_gate_compress = None
-
-    # Optional gate parameters are absent from public FullAttn/DMD checkpoints.
-    loaded_weights.update(gate_param_names)
-    return loaded_weights
-
-
 def resolve_wan_sample_solver(req: OmniDiffusionRequest, default: str = "unipc") -> str:
     extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
     raw = extra_args.get("sample_solver", default)
@@ -340,7 +314,6 @@ class Wan22Pipeline(
 
         # Read model_index.json to detect expand_timesteps mode (for TI2V-5B)
         self.expand_timesteps = False
-        self.is_dmd = False
         self.has_transformer_2 = False
         if local_files_only:
             model_index_path = os.path.join(model, "model_index.json")
@@ -348,7 +321,6 @@ class Wan22Pipeline(
                 with open(model_index_path) as f:
                     model_index = json.load(f)
                     self.expand_timesteps = model_index.get("expand_timesteps", False)
-                    self.is_dmd = model_index.get("_class_name") == "WanDMDPipeline"
             # Check if this is a two-stage model (MoE with transformer_2)
             transformer_2_path = os.path.join(model, "transformer_2")
             self.has_transformer_2 = os.path.exists(transformer_2_path)
@@ -361,7 +333,6 @@ class Wan22Pipeline(
                 with open(model_index_path) as f:
                     model_index = json.load(f)
                     self.expand_timesteps = model_index.get("expand_timesteps", False)
-                    self.is_dmd = model_index.get("_class_name") == "WanDMDPipeline"
                     # Check transformer_2 from model_index
                     transformer_2_info = model_index.get("transformer_2", [None, None])
                     self.has_transformer_2 = transformer_2_info[0] is not None
@@ -451,10 +422,6 @@ class Wan22Pipeline(
         else:
             self.transformer_2 = None
 
-        for transformer in (self.transformer, self.transformer_2):
-            if transformer is not None:
-                transformer.preserve_vsa_all_blocks = self.is_dmd
-
         # Store the active transformer config
         if load_transformer:
             self.transformer_config = self.transformer.config
@@ -463,14 +430,8 @@ class Wan22Pipeline(
         else:
             raise RuntimeError("No transformer loaded")
 
-        self._sample_solver = "euler" if self.is_dmd else "unipc"
-        self._flow_shift = (
-            FASTWAN_DMD_SCHEDULER_SHIFT
-            if self.is_dmd
-            else od_config.flow_shift
-            if od_config.flow_shift is not None
-            else 5.0
-        )
+        self._sample_solver = "unipc"
+        self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0
         self.scheduler = build_wan_scheduler(self._sample_solver, self._flow_shift)
 
         self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4
@@ -519,7 +480,6 @@ class Wan22Pipeline(
         attention_kwargs: dict[str, Any],
         latent_condition: torch.Tensor | None = None,
         first_frame_mask: torch.Tensor | None = None,
-        generator: torch.Generator | None = None,
     ) -> torch.Tensor | AsyncLatents:
         if attention_kwargs is None:
             attention_kwargs = {}
@@ -601,20 +561,7 @@ class Wan22Pipeline(
                     cfg_normalize=False,
                 )
 
-                if self.is_dmd:
-                    pred_clean = self.scheduler.predict_clean(noise_pred, latents, t).to(noise_pred.dtype)
-                    if step_idx + 1 < len(timesteps):
-                        noise = randn_tensor(
-                            latents.shape,
-                            generator=generator,
-                            device=latents.device,
-                            dtype=pred_clean.dtype,
-                        )
-                        latents = self.scheduler.add_noise(pred_clean, noise, timesteps[step_idx + 1])
-                    else:
-                        latents = pred_clean
-                else:
-                    latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
                 pbar.update()
 
         return latents
@@ -653,21 +600,13 @@ class Wan22Pipeline(
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
-        if self.is_dmd:
-            # The checkpoint was distilled for these three transitions. Ignore
-            # request-level step counts, including the engine's 1-step warmup.
-            num_steps = len(FASTWAN_DMD_TIMESTEPS)
-        else:
-            num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
+        num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
 
         output_type = common.output_type or "np"
         num_outputs_per_prompt = common.num_outputs_per_prompt or 1
         attention_kwargs: dict | None = None
 
-        guidance_low, guidance_high = resolve_wan_guidance_scales(
-            common,
-            default_guidance_scale=1.0 if self.is_dmd else 4.0,
-        )
+        guidance_low, guidance_high = resolve_wan_guidance_scales(common, default_guidance_scale=4.0)
 
         # record guidance for properties
         self._guidance_scale = guidance_low
@@ -746,19 +685,17 @@ class Wan22Pipeline(
             current_omni_platform.synchronize()
             _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000
 
-        if self.is_dmd:
-            timesteps = torch.tensor(FASTWAN_DMD_TIMESTEPS, device=device, dtype=torch.float32)
-        else:
-            first_request = req.requests[0]
-            sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver)
-            flow_shift = resolve_wan_flow_shift(first_request, self.od_config)
-            if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
-                self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
-                self._sample_solver = sample_solver
-                self._flow_shift = flow_shift
+        first_request = req.requests[0]
+        sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver)
+        flow_shift = resolve_wan_flow_shift(first_request, self.od_config)
+        if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
+            self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
+            self._sample_solver = sample_solver
+            self._flow_shift = flow_shift
 
-            self.scheduler.set_timesteps(num_steps, device=device)
-            timesteps = self.scheduler.timesteps
+        # Timesteps
+        self.scheduler.set_timesteps(num_steps, device=device)
+        timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
         boundary_timestep = None
         if boundary_ratio is not None:
@@ -880,7 +817,6 @@ class Wan22Pipeline(
             attention_kwargs=attention_kwargs,
             latent_condition=latent_condition,
             first_frame_mask=first_frame_mask,
-            generator=generator,
         )
 
         # Wan2.2 is prone to out of memory errors when predicting large videos
@@ -1067,7 +1003,9 @@ class Wan22Pipeline(
         return latents
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        return load_wan_weights_with_optional_gate(self, weights)
+        """Load weights using AutoWeightsLoader for vLLM integration."""
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
 
     def check_inputs(
         self,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from __future__ import annotations
 
@@ -16,6 +16,7 @@ import torchvision.transforms.functional as TF
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
+from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -35,7 +36,6 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     build_wan_scheduler,
     create_transformer_from_config,
     load_transformer_config,
-    load_wan_weights_with_optional_gate,
     resolve_wan_flow_shift,
     resolve_wan_guidance_scales,
     resolve_wan_sample_solver,
@@ -1014,7 +1014,9 @@ class Wan22I2VPipeline(
             raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.")
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        return load_wan_weights_with_optional_gate(self, weights)
+        """Load weights using AutoWeightsLoader for vLLM integration."""
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
 
 
 # ---------------------------------------------------------------------------

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
 Wan2.2 Speech-to-Video (S2V) Pipeline for vLLM-Omni.
@@ -22,6 +22,7 @@ from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from torchvision import transforms
 from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel, Wav2Vec2ForCTC, Wav2Vec2Processor
+from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -35,7 +36,6 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     load_transformer_config,
-    load_wan_weights_with_optional_gate,
     retrieve_latents,
 )
 from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import (
@@ -1475,4 +1475,6 @@ class Wan22S2VPipeline(
         return outputs
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        return load_wan_weights_with_optional_gate(self, weights)
+        """Load weights using AutoWeightsLoader for vLLM integration."""
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
 VACE (Video Creation and Editing) Pipeline for WAN models.
@@ -22,12 +22,12 @@ import PIL.Image
 import torch
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     Wan22Pipeline,
-    load_wan_weights_with_optional_gate,
     retrieve_latents,
 )
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
@@ -747,4 +747,6 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        return load_wan_weights_with_optional_gate(self, weights)
+        """Load weights using AutoWeightsLoader for vLLM integration."""
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from __future__ import annotations
 
@@ -110,42 +110,6 @@ class WanEulerScheduler:
         self._step_index = None
         self._begin_index = None
 
-    def sigma_for_timestep(
-        self,
-        timestep: float | torch.Tensor,
-        *,
-        device: torch.device | str | None = None,
-        dtype: torch.dtype = torch.float32,
-    ) -> torch.Tensor:
-        """Return the flow sigma nearest to a (possibly distilled) timestep."""
-        timestep_t = torch.as_tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype)
-        if timestep_t.numel() != 1:
-            raise ValueError(f"Expected a scalar timestep, got shape {tuple(timestep_t.shape)}")
-        index = torch.argmin(torch.abs(self.timesteps - timestep_t))
-        return self.sigmas[index].to(device=device or self.timesteps.device, dtype=dtype)
-
-    def predict_clean(
-        self,
-        model_output: torch.Tensor,
-        sample: torch.Tensor,
-        timestep: float | torch.Tensor,
-    ) -> torch.Tensor:
-        """Convert a flow/noise prediction into the corresponding clean latent."""
-        sample_fp32 = sample.to(torch.float32)
-        sigma = self.sigma_for_timestep(timestep, device=sample.device, dtype=torch.float32)
-        return sample_fp32 - sigma * model_output.to(torch.float32)
-
-    def add_noise(
-        self,
-        clean_sample: torch.Tensor,
-        noise: torch.Tensor,
-        timestep: float | torch.Tensor,
-    ) -> torch.Tensor:
-        """Re-noise a clean latent at a distilled flow timestep."""
-        sigma = self.sigma_for_timestep(timestep, device=noise.device, dtype=torch.float32)
-        noisy = (1.0 - sigma) * clean_sample.to(torch.float32) + sigma * noise.to(torch.float32)
-        return noisy.to(noise.dtype)
-
     def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:  # noqa: ARG002
         return sample
 

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import math
 from collections.abc import Iterable
@@ -419,24 +419,6 @@ class WanSelfAttention(nn.Module):
             prefix=prefix,
         )
 
-        # FastVideo VSA checkpoints may add a learned projection that gates
-        # the compressed global branch. Zero initialization makes checkpoints
-        # without these weights sparse-only, with no user-facing mode switch.
-        self.to_gate_compress: ColumnParallelLinear | None = None
-        if self.attn.attn_backend.get_name() == "FASTVIDEO_VSA":
-            self.to_gate_compress = ColumnParallelLinear(
-                dim,
-                self.inner_dim,
-                bias=True,
-                gather_output=False,
-                return_bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.to_gate_compress" if prefix else "to_gate_compress",
-            )
-            nn.init.zeros_(self.to_gate_compress.weight)
-            if self.to_gate_compress.bias is not None:
-                nn.init.zeros_(self.to_gate_compress.bias)
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -459,15 +441,6 @@ class WanSelfAttention(nn.Module):
         key = key.unflatten(2, (self.num_kv_heads, self.head_dim))
         value = value.unflatten(2, (self.num_kv_heads, self.head_dim))
 
-        if self.to_gate_compress is not None:
-            gate_result = self.to_gate_compress(hidden_states)
-            gate_compress = gate_result[0] if isinstance(gate_result, tuple) else gate_result
-            gate_compress = gate_compress.unflatten(2, (self.num_heads, self.head_dim))
-            if attn_metadata is None:
-                attn_metadata = AttentionMetadata(extra={"gate_compress": gate_compress})
-            else:
-                attn_metadata.extra["gate_compress"] = gate_compress
-
         # Apply rotary embeddings
         if rotary_emb is not None:
             freqs_cos, freqs_sin = rotary_emb
@@ -736,8 +709,6 @@ class WanTransformerBlock(nn.Module):
         temb: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
         hidden_states_mask: torch.Tensor | None = None,
-        vsa_dit_seq_shape: tuple[int, int, int] | None = None,
-        preserve_vsa_all_blocks: bool = False,
     ) -> torch.Tensor:
         if temb.ndim == 4:
             # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
@@ -758,12 +729,7 @@ class WanTransformerBlock(nn.Module):
 
         # 1. Self-attention
         norm_hidden_states = self.norm1(hidden_states, scale_msa, shift_msa).type_as(hidden_states)
-        self_attn_extra = {}
-        if vsa_dit_seq_shape is not None:
-            self_attn_extra["vsa_dit_seq_shape"] = vsa_dit_seq_shape
-        if preserve_vsa_all_blocks:
-            self_attn_extra["preserve_vsa_all_blocks"] = True
-        self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask, extra=self_attn_extra)
+        self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask)
         attn_output = self.attn1(norm_hidden_states, rotary_emb, self_attn_metadata)
         hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states)
 
@@ -944,9 +910,6 @@ class WanTransformer3DModel(nn.Module):
             pos_embed_seq_len=pos_embed_seq_len,
         )
 
-        # DMD/FastVideo checkpoints retain VSA semantics when top-k selects every block.
-        self.preserve_vsa_all_blocks = False
-
         # 3. Transformer blocks — partitioned across PP stages via vLLM's `make_layers`.
         # It computes the [start_layer, end_layer) slice for this rank and fills the remaining slots
         # with PPMissingLayer so that weight names stay globally consistent.
@@ -1081,19 +1044,8 @@ class WanTransformer3DModel(nn.Module):
             )
 
         # Transformer blocks
-        # Preserve the post-patch (T, H, W) grid so VSA can partition
-        # the flattened DiT sequence into spatiotemporal blocks.
-        vsa_dit_seq_shape = (post_patch_num_frames, post_patch_height, post_patch_width)
         for block in self.blocks[self.start_layer : self.end_layer]:
-            hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                timestep_proj,
-                rotary_emb,
-                hidden_states_mask,
-                vsa_dit_seq_shape,
-                self.preserve_vsa_all_blocks,
-            )
+            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask)
 
         if not is_pipeline_last_stage():
             # Non-last PP stage: hand the token sequence to the caller via IntermediateTensors.

vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """VACE variant of WanTransformer3DModel for conditional video generation."""
 
@@ -61,7 +61,6 @@ class VaceWanTransformerBlock(WanTransformerBlock):
         temb: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
         hidden_states_mask: torch.Tensor | None = None,
-        vsa_dit_seq_shape: tuple[int, int, int] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if self.proj_in is not None:
             control_hidden_states = self.proj_in(control_hidden_states)
@@ -73,7 +72,6 @@ class VaceWanTransformerBlock(WanTransformerBlock):
             temb,
             rotary_emb,
             hidden_states_mask,
-            vsa_dit_seq_shape,
         )
 
         conditioning_states = self.proj_out(control_hidden_states)
@@ -252,7 +250,6 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
             )
 
         # VACE: embed context and run conditioning blocks
-        vsa_dit_seq_shape = (post_patch_num_frames, post_patch_height, post_patch_width)
         vace_hints = None
         if vace_context is not None and self.vace_blocks is not None:
             full_seq_len = hidden_states.shape[1] * sp_size
@@ -266,7 +263,6 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
                     timestep_proj,
                     rotary_emb,
                     hidden_states_mask,
-                    vsa_dit_seq_shape,
                 )
                 vace_hints.append(conditioning_states)
 
@@ -282,7 +278,6 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
                 timestep_proj,
                 rotary_emb,
                 hidden_states_mask,
-                vsa_dit_seq_shape,
             )
             if vace_hints is not None and self.vace_layers_mapping is not None and i in self.vace_layers_mapping:
                 vace_idx = self.vace_layers_mapping[i]

vllm_omni/diffusion/registry.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import importlib
 
@@ -61,11 +61,6 @@ _DIFFUSION_MODELS = {
         "pipeline_wan2_2",
         "Wan22Pipeline",
     ),
-    "WanDMDPipeline": (
-        "wan2_2",
-        "pipeline_wan2_2",
-        "Wan22Pipeline",
-    ),
     "WanVACEPipeline": (
         "wan2_2",
         "pipeline_wan2_2_vace",
@@ -464,7 +459,6 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
 
     try:
         sp_size = od_config.parallel_config.sequence_parallel_size
-        assert sp_size is not None
         if sp_size <= 1:
             return
 
@@ -545,7 +539,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "OvisImagePipeline": "get_ovis_image_post_process_func",
     "BooguImagePipeline": "get_boogu_image_post_process_func",
     "WanPipeline": "get_wan22_post_process_func",
-    "WanDMDPipeline": "get_wan22_post_process_func",
     "WanVACEPipeline": "get_wan22_vace_post_process_func",
     "LTX2Pipeline": "get_ltx2_post_process_func",
     "LTX2TwoStagePipeline": "get_ltx2_post_process_func",
@@ -619,7 +612,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func",
     "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",
     "WanPipeline": "get_wan22_pre_process_func",
-    "WanDMDPipeline": "get_wan22_pre_process_func",
     "WanVACEPipeline": "get_wan22_vace_pre_process_func",
     "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func",
     "WanS2VPipeline": "get_wan22_s2v_pre_process_func",
@@ -711,7 +703,6 @@ def register_diffusion_model(
 
 def _load_process_func(od_config: OmniDiffusionConfig, func_name: str):
     """Load and return a process function from the appropriate module."""
-    assert od_config.model_class_name is not None
     mod_folder, mod_relname, _ = _DIFFUSION_MODELS[od_config.model_class_name]
     if mod_relname == "":
         # Full module path (registered via register_diffusion_model)

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -95,6 +95,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         #   been read").
         self._held_non_active: deque[Any] = deque()
         self.requests_num_chunks_sent: dict[str, int] = defaultdict(int)
+        # Boundary tasks advance this before the background sender runs, so a
+        # late old-segment frame cannot restore the previous watermark.
+        self._segment_generation: dict[str, int] = defaultdict(int)
         self._pending_streaming_prefills: dict[str, dict] = {}
         # Monotonic timestamp of when each request last began waiting for a
         # chunk, refreshed every time one arrives.  Read by
@@ -188,11 +191,31 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             is_segment_finished: whether the segment of request is finished
         """
         is_finished = request.is_finished() and not request.resumable
+        if not hasattr(self, "_segment_generation"):
+            self._segment_generation = defaultdict(int)
+        external_req_id = request.external_req_id
+        raw_generation = getattr(request, "_omni_segment_generation", 0)
+        try:
+            generation = int(raw_generation)
+        except (TypeError, ValueError):
+            generation = 0
+        expected_generation = self._segment_generation.get(external_req_id, generation)
+        if generation < expected_generation:
+            logger.warning(
+                "Skip late save_async for request %s, segment_generation=%s, expected=%s",
+                external_req_id,
+                generation,
+                expected_generation,
+            )
+            return
+        self._segment_generation[external_req_id] = generation
 
         confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
 
         # If the request is preempted, skip the already saved chunks.
-        if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
+        if not is_segment_finished and confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(
+            request.external_req_id, 0
+        ):
             logger.warning(
                 f"Enqueue save_async for request {request.external_req_id}, "
                 f"request.num_computed_tokens={request.num_computed_tokens}, "
@@ -201,13 +224,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             )
             return
 
-        self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
         task = {
             "multimodal_output": multimodal_output,
             "request": request,
             "is_finished": is_finished,
             "is_segment_finished": is_segment_finished,
+            "segment_generation": generation,
         }
+        if is_segment_finished:
+            self._segment_generation[external_req_id] = generation + 1
+            self.requests_num_chunks_sent.pop(external_req_id, None)
+        else:
+            self.requests_num_chunks_sent[external_req_id] = confirmed_num_computed_tokens
         self._pending_save_reqs.append(task)
         with self._save_cond:
             self._save_cond.notify()
@@ -425,7 +453,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
-            self.requests_num_chunks_sent.pop(external_req_id, None)
             self.ramp_chunk_count.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
             if cached_ic is not None:
@@ -501,6 +528,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.request_payload.pop(external_req_id, None)
         self.code_prompt_token_ids.pop(external_req_id, None)
         self.requests_num_chunks_sent.pop(external_req_id, None)
+        self._segment_generation.pop(external_req_id, None)
         self.ramp_chunk_count.pop(external_req_id, None)
         self._pending_streaming_prefills.pop(external_req_id, None)
 

vllm_omni/engine/arg_utils.py

@@ -1,13 +1,9 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 import argparse
 import json
 import os
 import tempfile
-from collections.abc import Callable
 from dataclasses import dataclass, field, fields
-from typing import Any, cast
+from typing import Any
 
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.logger import init_logger
@@ -302,26 +298,24 @@ class OmniEngineArgs(EngineArgs):
         }
         stage_connector_config["extra"]["stage_id"] = self.stage_id
 
-        hf_overrides = cast(dict[str, Any] | Callable[[Any], Any] | None, getattr(self, "hf_overrides", None))
         # If model_arch is specified, inject it into hf_overrides so vLLM can
         # resolve the architecture even when config.json lacks 'architectures'.
         # Also inject model_type so AutoConfig can resolve the correct config
         # class for models with empty or missing config.json (e.g. CosyVoice3).
         if self.model_arch:
-            if hf_overrides is None:
-                hf_overrides = {}
-                self.hf_overrides = hf_overrides
-            if isinstance(hf_overrides, dict):
-                hf_overrides.setdefault("architectures", [self.model_arch])
-                if "model_type" not in hf_overrides:
+            if self.hf_overrides is None:
+                self.hf_overrides = {}
+            if isinstance(self.hf_overrides, dict):
+                self.hf_overrides.setdefault("architectures", [self.model_arch])
+                if "model_type" not in self.hf_overrides:
                     model_type = _ARCH_TO_MODEL_TYPE.get(self.model_arch)
                     if model_type is not None:
-                        hf_overrides.setdefault("model_type", model_type)
+                        self.hf_overrides.setdefault("model_type", model_type)
 
                 # Stage wrappers (e.g. Code2Wav) may need max_model_len larger
                 # than the base checkpoint's text max_position_embeddings.
                 if self.model_arch == "Qwen3TTSCode2Wav" and self.max_model_len is not None:
-                    hf_overrides.setdefault("talker_config", {}).setdefault(
+                    self.hf_overrides.setdefault("talker_config", {}).setdefault(
                         "max_position_embeddings", int(self.max_model_len)
                     )
 
@@ -337,37 +331,33 @@ class OmniEngineArgs(EngineArgs):
                 if model_type is not None:
                     self._patch_empty_hf_config(model_type)
 
-        tokenizer = cast(str | None, getattr(self, "tokenizer", None))
-
         # NemotronVoiceChat: the checkpoint's only tokenizer subfolder is the
         # auxiliary rnnt_tokenizer/ (ASR loss vocab), which the generic
         # subfolder auto-detect below would wrongly pick up. The text tokenizer
         # is the Nemotron backbone's, resolved from a local override or its
         # HF id (cache-friendly).
-        if not tokenizer and self.model_arch in (
+        if not self.tokenizer and self.model_arch in (
             "NemotronVoiceChatThinkerForConditionalGeneration",
             "NemotronVoiceChatTalkerForConditionalGeneration",
             "NemotronVoiceChatCode2Wav",
         ):
             self.tokenizer = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
-            tokenizer = self.tokenizer
-            logger.info("NemotronVoiceChat: using text tokenizer from %s", tokenizer)
+            logger.info("NemotronVoiceChat: using text tokenizer from %s", self.tokenizer)
 
         # Auto-detect tokenizer for models that store it in a subdirectory
         # rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/).
-        if not tokenizer and self.model:
+        if not self.tokenizer and self.model:
             model_path = self.model
             if os.path.isdir(model_path) and not os.path.isfile(os.path.join(model_path, "tokenizer_config.json")):
                 for subfolder in sorted(os.listdir(model_path)):
                     candidate = os.path.join(model_path, subfolder)
                     if os.path.isdir(candidate) and os.path.isfile(os.path.join(candidate, "tokenizer_config.json")):
                         self.tokenizer = candidate
-                        tokenizer = candidate
                         logger.info("Auto-detected tokenizer at %s", candidate)
                         break
             elif not os.path.isdir(model_path):
-                tokenizer_subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch or "")
-                if tokenizer_subfolder:
+                subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch)
+                if subfolder:
                     # Download just the tokenizer files from the subfolder
                     try:
                         from huggingface_hub import snapshot_download
@@ -375,15 +365,15 @@ class OmniEngineArgs(EngineArgs):
                         local_dir = snapshot_download(
                             model_path,
                             allow_patterns=[
-                                f"{tokenizer_subfolder}/tokenizer*",
-                                f"{tokenizer_subfolder}/tokenization_*",
-                                f"{tokenizer_subfolder}/special_tokens*",
-                                f"{tokenizer_subfolder}/vocab*",
-                                f"{tokenizer_subfolder}/merges*",
-                                f"{tokenizer_subfolder}/added_tokens*",
+                                f"{subfolder}/tokenizer*",
+                                f"{subfolder}/tokenization_*",
+                                f"{subfolder}/special_tokens*",
+                                f"{subfolder}/vocab*",
+                                f"{subfolder}/merges*",
+                                f"{subfolder}/added_tokens*",
                             ],
                         )
-                        candidate = os.path.join(local_dir, tokenizer_subfolder)
+                        candidate = os.path.join(local_dir, subfolder)
                         if os.path.isdir(candidate):
                             self.tokenizer = candidate
                             logger.info("Downloaded tokenizer from %s/%s", model_path, subfolder)
@@ -542,7 +532,6 @@ class OrchestratorArgs:
     hsdp_shard_size: int = -1
     hsdp_replicate_size: int = 1
     diffusion_attention_backend: str | None = None
-    fastvideo_vsa_topk: int | None = None
     diffusion_attention_config: str | None = None
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 """
 Async Omni Engine for vLLM-Omni multi-stage runtime.
 
@@ -378,13 +375,9 @@ class AsyncOmniEngine:
                     getattr(self, "_duplex_runtime_extension_path", None)
                 )
                 if duplex_runtime_extension is not None:
-                    stage_clients = []
-                    for pool in self.stage_pools:
-                        assert pool.stage_client is not None
-                        stage_clients.append(pool.stage_client)
                     validate_duplex_runtime_extension(
                         duplex_runtime_extension,
-                        sampling_defaults=tuple(client.default_sampling_params for client in stage_clients),
+                        sampling_defaults=tuple(pool.stage_client.default_sampling_params for pool in self.stage_pools),
                     )
 
             orchestrator = Orchestrator(
@@ -692,7 +685,6 @@ class AsyncOmniEngine:
             )
 
             # Full input processing (tokenization, multimodal, etc.)
-            assert self.input_processor is not None
             _t_preprocess = time.perf_counter()
             try:
                 request = self.input_processor.process_inputs(
@@ -771,13 +763,11 @@ class AsyncOmniEngine:
             Exception: Whatever prompt expansion or input processing raised.
                 The caller is expected to let it reach the client.
         """
-        assert self.prompt_expand_func is not None
         expanded = self.prompt_expand_func(original_prompt, stage0_params)
         if not expanded:
             return []
 
         companions: list[AddCompanionRequestMessage] = []
-        assert self.input_processor is not None
         for ep in expanded:
             cid = f"{parent_id}{ep.request_id_suffix}"
             companion_prompt = ep.prompt
@@ -874,7 +864,7 @@ class AsyncOmniEngine:
                 "mag_max_skip_steps": 5,
                 "mag_retention_ratio": 0.1,
             }
-        if cache_backend == "step_cache":
+        if cache_backend in ("step_cache"):
             return {
                 "step_cache_dit_enabled": True,
                 "velocity_sim_thresholds": [0.95, 0.93],
@@ -1028,12 +1018,10 @@ class AsyncOmniEngine:
         if (
             kwargs.get("diffusion_attention_config") is not None
             or kwargs.get("diffusion_attention_backend") is not None
-            or kwargs.get("fastvideo_vsa_topk") is not None
         ):
             attention_config = parse_attention_config(
                 kwargs.get("diffusion_attention_config"),
                 attention_backend=kwargs.get("diffusion_attention_backend"),
-                fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
             )
 
         stage_engine_args = {
@@ -1235,7 +1223,6 @@ class AsyncOmniEngine:
                 if (
                     kwargs.get("diffusion_attention_config") is not None
                     or kwargs.get("diffusion_attention_backend") is not None
-                    or kwargs.get("fastvideo_vsa_topk") is not None
                 ):
                     has_stage_attention = (
                         hasattr(cfg.engine_args, "diffusion_attention_config")
@@ -1245,7 +1232,6 @@ class AsyncOmniEngine:
                         cfg.engine_args.diffusion_attention_config = parse_attention_config(
                             kwargs.get("diffusion_attention_config"),
                             attention_backend=kwargs.get("diffusion_attention_backend"),
-                            fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
                         )
                 quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
                 if quantization_config is not None:

vllm_omni/engine/orchestrator.py

@@ -199,17 +199,37 @@ class OrchestratorRequestState:
 
 
 @dataclass
-class StreamingInputState:
-    # Flag of streaming input request
-    enabled: bool = False
+class StreamingSegmentState:
+    """Streaming segment-boundary state for one stage of a request.
+
+    Every stage of a multi-stage pipeline reaches its own segment boundaries
+    independently, so this is tracked per stage rather than per request.
+    """
+
     # Flag of segment of streaming input finished
-    segment_finished: bool = False
+    finished: bool = False
     # Tokens from the current raw segment boundary. The vLLM output processor
     # does not guarantee that EngineCoreOutput.new_token_ids survives on the
     # processed RequestOutput used by the routing layer.
-    segment_token_ids: list[int] = field(default_factory=list)
-    segment_output_metadata: dict[str, Any] = field(default_factory=dict)
-    # Streaming update prompt length
+    token_ids: list[int] = field(default_factory=list)
+    output_metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class StreamingInputState:
+    # Flag of streaming input request
+    enabled: bool = False
+    # Keyed by stage. ``_orchestration_loop`` records a boundary per stage poll
+    # into this one ``req_state``, while every consumer reads it for a specific
+    # ``stage_id``, so a flat slot is only correct as long as each read happens in
+    # the same iteration as its own stage's write. It is not for a stage that
+    # never records one (the ``poll_diffusion_output`` branch reaches
+    # ``_handle_processed_outputs`` without the raw-output loop) or for any
+    # consumer reading outside the writing iteration.
+    segments: dict[int, StreamingSegmentState] = field(default_factory=dict)
+    # Streaming update prompt length. NOT stage-keyed: its only consumer reads it
+    # through the streaming context on the stage-input-bridge path, not by stage
+    # id. Do not assume the rest of this struct is stage-safe.
     new_prompt_len_snapshot: int | None = None
     # Model/bridge-specific runtime states (e.g., thinker->talker)
     bridge_states: dict[str, Any] = field(default_factory=dict)
@@ -217,6 +237,14 @@ class StreamingInputState:
     # while the downstream input processor consumes upstream token output.
     source_token_decoder: Callable[..., str] | None = None
 
+    def segment(self, stage_id: int | None) -> StreamingSegmentState:
+        """Return ``stage_id``'s segment state, or a fresh empty one if unreported.
+
+        Fresh rather than shared: ``output_metadata`` is handed out by reference.
+        """
+        segment = self.segments.get(stage_id)
+        return segment if segment is not None else StreamingSegmentState()
+
 
 class _OrchestratorDuplexStagePort:
     """Adapts generic stage pools to the model-neutral duplex control plane."""
@@ -966,17 +994,18 @@ class Orchestrator:
                                 req_state = self.request_states.get(getattr(eco, "request_id", None))
                                 if req_state is None or not req_state.streaming.enabled:
                                     continue
-                                req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False))
-                                req_state.streaming.segment_token_ids = (
-                                    self._coerce_int_list(getattr(eco, "new_token_ids", None))
-                                    if req_state.streaming.segment_finished
-                                    else []
-                                )
+                                segment_finished = bool(getattr(eco, "is_segment_finished", False))
                                 raw_mm = self._completion_multimodal_output(eco, None)
-                                req_state.streaming.segment_output_metadata = (
-                                    dict(raw_mm)
-                                    if req_state.streaming.segment_finished and isinstance(raw_mm, dict)
-                                    else {}
+                                req_state.streaming.segments[stage_id] = StreamingSegmentState(
+                                    finished=segment_finished,
+                                    token_ids=(
+                                        self._coerce_int_list(getattr(eco, "new_token_ids", None))
+                                        if segment_finished
+                                        else []
+                                    ),
+                                    output_metadata=(
+                                        dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}
+                                    ),
                                 )
                                 req_state.streaming.new_prompt_len_snapshot = getattr(
                                     eco,
@@ -1044,7 +1073,7 @@ class Orchestrator:
                 continue
 
             stage_metrics = None
-            segment_finished = req_state.streaming.enabled and req_state.streaming.segment_finished
+            segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished
             if output.finished or segment_finished:
                 stage_metrics = pool.build_stage_metrics(
                     [output],
@@ -1378,6 +1407,7 @@ class Orchestrator:
         req_id = output.request_id
         finished = output.finished
         submit_ts = req_state.stage_submit_ts.get(stage_id)
+        segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished
         # CFG companion: stash output so parent can bundle [parent, *companions]
         # into source_outputs for the bridge (e.g. thinker2imagegen).
         if finished and self._cfg_tracker.is_companion(req_id):
@@ -1387,11 +1417,7 @@ class Orchestrator:
             return
 
         request_finished = False
-        if (
-            finished
-            and self.stage_pools[stage_id].final_output
-            and not (req_state.streaming.enabled and req_state.streaming.segment_finished)
-        ):
+        if finished and self.stage_pools[stage_id].final_output and not segment_finished:
             req_state.finished_final_output_stage_ids.add(stage_id)
             final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
             request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids)
@@ -1403,7 +1429,7 @@ class Orchestrator:
         # filter out again (the official implementation returns exactly one
         # result per audio chunk).
         is_duplex_stage0_segment = (
-            stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished
+            stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment(0).finished
         )
         if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment:
             await self.output_async_queue.put(
@@ -1415,7 +1441,10 @@ class Orchestrator:
                     metrics=stage_metrics,
                     finished=(
                         request_finished
-                        or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished)
+                        or (
+                            self._is_duplex_session_request(req_state)
+                            and req_state.streaming.segment(stage_id).finished
+                        )
                     ),
                     stage_submit_ts=submit_ts,
                 )
@@ -1450,7 +1479,7 @@ class Orchestrator:
             return
 
         if (
-            (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished))
+            (finished or segment_finished)
             and stage_id < req_state.final_stage_id
             and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1))
             and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled)
@@ -1589,7 +1618,9 @@ class Orchestrator:
     def _duplex_output_context(
         req_state: OrchestratorRequestState,
         *,
-        stage_id: int | None = None,
+        # Required: segment state is stage-keyed, so a caller that omitted this
+        # would silently read "no boundary" rather than the intended stage's.
+        stage_id: int | None,
     ) -> DuplexOutputContext | None:
         identity = req_state.duplex_identity
         if identity is None:
@@ -1600,15 +1631,16 @@ class Orchestrator:
         )
 
         fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence
+        segment = req_state.streaming.segment(stage_id)
         return DuplexOutputContext(
             identity=DuplexRequestIdentity(
                 session_id=identity.session_id,
                 fence=fence,
             ),
             final_stage_id=req_state.final_stage_id,
-            segment_finished=req_state.streaming.enabled and req_state.streaming.segment_finished,
-            segment_token_ids=tuple(req_state.streaming.segment_token_ids),
-            segment_output_metadata=req_state.streaming.segment_output_metadata,
+            segment_finished=req_state.streaming.enabled and segment.finished,
+            segment_token_ids=tuple(segment.token_ids),
+            segment_output_metadata=segment.output_metadata,
         )
 
     @staticmethod

vllm_omni/entrypoints/cli/serve.py

@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
 """
 Omni serve command for vLLM-Omni.
 
@@ -609,12 +606,6 @@ class OmniServeCommand(CLISubcommand):
             "May be combined with --diffusion-attention-config.per_role.* overrides, "
             "but mutually exclusive with --diffusion-attention-config.default.backend.",
         )
-        omni_config_group.add_argument(
-            "--fastvideo-vsa-topk",
-            type=int,
-            default=None,
-            help="Number of key/value blocks selected per query block by FASTVIDEO_VSA.",
-        )
         omni_config_group.add_argument(
             "--diffusion-attention-config",
             "-dac",

vllm_omni/model_executor/models/wan2_2/pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Wan2.2 TI2V pipeline topology."""
 
 from vllm_omni.config.stage_config import (
@@ -12,7 +12,6 @@ WAN2_2_TI2V_PIPELINE = PipelineConfig(
     model_type="wan2_2_ti2v",
     model_arch="WanPipeline",
     diffusers_class_name="WanPipeline",
-    diffusers_class_aliases=("WanDMDPipeline",),
     stages=(
         StagePipelineConfig(
             stage_id=0,