diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md index 18e05517c..41cbaf3ff 100644 --- a/docs/design/feature/attention_backend_selection.md +++ b/docs/design/feature/attention_backend_selection.md @@ -61,7 +61,8 @@ 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. +- `block_sparse` is consumed by block-sparse backends such as RainFusion; and +- `fastvideo_vsa_topk` is consumed by FastVideo VSA. 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 @@ -87,3 +88,19 @@ 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. diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 6897d633e..e2a3b9bfc 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -20,6 +20,7 @@ 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 @@ -35,6 +36,7 @@ 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 @@ -105,6 +107,7 @@ 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 diff --git a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md new file mode 100644 index 000000000..2439633d4 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md @@ -0,0 +1,125 @@ +# 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 \ + --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 +``` + +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 \ + --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. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index ac5e5aea9..374a6d212 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1304,6 +1304,7 @@ 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( diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index f45f62ec1..0db9803a6 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the additive structured Omni config.""" from __future__ import annotations @@ -1152,6 +1152,35 @@ 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( diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 4a2c65cbb..7d0b2a5fa 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for per-role attention backend configuration (RFC: per-role-attention-backend). @@ -90,6 +90,14 @@ 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 diff --git a/tests/diffusion/attention/test_fastvideo_vsa.py b/tests/diffusion/attention/test_fastvideo_vsa.py new file mode 100644 index 000000000..d6209bf42 --- /dev/null +++ b/tests/diffusion/attention/test_fastvideo_vsa.py @@ -0,0 +1,203 @@ +# 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" diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 325bd59da..de5abeab0 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import importlib from contextlib import contextmanager from types import SimpleNamespace @@ -9,6 +10,7 @@ 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 @@ -82,6 +84,7 @@ 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 @@ -94,7 +97,7 @@ def _make_pipeline() -> Wan22Pipeline: def _make_sampling(**overrides): - values = { + values: dict[str, object] = { "height": None, "width": None, "num_frames": 1, @@ -338,3 +341,125 @@ 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 diff --git a/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py b/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py new file mode 100644 index 000000000..8c9208014 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py @@ -0,0 +1,32 @@ +# 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])) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 28667227d..64c1ce99a 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace @@ -400,7 +400,9 @@ def test_serve_cli_accepts_diffusion_attention_backend(): "Qwen/Qwen-Image", "--omni", "--diffusion-attention-backend", - "FLASH_ATTN", + "FASTVIDEO_VSA", + "--fastvideo-vsa-topk", + "96", ] ) @@ -408,10 +410,12 @@ 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 == "FLASH_ATTN" + assert args.diffusion_attention_backend == "FASTVIDEO_VSA" + assert args.fastvideo_vsa_topk == 96 assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None - assert diffusion_attention_config.default.backend == "FLASH_ATTN" + assert diffusion_attention_config.default.backend == "FASTVIDEO_VSA" + assert diffusion_attention_config.default.backend_kwargs() == {"topk": 96} def test_serve_cli_accepts_request_batch_max_wait_ms(): diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index 5adf218c7..48bd0b476 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Config factories for vllm-omni, e.g., StageConfigFactory.""" from __future__ import annotations @@ -226,7 +226,10 @@ 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 pipeline_cfg.diffusers_class_name == class_name: + if pipeline_cfg is not None and class_name in ( + pipeline_cfg.diffusers_class_name, + *pipeline_cfg.diffusers_class_aliases, + ): logger.info( "Detected pipeline %r from model_index.json (_class_name=%r)", pipeline_cfg.model_type, @@ -433,6 +436,7 @@ 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: @@ -440,9 +444,12 @@ 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) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 5a1dbb12f..6f3a2ab65 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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 - cls.__match_args__ = () + setattr(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 + is_encoder_decoder: InitVar[bool] = False # type: ignore[assignment] 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, - self.all2all_backend, - self.disable_custom_all_reduce, + getattr(self, "all2all_backend", None), + getattr(self, "disable_custom_all_reduce", None), ) 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(cls)} + valid_fields = {f.name for f in fields(cast(Any, 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(_DiffusionConfigProjection)) +_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(cast(Any, _DiffusionConfigProjection))) # Current OmniDiffusionConfig still contains a flat mix of shared engine, # runtime, parallel, and diffusion-specific knobs. Keep this classification @@ -956,6 +956,7 @@ _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", @@ -1140,7 +1141,7 @@ def _global_stage_cli_fields() -> frozenset[str]: ) externally_consumed = ( _NON_STAGE_ENGINE_CLI_FIELDS - | frozenset(f.name for f in fields(VllmOmniOrchestratorConfig)) + | frozenset(f.name for f in fields(cast(Any, VllmOmniOrchestratorConfig))) | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS) ) return candidates - externally_consumed @@ -1238,7 +1239,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(VllmOmniOrchestratorConfig): + for config_field in fields(cast(Any, VllmOmniOrchestratorConfig)): name = config_field.name if name == "deploy_config_path": continue @@ -1443,8 +1444,8 @@ def _with_resolved_processors( input_proc: str | None, next_stage_proc: str | None, ) -> StageConfigType: - stage_config._resolved_custom_process_input_func = input_proc - stage_config._resolved_custom_process_next_stage_input_func = next_stage_proc + setattr(stage_config, "_resolved_custom_process_input_func", input_proc) + setattr(stage_config, "_resolved_custom_process_next_stage_input_func", next_stage_proc) return stage_config @@ -1558,13 +1559,16 @@ 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 builder( - pipeline, - deploy, - topology, - stage_deploy, - engine, - model=model, + return cast( + StageConfigType, + builder( + pipeline, + deploy, + topology, + stage_deploy, + engine, + model=model, + ), ) @@ -1608,7 +1612,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 OmniStageModelConfig( + return cast(Any, OmniStageModelConfig)( default_sampling_params=default_sampling_params, duplex_max_sessions=duplex_max_sessions, **kwargs, @@ -1699,7 +1703,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 OmniStageConnectorConfig( + return cast(Any, 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, @@ -1848,11 +1852,11 @@ class VllmOmniConfig: for topology in pipeline_cfg.stages ) - orchestrator_config = VllmOmniOrchestratorConfig( + orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)( deploy_config_path=loaded_deploy_config_path, **_orchestrator_cli_overrides(cli_overrides), ) - return cls( + return cast(Any, cls)( pipeline_config=pipeline_cfg, stage_configs=stage_configs, orchestrator_config=orchestrator_config, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 70ea7726a..fdfd7c288 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage configuration system for vLLM-Omni.""" from __future__ import annotations @@ -234,6 +234,7 @@ 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 @@ -356,6 +357,7 @@ 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. diff --git a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py new file mode 100644 index 000000000..fcd592dd8 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py @@ -0,0 +1,431 @@ +# 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) diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index fdfaf5f9c..1af079696 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/vllm_omni/diffusion/attention/backends/registry.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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": + def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum": # type: ignore[override] """Get backend by name with helpful error messages.""" try: return super().__getitem__(name) # type: ignore[return-value] @@ -62,6 +62,7 @@ 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" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9adda521a..0f6ae494a 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import copy import math import os @@ -77,11 +77,13 @@ 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) - if diffusion_attn_backend is not None: + fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None) + if diffusion_attn_backend is not None or fastvideo_vsa_topk 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. @@ -104,6 +106,7 @@ 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)): @@ -112,7 +115,7 @@ def parse_kv_cache_skip_selector( text = selector.strip() if not text: return None - values: set[int] = set() + values = set() for chunk in text.split(","): token = chunk.strip() if not token: @@ -245,6 +248,7 @@ 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" @@ -362,6 +366,7 @@ 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 @@ -1063,6 +1068,7 @@ 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: @@ -1262,6 +1268,7 @@ 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, @@ -1338,6 +1345,7 @@ 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") @@ -1395,6 +1403,7 @@ 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()) @@ -1593,10 +1602,11 @@ 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 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: + 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: raise ValueError( - f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}." + f"quant.{block_name}={block_value!r} unsupported; " + f"kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}." ) @property @@ -1638,6 +1648,7 @@ 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) @@ -1657,6 +1668,11 @@ 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. @@ -1698,12 +1714,15 @@ 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 @@ -1773,7 +1792,7 @@ class AttentionConfig: normalized[role] = node return - spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"} + spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse"} node_dict = dict(node) node_keys = set(node_dict) if node_keys & spec_keys: @@ -1812,6 +1831,7 @@ 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. @@ -1839,6 +1859,14 @@ 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 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 7ec6c3060..f92053e52 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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 WanTransformer3DModel +from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, 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,6 +44,8 @@ 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: @@ -64,6 +66,30 @@ 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) @@ -314,6 +340,7 @@ 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") @@ -321,6 +348,7 @@ 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) @@ -333,6 +361,7 @@ 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 @@ -422,6 +451,10 @@ 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 @@ -430,8 +463,14 @@ class Wan22Pipeline( else: raise RuntimeError("No transformer loaded") - self._sample_solver = "unipc" - self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0 + 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.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 @@ -480,6 +519,7 @@ 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 = {} @@ -561,7 +601,20 @@ class Wan22Pipeline( cfg_normalize=False, ) - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + 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) pbar.update() return latents @@ -600,13 +653,21 @@ 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 - num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps + 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 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=4.0) + guidance_low, guidance_high = resolve_wan_guidance_scales( + common, + default_guidance_scale=1.0 if self.is_dmd else 4.0, + ) # record guidance for properties self._guidance_scale = guidance_low @@ -685,17 +746,19 @@ class Wan22Pipeline( current_omni_platform.synchronize() _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000 - 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 + 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 - # Timesteps - self.scheduler.set_timesteps(num_steps, device=device) - timesteps = self.scheduler.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: @@ -817,6 +880,7 @@ 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 @@ -1003,9 +1067,7 @@ class Wan22Pipeline( return latents def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) def check_inputs( self, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 8e149d389..b0cd35458 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -16,7 +16,6 @@ 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 @@ -36,6 +35,7 @@ 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,9 +1014,7 @@ 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]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) # --------------------------------------------------------------------------- diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index bce71633c..bc688fbc4 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Wan2.2 Speech-to-Video (S2V) Pipeline for vLLM-Omni. @@ -22,7 +22,6 @@ 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 @@ -36,6 +35,7 @@ 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,6 +1475,4 @@ class Wan22S2VPipeline( return outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index ae869fd57..58cd98c3e 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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,6 +747,4 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) diff --git a/vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py b/vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py index 25444044c..4a365c2be 100644 --- a/vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -110,6 +110,42 @@ 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 diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index d61d1c82b..dbc555604 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math from collections.abc import Iterable @@ -419,6 +419,24 @@ 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, @@ -441,6 +459,15 @@ 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 @@ -709,6 +736,8 @@ 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) @@ -729,7 +758,12 @@ class WanTransformerBlock(nn.Module): # 1. Self-attention norm_hidden_states = self.norm1(hidden_states, scale_msa, shift_msa).type_as(hidden_states) - self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask) + 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) attn_output = self.attn1(norm_hidden_states, rotary_emb, self_attn_metadata) hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states) @@ -910,6 +944,9 @@ 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. @@ -1044,8 +1081,19 @@ 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) + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + hidden_states_mask, + vsa_dit_seq_shape, + self.preserve_vsa_all_blocks, + ) if not is_pipeline_last_stage(): # Non-last PP stage: hand the token sequence to the caller via IntermediateTensors. diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index a13c04c99..7d04580e4 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """VACE variant of WanTransformer3DModel for conditional video generation.""" @@ -61,6 +61,7 @@ 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) @@ -72,6 +73,7 @@ class VaceWanTransformerBlock(WanTransformerBlock): temb, rotary_emb, hidden_states_mask, + vsa_dit_seq_shape, ) conditioning_states = self.proj_out(control_hidden_states) @@ -250,6 +252,7 @@ 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 @@ -263,6 +266,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel): timestep_proj, rotary_emb, hidden_states_mask, + vsa_dit_seq_shape, ) vace_hints.append(conditioning_states) @@ -278,6 +282,7 @@ 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] diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 5178371fc..921ceca1b 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import importlib @@ -61,6 +61,11 @@ _DIFFUSION_MODELS = { "pipeline_wan2_2", "Wan22Pipeline", ), + "WanDMDPipeline": ( + "wan2_2", + "pipeline_wan2_2", + "Wan22Pipeline", + ), "WanVACEPipeline": ( "wan2_2", "pipeline_wan2_2_vace", @@ -459,6 +464,7 @@ 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 @@ -539,6 +545,7 @@ _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", @@ -612,6 +619,7 @@ _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", @@ -703,6 +711,7 @@ 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) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 152e5b3ee..f1c483d48 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,9 +1,13 @@ +# 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 +from typing import Any, cast from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.logger import init_logger @@ -298,24 +302,26 @@ 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 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: + 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: model_type = _ARCH_TO_MODEL_TYPE.get(self.model_arch) if model_type is not None: - self.hf_overrides.setdefault("model_type", model_type) + 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: - self.hf_overrides.setdefault("talker_config", {}).setdefault( + hf_overrides.setdefault("talker_config", {}).setdefault( "max_position_embeddings", int(self.max_model_len) ) @@ -331,33 +337,37 @@ 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 self.tokenizer and self.model_arch in ( + if not 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" - logger.info("NemotronVoiceChat: using text tokenizer from %s", self.tokenizer) + tokenizer = self.tokenizer + logger.info("NemotronVoiceChat: using text tokenizer from %s", tokenizer) # Auto-detect tokenizer for models that store it in a subdirectory # rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/). - if not self.tokenizer and self.model: + if not 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): - subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch) - if subfolder: + tokenizer_subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch or "") + if tokenizer_subfolder: # Download just the tokenizer files from the subfolder try: from huggingface_hub import snapshot_download @@ -365,15 +375,15 @@ class OmniEngineArgs(EngineArgs): local_dir = snapshot_download( model_path, allow_patterns=[ - f"{subfolder}/tokenizer*", - f"{subfolder}/tokenization_*", - f"{subfolder}/special_tokens*", - f"{subfolder}/vocab*", - f"{subfolder}/merges*", - f"{subfolder}/added_tokens*", + 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*", ], ) - candidate = os.path.join(local_dir, subfolder) + candidate = os.path.join(local_dir, tokenizer_subfolder) if os.path.isdir(candidate): self.tokenizer = candidate logger.info("Downloaded tokenizer from %s/%s", model_path, subfolder) @@ -532,6 +542,7 @@ 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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 12236294e..af366b2c5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Async Omni Engine for vLLM-Omni multi-stage runtime. @@ -375,9 +378,13 @@ 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(pool.stage_client.default_sampling_params for pool in self.stage_pools), + sampling_defaults=tuple(client.default_sampling_params for client in stage_clients), ) orchestrator = Orchestrator( @@ -685,6 +692,7 @@ 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( @@ -763,11 +771,13 @@ 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 @@ -864,7 +874,7 @@ class AsyncOmniEngine: "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1, } - if cache_backend in ("step_cache"): + if cache_backend == "step_cache": return { "step_cache_dit_enabled": True, "velocity_sim_thresholds": [0.95, 0.93], @@ -1018,10 +1028,12 @@ 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 = { @@ -1223,6 +1235,7 @@ 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") @@ -1232,6 +1245,7 @@ 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: diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index c104f1e72..122cc862b 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Omni serve command for vLLM-Omni. @@ -606,6 +609,12 @@ 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", diff --git a/vllm_omni/model_executor/models/wan2_2/pipeline.py b/vllm_omni/model_executor/models/wan2_2/pipeline.py index 7059845fa..2d2bc3064 100644 --- a/vllm_omni/model_executor/models/wan2_2/pipeline.py +++ b/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 project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Wan2.2 TI2V pipeline topology.""" from vllm_omni.config.stage_config import ( @@ -12,6 +12,7 @@ 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,