From 094abd70c1d98edd965d5e2fd392993318092bd3 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 11 Aug 2026 08:24:32 +0800 Subject: [PATCH] feat(diffusion): add request-gated VAE fast paths Signed-off-by: Hongsheng Liu --- .../test_diffusion_config_propagation.py | 8 + .../diffusion/test_diffusion_model_runner.py | 62 ++- tests/diffusion/test_diffusion_worker.py | 2 + .../vae_optimizations/test_vae_fastpaths.py | 300 +++++++++++++++ tests/entrypoints/test_serve.py | 11 + vllm_omni/config/omni_config.py | 1 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 8 + .../autoencoders/autoencoder_kl_wan.py | 2 +- .../diffusion/vae_optimizations/__init__.py | 19 + vllm_omni/diffusion/vae_optimizations/gate.py | 67 ++++ .../diffusion/vae_optimizations/image.py | 354 ++++++++++++++++++ .../diffusion/vae_optimizations/optimize.py | 70 ++++ .../diffusion/vae_optimizations/precision.py | 91 +++++ .../triton_group_norm_silu.py | 237 ++++++++++++ .../triton_wan_rmsnorm_silu.py | 160 ++++++++ vllm_omni/diffusion/vae_optimizations/wan.py | 162 ++++++++ .../worker/diffusion_model_runner.py | 21 +- .../diffusion/worker/diffusion_worker.py | 8 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 9 + 22 files changed, 1591 insertions(+), 4 deletions(-) create mode 100644 tests/diffusion/vae_optimizations/test_vae_fastpaths.py create mode 100644 vllm_omni/diffusion/vae_optimizations/__init__.py create mode 100644 vllm_omni/diffusion/vae_optimizations/gate.py create mode 100644 vllm_omni/diffusion/vae_optimizations/image.py create mode 100644 vllm_omni/diffusion/vae_optimizations/optimize.py create mode 100644 vllm_omni/diffusion/vae_optimizations/precision.py create mode 100644 vllm_omni/diffusion/vae_optimizations/triton_group_norm_silu.py create mode 100644 vllm_omni/diffusion/vae_optimizations/triton_wan_rmsnorm_silu.py create mode 100644 vllm_omni/diffusion/vae_optimizations/wan.py diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index e25d6497c..7f6d36aae 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -113,6 +113,14 @@ class TestCreateDefaultDiffusion: stages = StageConfigFactory.create_default_diffusion({"model": "x"}) assert stages[0]["engine_args"]["cache_backend"] == "none" + def test_vae_decode_precision_roundtrip(self): + od = _roundtrip_diffusion_config(model="x", vae_decode_precision="bf16") + assert od.vae_decode_precision == "bf16" + + def test_invalid_vae_decode_precision_is_rejected(self): + with pytest.raises(ValueError, match="vae_decode_precision"): + OmniDiffusionConfig(model="x", vae_decode_precision="int8") + def test_single_gpu_default_devices(self): stages = StageConfigFactory.create_default_diffusion({"model": "x"}) assert stages[0]["runtime"]["devices"] == "0" diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index a1f8440e3..5060c702e 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -512,6 +512,35 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin assert runner.pipeline.last_req.num_reqs == 1 +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("quality", "expected_enabled"), + [(None, False), ("lossless", False), ("high", True)], +) +def test_execute_model_scopes_vae_fast_path_to_request_quality(monkeypatch, quality, expected_enabled): + runner = _make_runner(cache_backend=None, cache_backend_name="none") + req = _make_request() + req.sampling_params.quality = quality + gate_calls = [] + + @contextmanager + def _record_gate(pipeline, enabled): + gate_calls.append((pipeline, enabled, "enter")) + yield + gate_calls.append((pipeline, enabled, "exit")) + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "use_pipeline_vae_fast_path", _record_gate) + + DiffusionModelRunner.execute_model(runner, req) + + assert gate_calls == [ + (runner.pipeline, expected_enabled, "enter"), + (runner.pipeline, expected_enabled, "exit"), + ] + + class _BatchPipeline: """Pipeline returning a configurable list of outputs from forward().""" @@ -579,6 +608,18 @@ def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch): DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_rejects_mixed_vae_quality() -> None: + runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")])) + sched = _make_scheduler_output(num_reqs=2) + sched.scheduled_new_reqs[0].req.sampling_params.quality = "lossless" + sched.scheduled_new_reqs[1].req.sampling_params.quality = "high" + + with pytest.raises(RuntimeError, match="cannot mix request quality levels"): + DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) + + @pytest.mark.core_model @pytest.mark.cpu def test_execute_model_batch_routes_one_output_per_request(monkeypatch): @@ -728,6 +769,19 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch): @pytest.mark.core_model @pytest.mark.cpu def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): + memory_pool_active = False + optimization_pool_states = [] + + @contextmanager + def _memory_pool_context(tag): + nonlocal memory_pool_active + assert tag == "weights" + memory_pool_active = True + try: + yield + finally: + memory_pool_active = False + class _DummyLoader: def __init__(self, load_config, od_config=None): del load_config, od_config @@ -775,16 +829,22 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) + monkeypatch.setattr( + model_runner_module, + "optimize_pipeline_vaes", + lambda pipeline, od_config: optimization_pool_states.append(memory_pool_active), + ) monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None) monkeypatch.setattr( model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend ) - DiffusionModelRunner.load_model(runner) + DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context) assert runner.cache_backend is None assert runner.od_config.cache_backend is None assert dummy_cache_backend.enabled is False + assert optimization_pool_states == [True] @pytest.mark.core_model diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 6a49d0909..82537550e 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -102,6 +102,7 @@ class TestDiffusionWorkerSleep: def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" + clear_vae_caches = mocker.patch("vllm_omni.diffusion.vae_optimizations.clear_pipeline_vae_fast_path_caches") mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3] @@ -128,6 +129,7 @@ class TestDiffusionWorkerSleep: # Verify sleep was called with correct tags mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",)) + clear_vae_caches.assert_called_once_with(mock_gpu_worker.model_runner.pipeline) assert bool(result) is True # Verify buffers were NOT saved (level 1 doesn't save buffers) assert len(mock_gpu_worker._sleep_saved_buffers) == 0 diff --git a/tests/diffusion/vae_optimizations/test_vae_fastpaths.py b/tests/diffusion/vae_optimizations/test_vae_fastpaths.py new file mode 100644 index 000000000..813ac84bb --- /dev/null +++ b/tests/diffusion/vae_optimizations/test_vae_fastpaths.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness contracts for request-gated VAE decoder fast paths.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.vae_optimizations.gate import ( + VaeFastPathGate, + register_vae_fast_path_gate, + use_pipeline_vae_fast_path, + use_vae_fast_path, +) + +pytestmark = [pytest.mark.diffusion] + + +class _PipelineWithVae(nn.Module): + def __init__(self, vae: nn.Module) -> None: + super().__init__() + self.vae = vae + + +class _DistributedExecutorConfig: + def __init__(self, parallel_mode: str) -> None: + self.parallel_mode = parallel_mode + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_vae_fast_path_gate_restores_nested_state() -> None: + vae = nn.Identity() + gate = VaeFastPathGate() + register_vae_fast_path_gate(vae, gate) + + with use_vae_fast_path(vae, True): + assert gate.enabled + with use_vae_fast_path(vae, False): + assert not gate.enabled + assert gate.enabled + + assert not gate.enabled + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_pipeline_vae_fast_path_controls_discovered_vae() -> None: + vae = nn.Identity() + pipeline = _PipelineWithVae(vae) + gate = VaeFastPathGate() + register_vae_fast_path_gate(vae, gate) + + with use_pipeline_vae_fast_path(pipeline, True): + assert gate.enabled + + assert not gate.enabled + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_image_group_norm_fast_path_uses_compile_safe_fallback(monkeypatch) -> None: + from vllm_omni.diffusion.vae_optimizations import image + + gate = VaeFastPathGate() + gate.enabled = True + norm = image.FusedGroupNormSiLU(nn.GroupNorm(2, 4), gate) + inputs = torch.randn(1, 4, 3, 3) + + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + monkeypatch.setattr( + image, + "group_norm_silu_4d", + lambda *args, **kwargs: pytest.fail("eager Triton kernel entered during compilation"), + ) + + actual = norm(inputs) + expected = F.silu(F.group_norm(inputs, 2, norm.weight, norm.bias, norm.eps)) + torch.testing.assert_close(actual, expected) + + +def _small_image_vae() -> nn.Module: + from diffusers.models.autoencoders import AutoencoderKL + + return AutoencoderKL( + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"), + up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"), + block_out_channels=(32, 32), + layers_per_block=1, + latent_channels=4, + norm_num_groups=8, + sample_size=16, + ) + + +def _small_wan_vae(vae_cls: type[nn.Module] | None = None) -> nn.Module: + if vae_cls is None: + from diffusers.models.autoencoders import AutoencoderKLWan + + vae_cls = AutoencoderKLWan + + return vae_cls( + base_dim=8, + decoder_base_dim=8, + z_dim=4, + dim_mult=[1, 2], + num_res_blocks=1, + attn_scales=[], + temperal_downsample=[False], + latents_mean=[0.0] * 4, + latents_std=[1.0] * 4, + scale_factor_temporal=1, + scale_factor_spatial=2, + ) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.no_grad() +def test_image_vae_install_preserves_lossless_path_and_state_dict() -> None: + from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches + from vllm_omni.diffusion.vae_optimizations.image import ( + FusedGroupNormSiLU, + FusedUpsample2xConv2d, + maybe_optimize_image_vae, + ) + + torch.manual_seed(0) + vae = _small_image_vae().to(device="cuda", dtype=torch.bfloat16).eval() + latents = torch.randn(1, 4, 8, 8, device="cuda", dtype=torch.bfloat16) + parameter_names = {name for name, _ in vae.named_parameters()} + state_dict = {name: value.clone() for name, value in vae.state_dict().items()} + reference = vae.decode(latents).sample + + maybe_optimize_image_vae(vae) + maybe_optimize_image_vae(vae) + + assert {name for name, _ in vae.named_parameters()} == parameter_names + vae.load_state_dict(state_dict, strict=True) + assert any(isinstance(module, FusedGroupNormSiLU) for module in vae.modules()) + assert any(isinstance(module, FusedUpsample2xConv2d) for module in vae.modules()) + assert torch.equal(vae.decode(latents).sample, reference) + + with use_vae_fast_path(vae, True): + fast = vae.decode(latents).sample + + torch.testing.assert_close(fast.float(), reference.float(), atol=0.1, rtol=0) + assert any( + buffer is not None + for module in vae.modules() + for name, buffer in module._buffers.items() + if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"} + ) + clear_pipeline_vae_fast_path_caches(_PipelineWithVae(vae)) + assert all( + buffer is None + for module in vae.modules() + for name, buffer in module._buffers.items() + if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"} + ) + with use_vae_fast_path(vae, True): + torch.testing.assert_close(vae.decode(latents).sample.float(), reference.float(), atol=0.1, rtol=0) + assert torch.equal(vae.decode(latents).sample, reference) + assert vae.decoder.conv_in.weight.is_contiguous() + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.no_grad() +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_channels_last_group_norm_silu_matches_reference(dtype: torch.dtype) -> None: + from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import group_norm_silu_4d + + torch.manual_seed(0) + x = torch.randn(1, 64, 16, 16, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last) + weight = torch.randn(64, device="cuda", dtype=dtype) + bias = torch.randn(64, device="cuda", dtype=dtype) + + actual = group_norm_silu_4d(x, weight, bias, num_groups=8, eps=1e-5) + expected = F.silu(F.group_norm(x, 8, weight, bias)) + + assert actual is not None + atol, rtol = (1e-5, 1e-5) if dtype == torch.float32 else (7e-2, 2e-2) + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.no_grad() +@pytest.mark.parametrize( + "activation_dtype,affine_dtype,atol,rtol", + [ + (torch.float32, torch.float32, 1e-5, 1e-5), + (torch.bfloat16, torch.float32, 1.5e-1, 3e-2), + ], +) +def test_wan_rmsnorm_silu_matches_reference( + activation_dtype: torch.dtype, + affine_dtype: torch.dtype, + atol: float, + rtol: float, +) -> None: + from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import wan_rmsnorm_silu + + torch.manual_seed(0) + x = torch.randn(1, 96, 3, 10, 14, device="cuda", dtype=activation_dtype).contiguous( + memory_format=torch.channels_last_3d + ) + gamma = torch.randn(96, 1, 1, 1, device="cuda", dtype=affine_dtype) + bias = torch.randn_like(gamma) + normalized = F.normalize(x.float() if activation_dtype != torch.float32 else x, dim=1).to(activation_dtype) + expected = F.silu(normalized * 96**0.5 * gamma + bias) + + actual = wan_rmsnorm_silu(x, gamma, bias) + + assert actual is not None + assert actual.stride() == x.stride() + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.no_grad() +def test_wan_vae_install_preserves_lossless_path_and_state_dict() -> None: + from vllm_omni.diffusion.vae_optimizations.wan import ( + FusedWanRMSNormSiLU, + maybe_optimize_wan_vae, + ) + + torch.manual_seed(0) + vae = _small_wan_vae().cuda().eval() + latents = torch.randn(1, 4, 1, 8, 8, device="cuda") + parameter_names = {name for name, _ in vae.named_parameters()} + state_dict = {name: value.clone() for name, value in vae.state_dict().items()} + reference = vae.decode(latents).sample + + maybe_optimize_wan_vae(vae) + maybe_optimize_wan_vae(vae) + + assert {name for name, _ in vae.named_parameters()} == parameter_names + vae.load_state_dict(state_dict, strict=True) + assert any(isinstance(module, FusedWanRMSNormSiLU) for module in vae.modules()) + assert torch.equal(vae.decode(latents).sample, reference) + + with use_vae_fast_path(vae, True): + fast = vae.decode(latents).sample + + torch.testing.assert_close(fast, reference, atol=5e-4, rtol=1e-4) + assert torch.equal(vae.decode(latents).sample, reference) + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("parallel_mode", ["auto", "spatial_shard_height", "spatial_shard_width"]) +def test_wan_fast_path_skips_spatial_shard_mode(monkeypatch: pytest.MonkeyPatch, parallel_mode: str) -> None: + import vllm_omni.diffusion.vae_optimizations.wan as wan_optimizations + + monkeypatch.setattr(wan_optimizations, "_HAS_TRITON", True) + vae = _small_wan_vae() + vae.distributed_executor = _DistributedExecutorConfig(parallel_mode) + + wan_optimizations.maybe_optimize_wan_vae(vae) + + assert not any(isinstance(module, wan_optimizations.FusedWanRMSNormSiLU) for module in vae.modules()) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.no_grad() +@pytest.mark.parametrize( + ("precision", "decode_dtype"), + [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)], +) +def test_wan_decode_precision_keeps_encode_fp32(precision: str, decode_dtype: torch.dtype) -> None: + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan + from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision + + torch.manual_seed(0) + vae = _small_wan_vae(DistributedAutoencoderKLWan).to(device="cuda", dtype=torch.bfloat16).eval() + configure_wan_decode_precision(vae, precision) + + assert next(vae.encoder.parameters()).dtype == torch.float32 + assert next(vae.quant_conv.parameters()).dtype == torch.float32 + assert next(vae.decoder.parameters()).dtype == decode_dtype + assert next(vae.post_quant_conv.parameters()).dtype == decode_dtype + + _, grid_spec = vae.tile_split(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=decode_dtype)) + assert grid_spec.output_dtype == decode_dtype + + encoded = vae.encode(torch.randn(1, 3, 1, 16, 16, device="cuda", dtype=torch.bfloat16)) + decoded = vae.decode(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=torch.float32)) + + assert encoded.latent_dist.parameters.dtype == torch.float32 + assert decoded.sample.dtype == decode_dtype diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 04a29edd1..51541aa96 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -76,6 +76,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: assert args.cfg_parallel_size == 4 +def test_serve_parser_accepts_wan_decode_precision() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"]) + + assert args.vae_decode_precision == "bf16" + assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16" + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 8997b12e1..e4f44265d 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -573,6 +573,7 @@ class _DiffusionConfigProjection: fa_deterministic: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False + vae_decode_precision: Literal["fp32", "fp16", "bf16"] | None = None mask_strategy_file_path: str | None = None skip_time_steps: int = 15 VSA_sparsity: float = 0.0 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5b6b9a01a..d7e1a5a4c 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -401,6 +401,7 @@ class StageDeployConfig: step_execution: bool | None = None vae_use_slicing: bool | None = None vae_use_tiling: bool | None = None + vae_decode_precision: str | None = None boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 73051692d..a229fc73e 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -711,6 +711,9 @@ class OmniDiffusionConfig: # VAE memory optimization parameters vae_use_slicing: bool = False vae_use_tiling: bool = False + # Optional Wan decode-only precision. When set, the VAE encoder remains + # FP32 while decoder-side modules use the selected precision. + vae_decode_precision: str | None = None # STA (Sliding Tile Attention) parameters mask_strategy_file_path: str | None = None @@ -923,6 +926,11 @@ class OmniDiffusionConfig: ) def __post_init__(self): + if self.vae_decode_precision not in (None, "fp32", "fp16", "bf16"): + raise ValueError( + "vae_decode_precision must be one of {'fp32', 'fp16', 'bf16'} or None, " + f"got {self.vae_decode_precision!r}" + ) if self.diffusion_compile_granularity not in {"regional", "full"}: raise ValueError( "diffusion_compile_granularity must be 'regional' or 'full', " diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 44c6efe66..817e67e49 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -106,7 +106,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): split_dims=(3, 4), grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), tile_spec=tile_spec, - output_dtype=self.dtype, + output_dtype=getattr(self, "_vllm_decode_dtype", self.dtype), ) return tiletask_list, grid_spec diff --git a/vllm_omni/diffusion/vae_optimizations/__init__.py b/vllm_omni/diffusion/vae_optimizations/__init__.py new file mode 100644 index 000000000..43e76cb50 --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Optional VAE decoder optimizations.""" + +from vllm_omni.diffusion.vae_optimizations.gate import ( + use_pipeline_vae_fast_path, + use_vae_fast_path, +) +from vllm_omni.diffusion.vae_optimizations.optimize import ( + clear_pipeline_vae_fast_path_caches, + optimize_pipeline_vaes, +) + +__all__ = [ + "clear_pipeline_vae_fast_path_caches", + "optimize_pipeline_vaes", + "use_pipeline_vae_fast_path", + "use_vae_fast_path", +] diff --git a/vllm_omni/diffusion/vae_optimizations/gate.py b/vllm_omni/diffusion/vae_optimizations/gate.py new file mode 100644 index 000000000..333b34dd6 --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/gate.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decode-scoped gates for optional, non-bit-exact VAE fast paths.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from weakref import WeakKeyDictionary + +from torch import nn + +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + + +class VaeFastPathGate: + """Mutable flag shared by wrappers installed on one VAE instance.""" + + __slots__ = ("enabled",) + + def __init__(self) -> None: + self.enabled = False + + +_VAE_FAST_PATH_GATES: WeakKeyDictionary[nn.Module, VaeFastPathGate] = WeakKeyDictionary() + + +def register_vae_fast_path_gate(vae: nn.Module, gate: VaeFastPathGate) -> None: + """Associate an installed fast path with its owning VAE.""" + + _VAE_FAST_PATH_GATES[vae] = gate + + +@contextmanager +def use_vae_fast_path(vae: nn.Module, enabled: bool) -> Iterator[None]: + """Enable an installed fast path for one decode and always restore it.""" + + gate = _VAE_FAST_PATH_GATES.get(vae) + if gate is None: + yield + return + + previous_enabled = gate.enabled + gate.enabled = enabled + try: + yield + finally: + gate.enabled = previous_enabled + + +@contextmanager +def use_pipeline_vae_fast_path(pipeline: nn.Module, enabled: bool) -> Iterator[None]: + """Apply one request's quality choice to every declared VAE component.""" + + vaes = ModuleDiscovery.discover(pipeline).vaes + with ExitStack() as stack: + for vae in vaes: + stack.enter_context(use_vae_fast_path(vae, enabled)) + yield + + +__all__ = [ + "VaeFastPathGate", + "register_vae_fast_path_gate", + "use_pipeline_vae_fast_path", + "use_vae_fast_path", +] diff --git a/vllm_omni/diffusion/vae_optimizations/image.py b/vllm_omni/diffusion/vae_optimizations/image.py new file mode 100644 index 000000000..6be34d92f --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/image.py @@ -0,0 +1,354 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quality-gated CUDA fast paths for Diffusers-style 2D KL VAE decoders.""" + +from __future__ import annotations + +from types import MethodType +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.vae_optimizations.gate import ( + VaeFastPathGate, + register_vae_fast_path_gate, +) + +logger = init_logger(__name__) + +try: + from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import ( + group_norm_silu_4d, + group_norm_silu_rows, + ) + + _HAS_TRITON = True +except ImportError: # pragma: no cover - exercised on builds without Triton + _HAS_TRITON = False + + +class FusedGroupNormSiLU(nn.Module): + """Quality-gated GroupNorm + SiLU with an exact eager fallback.""" + + def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None: + super().__init__() + self.num_groups = norm.num_groups + self.num_channels = norm.num_channels + self.eps = norm.eps + self.affine = norm.affine + # Register parameters directly so state-dict FQNs do not change. + self.weight = norm.weight + self.bias = norm.bias + self._vllm_gate = gate + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self._vllm_gate.enabled and x.dim() == 4 and not torch.compiler.is_compiling(): + output = group_norm_silu_4d( + x, + self.weight, + self.bias, + self.num_groups, + self.eps, + apply_silu=True, + ) + if output is not None: + return output + return F.silu(F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps)) + + +def _install_norm_silu(decoder: nn.Module, resnet_cls: type[nn.Module], gate: VaeFastPathGate) -> int: + def is_fusable(norm: object) -> bool: + return type(norm) is nn.GroupNorm and norm.affine and norm.weight is not None and norm.bias is not None + + count = 0 + for module in decoder.modules(): + if ( + type(module) is resnet_cls + and module.time_emb_proj is None + and module.time_embedding_norm in ("default", "group") + and module.upsample is None + and module.downsample is None + and isinstance(module.nonlinearity, nn.SiLU) + and is_fusable(module.norm1) + and is_fusable(module.norm2) + ): + module.norm1 = FusedGroupNormSiLU(module.norm1, gate) + module.norm2 = FusedGroupNormSiLU(module.norm2, gate) + module.nonlinearity = nn.Identity() + count += 2 + + if ( + type(getattr(decoder, "conv_norm_out", None)) is nn.GroupNorm + and isinstance(getattr(decoder, "conv_act", None), nn.SiLU) + and is_fusable(decoder.conv_norm_out) + ): + decoder.conv_norm_out = FusedGroupNormSiLU(decoder.conv_norm_out, gate) + decoder.conv_act = nn.Identity() + count += 1 + return count + + +_UPSAMPLE_TAP_MAP = {0: (2,), 1: (1, 2), 2: (0, 1), 3: (0,)} + + +def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor: + weight = conv.weight.detach().float() + output_channels, input_channels = weight.shape[:2] + folded = weight.new_zeros(input_channels, output_channels, 4, 4) + for row in range(4): + for column in range(4): + value = weight.new_zeros(output_channels, input_channels) + for source_row in _UPSAMPLE_TAP_MAP[row]: + for source_column in _UPSAMPLE_TAP_MAP[column]: + value += weight[:, :, source_row, source_column] + folded[:, :, row, column] = value.t() + folded = folded.to(conv.weight.dtype) + if conv.weight.is_contiguous(memory_format=torch.channels_last): + folded = folded.contiguous(memory_format=torch.channels_last) + return folded.to(conv.weight.device) + + +class FusedUpsample2xConv2d(nn.Module): + """Quality-gated nearest-2x + Conv2d replacement using ConvTranspose2d.""" + + def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None: + super().__init__() + # Keep only ``conv`` registered so checkpoint parameter names stay + # ``...upsamplers.N.conv.*``. The original module is an eager fallback. + object.__setattr__(self, "_original", upsample) + self.conv = upsample.conv + self.channels = upsample.channels + self._vllm_gate = gate + self.register_buffer("_fused_weight", None, persistent=False) + + def forward( + self, + hidden_states: torch.Tensor, + output_size: int | None = None, + *args: Any, + **kwargs: Any, + ) -> torch.Tensor: + if not self._vllm_gate.enabled or output_size is not None or hidden_states.shape[1] != self.channels: + return self._original(hidden_states, output_size=output_size, *args, **kwargs) + + weight = self._fused_weight + if weight is None or weight.device != self.conv.weight.device or weight.dtype != self.conv.weight.dtype: + weight = _fold_upsample2x_conv2d_weight(self.conv) + self._fused_weight = weight + return F.conv_transpose2d(hidden_states, weight, self.conv.bias, stride=2, padding=1) + + +def _install_fused_upsample(decoder: nn.Module, upsample_cls: type[nn.Module], gate: VaeFastPathGate) -> int: + count = 0 + for block in decoder.up_blocks: + upsamplers = getattr(block, "upsamplers", None) + if not upsamplers: + continue + for index, upsample in enumerate(upsamplers): + if type(upsample) is not upsample_cls: + continue + conv = getattr(upsample, "conv", None) + if ( + upsample.use_conv + and not upsample.use_conv_transpose + and upsample.interpolate + and upsample.norm is None + and upsample.name == "conv" + and type(conv) is nn.Conv2d + and conv.kernel_size == (3, 3) + and conv.stride == (1, 1) + and conv.padding == (1, 1) + and conv.dilation == (1, 1) + and conv.groups == 1 + and conv.padding_mode == "zeros" + ): + upsamplers[index] = FusedUpsample2xConv2d(upsample, gate) + count += 1 + return count + + +def _fold_attention_value_projection(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor]: + value_weight = module.to_v.weight.detach().float() + value_bias = module.to_v.bias.detach().float() + output_weight = module.to_out[0].weight.detach().float() + output_bias = module.to_out[0].bias.detach().float() + dtype = module.to_v.weight.dtype + return ( + (output_weight @ value_weight).to(dtype).contiguous(), + (output_weight @ value_bias + output_bias).to(dtype), + ) + + +def _attention_fast_forward( + self: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + temb: torch.Tensor | None = None, + **cross_attention_kwargs: Any, +) -> torch.Tensor: + if ( + not self._vllm_gate.enabled + or encoder_hidden_states is not None + or attention_mask is not None + or temb is not None + or hidden_states.ndim != 4 + ): + return type(self).forward( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + temb=temb, + **cross_attention_kwargs, + ) + + residual = hidden_states + batch_size, channels, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels) + + if self.group_norm is not None: + norm = self.group_norm + normalized = None + if not torch.compiler.is_compiling(): + normalized = group_norm_silu_rows( + hidden_states, + norm.weight, + norm.bias, + norm.num_groups, + norm.eps, + apply_silu=False, + ) + hidden_states = normalized if normalized is not None else norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + folded_weight = self._vllm_folded_value_weight + folded_bias = self._vllm_folded_value_bias + if ( + folded_weight is None + or folded_bias is None + or folded_weight.device != hidden_states.device + or folded_weight.dtype != self.to_v.weight.dtype + ): + folded_weight, folded_bias = _fold_attention_value_projection(self) + self._vllm_folded_value_weight = folded_weight + self._vllm_folded_value_bias = folded_bias + value = F.linear(hidden_states, folded_weight, folded_bias) + + output = F.scaled_dot_product_attention(query.unsqueeze(1), key.unsqueeze(1), value.unsqueeze(1)) + output = output.squeeze(1).to(query.dtype) + output = self.to_out[1](output) + output = output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2) + if self.residual_connection: + output = output + residual + return output / self.rescale_output_factor + + +def _attention_fast_compatible(module: nn.Module, attention_cls: type, processor_cls: type) -> bool: + return ( + type(module) is attention_cls + and isinstance(module.processor, processor_cls) + and module.heads == 1 + and module.scale_qk + and module.spatial_norm is None + and not module.norm_cross + and getattr(module, "norm_q", None) is None + and getattr(module, "norm_k", None) is None + and getattr(module, "add_k_proj", None) is None + and type(module.to_q) is nn.Linear + and type(module.to_k) is nn.Linear + and type(module.to_v) is nn.Linear + and type(module.to_out[0]) is nn.Linear + and isinstance(module.to_out[1], nn.Dropout) + and module.to_v.bias is not None + and module.to_out[0].bias is not None + ) + + +def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor: + use_channels_last = self._vllm_gate.enabled + if use_channels_last != self._vllm_channels_last: + memory_format = torch.channels_last if use_channels_last else torch.contiguous_format + self.to(memory_format=memory_format) + self._vllm_channels_last = use_channels_last + logger.debug( + "%s decoder switched to %s layout", + self._vllm_label, + "channels_last" if use_channels_last else "contiguous", + ) + return type(self).forward(self, *args, **kwargs) + + +def maybe_optimize_image_vae(vae: nn.Module) -> nn.Module: + """Install image VAE decoder fast paths when the structure is fully supported.""" + + if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False): + return vae + + from diffusers.models.attention_processor import Attention, AttnProcessor2_0 + from diffusers.models.autoencoders import AutoencoderKL + from diffusers.models.autoencoders.vae import Decoder + from diffusers.models.resnet import ResnetBlock2D + from diffusers.models.upsampling import Upsample2D + + supported_types: tuple[type[nn.Module], ...] = (AutoencoderKL,) + try: + from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 + + supported_types += (AutoencoderKLFlux2,) + except ImportError: # pragma: no cover - older supported Diffusers builds + pass + + if not isinstance(vae, supported_types): + return vae + decoder = getattr(vae, "decoder", None) + if type(decoder) is not Decoder: + return vae + + attention_modules = [ + module for module in decoder.modules() if _attention_fast_compatible(module, Attention, AttnProcessor2_0) + ] + attention_count = sum(isinstance(module, Attention) for module in decoder.modules()) + if len(attention_modules) != attention_count: + logger.info( + "%s has %d/%d attention blocks without a layout-safe rewrite; skipping VAE fast paths", + type(vae).__name__, + attention_count - len(attention_modules), + attention_count, + ) + return vae + + gate = VaeFastPathGate() + decoder._vllm_gate = gate + decoder._vllm_label = type(vae).__name__ + decoder._vllm_channels_last = False + decoder.forward = MethodType(_decoder_layout_forward, decoder) + upsample_count = _install_fused_upsample(decoder, Upsample2D, gate) + for module in attention_modules: + module._vllm_gate = gate + module.register_buffer("_vllm_folded_value_weight", None, persistent=False) + module.register_buffer("_vllm_folded_value_bias", None, persistent=False) + module.forward = MethodType(_attention_fast_forward, module) + norm_count = _install_norm_silu(decoder, ResnetBlock2D, gate) + register_vae_fast_path_gate(vae, gate) + vae._vllm_vae_fast_path_installed = True + logger.info( + "%s installed quality-gated VAE decoder fast paths " + "(%d upsamplers, %d attention blocks, %d GroupNorm+SiLU fusions)", + type(vae).__name__, + upsample_count, + len(attention_modules), + norm_count, + ) + return vae + + +__all__ = [ + "FusedGroupNormSiLU", + "FusedUpsample2xConv2d", + "maybe_optimize_image_vae", +] diff --git a/vllm_omni/diffusion/vae_optimizations/optimize.py b/vllm_omni/diffusion/vae_optimizations/optimize.py new file mode 100644 index 000000000..e290e79f7 --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/optimize.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pipeline-level installation for optional VAE optimizations.""" + +from __future__ import annotations + +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery +from vllm_omni.diffusion.vae_optimizations.image import maybe_optimize_image_vae +from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision +from vllm_omni.diffusion.vae_optimizations.wan import maybe_optimize_wan_vae +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + +_FOLDED_CACHE_NAMES = ( + "_fused_weight", + "_vllm_folded_value_weight", + "_vllm_folded_value_bias", +) + + +def optimize_pipeline_vaes(pipeline: nn.Module, od_config: OmniDiffusionConfig) -> None: + """Configure declared VAE components once after checkpoint loading.""" + + precision = getattr(od_config, "vae_decode_precision", None) + vaes = ModuleDiscovery.discover(pipeline).vaes + for vae in vaes: + # An explicit precision override is a serving contract, so failures on + # a recognized Wan VAE must abort startup instead of being hidden. + configure_wan_decode_precision(vae, precision) + + if not current_omni_platform.is_cuda(): + return + + for vae in vaes: + for optimizer in (maybe_optimize_image_vae, maybe_optimize_wan_vae): + try: + optimizer(vae) + except Exception: + # A wrapper may already have installed its exact fallback + # before a later optional rewrite failed. With no registered + # enabled gate, those wrappers remain on the reference path. + logger.warning( + "Failed to install optional VAE fast paths on %s; using the reference path", + type(vae).__name__, + exc_info=True, + ) + + +def clear_pipeline_vae_fast_path_caches(pipeline: nn.Module) -> None: + """Release lazily folded buffers before the worker enters sleep mode.""" + + cleared = 0 + for vae in ModuleDiscovery.discover(pipeline).vaes: + if not getattr(vae, "_vllm_vae_fast_path_installed", False): + continue + for module in vae.modules(): + for name in _FOLDED_CACHE_NAMES: + if name in module._buffers and module._buffers[name] is not None: + setattr(module, name, None) + cleared += 1 + if cleared: + logger.debug("Released %d lazy VAE fast-path buffers before sleep", cleared) + + +__all__ = ["clear_pipeline_vae_fast_path_caches", "optimize_pipeline_vaes"] diff --git a/vllm_omni/diffusion/vae_optimizations/precision.py b/vllm_omni/diffusion/vae_optimizations/precision.py new file mode 100644 index 000000000..b5686cd57 --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/precision.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decode-only precision configuration for structurally split Wan VAEs.""" + +from __future__ import annotations + +from types import MethodType +from typing import Any + +import torch +from torch import nn +from vllm.logger import init_logger + +logger = init_logger(__name__) + +PRECISION_TO_DTYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +def _decode_with_precision( + self: nn.Module, + latents: torch.Tensor, + *args: Any, + **kwargs: Any, +) -> Any: + return type(self).decode( + self, + latents.to(dtype=self._vllm_decode_dtype), + *args, + **kwargs, + ) + + +def _encode_in_fp32( + self: nn.Module, + samples: torch.Tensor, + *args: Any, + **kwargs: Any, +) -> Any: + return type(self).encode( + self, + samples.to(dtype=torch.float32), + *args, + **kwargs, + ) + + +def configure_wan_decode_precision(vae: nn.Module, precision: str | None) -> nn.Module: + """Keep Wan encode in FP32 while materializing only decode modules lower.""" + + if precision is None: + return vae + + from diffusers.models.autoencoders import AutoencoderKLWan + + if not isinstance(vae, AutoencoderKLWan): + return vae + try: + decode_dtype = PRECISION_TO_DTYPE[precision] + except KeyError as exc: + raise ValueError( + f"Unsupported vae_decode_precision={precision!r}; expected one of {sorted(PRECISION_TO_DTYPE)}" + ) from exc + + encoder = getattr(vae, "encoder", None) + quant_conv = getattr(vae, "quant_conv", None) + decoder = getattr(vae, "decoder", None) + post_quant_conv = getattr(vae, "post_quant_conv", None) + if not all(isinstance(module, nn.Module) for module in (encoder, decoder, post_quant_conv)): + logger.info("Wan VAE lacks separate encoder/decoder modules; skipping decode-only precision") + return vae + + encoder.to(dtype=torch.float32) + if isinstance(quant_conv, nn.Module): + quant_conv.to(dtype=torch.float32) + decoder.to(dtype=decode_dtype) + post_quant_conv.to(dtype=decode_dtype) + vae.encode = MethodType(_encode_in_fp32, vae) + vae._vllm_decode_dtype = decode_dtype + vae.decode = MethodType(_decode_with_precision, vae) + logger.info( + "Wan VAE configured with FP32 encode and %s decode", + precision, + ) + return vae + + +__all__ = ["PRECISION_TO_DTYPE", "configure_wan_decode_precision"] diff --git a/vllm_omni/diffusion/vae_optimizations/triton_group_norm_silu.py b/vllm_omni/diffusion/vae_optimizations/triton_group_norm_silu.py new file mode 100644 index 000000000..1b1010d42 --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/triton_group_norm_silu.py @@ -0,0 +1,237 @@ +# ruff: noqa: N803 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Channels-last two-pass GroupNorm(+SiLU) Triton kernels. + +The kernels keep the channel dimension innermost throughout a 2D VAE +decoder. Statistics are accumulated in FP32, then a separate apply pass +performs the affine transform and optional SiLU epilogue. Unsupported inputs +return ``None`` so callers can fail closed to the reference implementation. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} +_MAX_CHANNELS = 2048 + + +@triton.jit +def _gn_partial_rows_kernel( + x_ptr, + psum_ptr, + psq_ptr, + rows, + rows_per_prog, + C: tl.constexpr, + BLOCK_R: tl.constexpr, +): + chunk = tl.program_id(0).to(tl.int64) + batch = tl.program_id(1).to(tl.int64) + num_chunks = tl.num_programs(0) + row0 = chunk * rows_per_prog + cols = tl.arange(0, C) + acc_sum = tl.zeros((C,), tl.float32) + acc_sq = tl.zeros((C,), tl.float32) + x_base = x_ptr + batch * rows * C + for row_offset in range(0, rows_per_prog, BLOCK_R): + row_indices = row0 + row_offset + tl.arange(0, BLOCK_R) + mask = row_indices < rows + x = tl.load( + x_base + row_indices[:, None] * C + cols[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + acc_sum += tl.sum(x, 0) + acc_sq += tl.sum(x * x, 0) + output_offsets = (batch * num_chunks + chunk) * C + cols + tl.store(psum_ptr + output_offsets, acc_sum) + tl.store(psq_ptr + output_offsets, acc_sq) + + +@triton.jit +def _gn_finalize_kernel( + psum_ptr, + psq_ptr, + weight_ptr, + bias_ptr, + scale_shift_ptr, + num_chunks, + group_numel, + eps, + channels, + CPG: tl.constexpr, + BLOCK_K: tl.constexpr, +): + group = tl.program_id(0).to(tl.int64) + batch = tl.program_id(1).to(tl.int64) + cols = group * CPG + tl.arange(0, CPG) + acc_sum = tl.zeros((), tl.float32) + acc_sq = tl.zeros((), tl.float32) + for chunk0 in range(0, num_chunks, BLOCK_K): + chunks = chunk0 + tl.arange(0, BLOCK_K) + mask = chunks < num_chunks + offsets = (batch * num_chunks + chunks)[:, None] * channels + cols[None, :] + acc_sum += tl.sum(tl.load(psum_ptr + offsets, mask=mask[:, None], other=0.0)) + acc_sq += tl.sum(tl.load(psq_ptr + offsets, mask=mask[:, None], other=0.0)) + mean = acc_sum / group_numel + variance = acc_sq / group_numel - mean * mean + variance = tl.maximum(variance, 0.0) + reciprocal_std = tl.rsqrt(variance + eps) + weight = tl.load(weight_ptr + cols).to(tl.float32) + bias = tl.load(bias_ptr + cols).to(tl.float32) + scale = weight * reciprocal_std + shift = bias - mean * scale + tl.store(scale_shift_ptr + batch * 2 * channels + cols, scale) + tl.store(scale_shift_ptr + (batch * 2 + 1) * channels + cols, shift) + + +@triton.jit +def _gn_apply_rows_kernel( + x_ptr, + scale_shift_ptr, + y_ptr, + rows, + C: tl.constexpr, + BLOCK_R: tl.constexpr, + SILU: tl.constexpr, +): + program = tl.program_id(0).to(tl.int64) + batch = tl.program_id(1).to(tl.int64) + cols = tl.arange(0, C) + scale = tl.load(scale_shift_ptr + batch * 2 * C + cols) + shift = tl.load(scale_shift_ptr + (batch * 2 + 1) * C + cols) + row_indices = program * BLOCK_R + tl.arange(0, BLOCK_R) + mask = row_indices < rows + offsets = batch * rows * C + row_indices[:, None] * C + cols[None, :] + x = tl.load(x_ptr + offsets, mask=mask[:, None], other=0.0).to(tl.float32) + y = x * scale[None, :] + shift[None, :] + if SILU: + y = y * tl.sigmoid(y) + tl.store(y_ptr + offsets, y.to(y_ptr.dtype.element_ty), mask=mask[:, None]) + + +def _group_norm_silu_rows( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + eps: float, + apply_silu: bool, +) -> torch.Tensor: + batch_size, rows, channels = x.shape + channels_per_group = channels // num_groups + block_rows = max(1, 8192 // channels) + rows_per_program = block_rows * 32 + num_chunks = triton.cdiv(rows, rows_per_program) + partial_sum = torch.empty((batch_size, num_chunks, channels), device=x.device, dtype=torch.float32) + partial_sq = torch.empty_like(partial_sum) + with torch.cuda.device(x.device): + _gn_partial_rows_kernel[(num_chunks, batch_size)]( + x, + partial_sum, + partial_sq, + rows, + rows_per_program, + C=channels, + BLOCK_R=block_rows, + num_warps=4, + ) + scale_shift = torch.empty((batch_size, 2, channels), device=x.device, dtype=torch.float32) + block_chunks = max(1, min(4096 // channels_per_group, triton.next_power_of_2(num_chunks))) + _gn_finalize_kernel[(num_groups, batch_size)]( + partial_sum, + partial_sq, + weight, + bias, + scale_shift, + num_chunks, + rows * channels_per_group, + eps, + channels, + CPG=channels_per_group, + BLOCK_K=block_chunks, + num_warps=4, + ) + output = torch.empty_like(x) + _gn_apply_rows_kernel[(triton.cdiv(rows, block_rows), batch_size)]( + x, + scale_shift, + output, + rows, + C=channels, + BLOCK_R=block_rows, + SILU=apply_silu, + num_warps=4, + ) + return output + + +def _is_supported( + x: torch.Tensor, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + num_groups: int, +) -> bool: + if not (x.is_cuda and x.numel() > 0 and not torch.is_grad_enabled()): + return False + if x.requires_grad or x.dtype not in _SUPPORTED_DTYPES: + return False + if weight is None or bias is None: + return False + channels = x.shape[1] if x.dim() == 4 else x.shape[-1] + if weight.shape != (channels,) or bias.shape != (channels,): + return False + if not ( + weight.device == bias.device == x.device + and weight.dtype == bias.dtype == x.dtype + and weight.is_contiguous() + and bias.is_contiguous() + ): + return False + if num_groups < 1 or channels % num_groups != 0: + return False + return triton.next_power_of_2(channels) == channels and channels <= _MAX_CHANNELS + + +def group_norm_silu_4d( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + eps: float, + apply_silu: bool = True, +) -> torch.Tensor | None: + """Run GroupNorm(+SiLU) on a channels-last ``(N, C, H, W)`` tensor.""" + + if x.dim() != 4 or not _is_supported(x, weight, bias, num_groups): + return None + batch_size, channels, height, width = x.shape + if not (channels > 1 and (height > 1 or width > 1) and x.is_contiguous(memory_format=torch.channels_last)): + return None + rows = x.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels) + output = _group_norm_silu_rows(rows, weight, bias, num_groups, eps, apply_silu) + return output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2) + + +def group_norm_silu_rows( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + eps: float, + apply_silu: bool = True, +) -> torch.Tensor | None: + """Run GroupNorm(+SiLU) over contiguous ``(N, rows, C)`` input.""" + + if x.dim() != 3 or not x.is_contiguous(): + return None + if not _is_supported(x, weight, bias, num_groups): + return None + return _group_norm_silu_rows(x, weight, bias, num_groups, eps, apply_silu) + + +__all__ = ["group_norm_silu_4d", "group_norm_silu_rows"] diff --git a/vllm_omni/diffusion/vae_optimizations/triton_wan_rmsnorm_silu.py b/vllm_omni/diffusion/vae_optimizations/triton_wan_rmsnorm_silu.py new file mode 100644 index 000000000..af7d597fd --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/triton_wan_rmsnorm_silu.py @@ -0,0 +1,160 @@ +# ruff: noqa: N803 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Channels-last-3D Wan VAE RMSNorm(+SiLU) Triton kernel.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} +_MAX_CHANNELS = 1024 + + +@triton.jit +def _wan_rmsnorm_silu_kernel( + x_ptr, + gamma_ptr, + bias_ptr, + output_ptr, + channels: tl.constexpr, + time_size, + height_size, + width_size, + x_stride_batch, + x_stride_channel, + x_stride_time, + x_stride_height, + x_stride_width, + output_stride_batch, + output_stride_channel, + output_stride_time, + output_stride_height, + output_stride_width, + rms_scale, + eps, + has_bias: tl.constexpr, + BLOCK_C: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + channel_offsets = tl.arange(0, BLOCK_C) + channel_mask = channel_offsets < channels + + width = row % width_size + remainder = row // width_size + height = remainder % height_size + remainder = remainder // height_size + time = remainder % time_size + batch = remainder // time_size + + x_base = batch * x_stride_batch + time * x_stride_time + height * x_stride_height + width * x_stride_width + output_base = ( + batch * output_stride_batch + + time * output_stride_time + + height * output_stride_height + + width * output_stride_width + ) + + x = tl.load( + x_ptr + x_base + channel_offsets * x_stride_channel, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + norm = tl.sqrt(tl.sum(x * x, axis=0)) + inverse_norm = 1.0 / tl.maximum(norm, eps) + + # Match the eager dtype boundaries: normalize and scale in activation + # dtype, affine in the promoted dtype, then SiLU in FP32. + y = (x * inverse_norm).to(x_ptr.dtype.element_ty) + gamma = tl.load(gamma_ptr + channel_offsets, mask=channel_mask, other=1.0) + y = (y * rms_scale).to(x_ptr.dtype.element_ty) + y = (y.to(tl.float32) * gamma.to(tl.float32)).to(output_ptr.dtype.element_ty) + if has_bias: + bias = tl.load(bias_ptr + channel_offsets, mask=channel_mask, other=0.0) + y = (y.to(tl.float32) + bias.to(tl.float32)).to(output_ptr.dtype.element_ty) + y = y.to(tl.float32) + y = y * tl.sigmoid(y) + tl.store( + output_ptr + output_base + channel_offsets * output_stride_channel, + y, + mask=channel_mask, + ) + + +def _affine_supported(x: torch.Tensor, affine: torch.Tensor) -> bool: + return ( + affine.is_cuda + and affine.device == x.device + and (affine.dtype == x.dtype or affine.dtype == torch.float32) + and affine.numel() == x.shape[1] + ) + + +def can_use_wan_rmsnorm_silu( + x: torch.Tensor, + gamma: torch.Tensor, + bias: torch.Tensor | None, +) -> bool: + """Return whether ``wan_rmsnorm_silu`` supports this exact input.""" + + return ( + x.is_cuda + and not torch.is_grad_enabled() + and not x.requires_grad + and x.dtype in _SUPPORTED_DTYPES + and x.ndim == 5 + and x.numel() > 0 + and 0 < x.shape[1] <= _MAX_CHANNELS + and x.is_contiguous(memory_format=torch.channels_last_3d) + and _affine_supported(x, gamma) + and (bias is None or _affine_supported(x, bias)) + ) + + +def wan_rmsnorm_silu( + x: torch.Tensor, + gamma: torch.Tensor, + bias: torch.Tensor | None = None, + rms_scale: float | None = None, + eps: float = 1e-12, +) -> torch.Tensor | None: + """Fuse ``SiLU(F.normalize(x, dim=1) * scale * gamma + bias)``.""" + + if not can_use_wan_rmsnorm_silu(x, gamma, bias): + return None + + batch_size, channels, time_size, height_size, width_size = x.shape + output_dtype = torch.promote_types(x.dtype, gamma.dtype) + output = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=output_dtype) + block_channels = triton.next_power_of_2(channels) + num_warps = 1 if block_channels <= 64 else 4 if block_channels <= 512 else 8 + gamma = gamma.reshape(channels).contiguous() + has_bias = bias is not None + bias = gamma if bias is None else bias.reshape(channels).contiguous() + if rms_scale is None: + rms_scale = channels**0.5 + + with torch.cuda.device(x.device): + _wan_rmsnorm_silu_kernel[(batch_size * time_size * height_size * width_size,)]( + x, + gamma, + bias, + output, + channels, + time_size, + height_size, + width_size, + *x.stride(), + *output.stride(), + float(rms_scale), + eps, + has_bias, + BLOCK_C=block_channels, + num_warps=num_warps, + ) + return output + + +__all__ = ["can_use_wan_rmsnorm_silu", "wan_rmsnorm_silu"] diff --git a/vllm_omni/diffusion/vae_optimizations/wan.py b/vllm_omni/diffusion/vae_optimizations/wan.py new file mode 100644 index 000000000..ac2fda3e4 --- /dev/null +++ b/vllm_omni/diffusion/vae_optimizations/wan.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quality-gated CUDA fast path for Diffusers Wan VAE decoders.""" + +from __future__ import annotations + +from types import MethodType +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.vae_optimizations.gate import ( + VaeFastPathGate, + register_vae_fast_path_gate, +) + +logger = init_logger(__name__) + +try: + from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import ( + wan_rmsnorm_silu, + ) + + _HAS_TRITON = True +except ImportError: # pragma: no cover - exercised on builds without Triton + _HAS_TRITON = False + + +class FusedWanRMSNormSiLU(nn.Module): + """Quality-gated Wan RMSNorm + SiLU with an exact eager fallback.""" + + def __init__(self, norm: nn.Module, gate: VaeFastPathGate) -> None: + super().__init__() + self.gamma = norm.gamma + self.bias = norm.bias + self.scale = float(norm.scale) + self._vllm_gate = gate + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self._vllm_gate.enabled and not torch.compiler.is_compiling(): + bias = self.bias if isinstance(self.bias, torch.Tensor) else None + output = wan_rmsnorm_silu(x, self.gamma, bias, rms_scale=self.scale) + if output is not None: + return output + + needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( + token in str(x.dtype) for token in ("float4_", "float8_") + ) + normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=1).to(x.dtype) + return F.silu(normalized * self.scale * self.gamma + self.bias) + + +def _plain_silu(activation: object) -> bool: + return isinstance(activation, nn.SiLU) and not activation.inplace + + +def _fusable_norm(norm: object, norm_cls: type) -> bool: + return ( + type(norm) is norm_cls + and getattr(norm, "channel_first", False) + and isinstance(getattr(norm, "gamma", None), torch.Tensor) + ) + + +def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None: + from diffusers.models.autoencoders.autoencoder_kl_wan import WanResidualBlock, WanRMS_norm + + residual_blocks = [module for module in decoder.modules() if isinstance(module, WanResidualBlock)] + eligible = [ + module + for module in residual_blocks + if type(module) is WanResidualBlock + and _plain_silu(module.nonlinearity) + and _fusable_norm(module.norm1, WanRMS_norm) + and _fusable_norm(module.norm2, WanRMS_norm) + ] + if len(eligible) != len(residual_blocks): + logger.info( + "Wan VAE has %d/%d non-standard residual blocks; skipping fast path", + len(residual_blocks) - len(eligible), + len(residual_blocks), + ) + return None + if not ( + _fusable_norm(getattr(decoder, "norm_out", None), WanRMS_norm) + and _plain_silu(getattr(decoder, "nonlinearity", None)) + ): + logger.info("Wan VAE has a non-standard output head; skipping fast path") + return None + + count = 0 + for module in eligible: + module.norm1 = FusedWanRMSNormSiLU(module.norm1, gate) + module.norm2 = FusedWanRMSNormSiLU(module.norm2, gate) + module.nonlinearity = nn.Identity() + count += 2 + decoder.norm_out = FusedWanRMSNormSiLU(decoder.norm_out, gate) + decoder.nonlinearity = nn.Identity() + return count + 1 + + +def _set_wan_decoder_memory_format(decoder: nn.Module, *, channels_last: bool) -> None: + for module in decoder.modules(): + if isinstance(module, nn.Conv3d): + module.to(memory_format=torch.channels_last_3d if channels_last else torch.contiguous_format) + elif isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)): + module.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format) + + +def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor: + use_channels_last = self._vllm_gate.enabled + if use_channels_last != self._vllm_channels_last: + _set_wan_decoder_memory_format(self, channels_last=use_channels_last) + self._vllm_channels_last = use_channels_last + logger.debug( + "Wan VAE decoder switched to %s layout", + "channels_last_3d" if use_channels_last else "contiguous", + ) + return type(self).forward(self, *args, **kwargs) + + +def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module: + """Install the Wan decoder fast path when its full structure is supported.""" + + if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False): + return vae + + from diffusers.models.autoencoders import AutoencoderKLWan + from diffusers.models.autoencoders.autoencoder_kl_wan import WanDecoder3d + + if not isinstance(vae, AutoencoderKLWan): + return vae + decoder = getattr(vae, "decoder", None) + if type(decoder) is not WanDecoder3d: + return vae + + distributed_executor = getattr(vae, "distributed_executor", None) + parallel_mode = getattr(distributed_executor, "parallel_mode", "tile") + if parallel_mode in ("auto", "spatial_shard_height", "spatial_shard_width"): + logger.info("Wan spatial-shard decode is configured; skipping quality-gated VAE fast path") + return vae + + gate = VaeFastPathGate() + norm_count = _install_norm_silu(decoder, gate) + if norm_count is None: + return vae + decoder._vllm_gate = gate + decoder._vllm_channels_last = False + decoder.forward = MethodType(_decoder_layout_forward, decoder) + register_vae_fast_path_gate(vae, gate) + vae._vllm_vae_fast_path_installed = True + logger.info( + "Wan VAE installed quality-gated decoder fast path (%d RMSNorm+SiLU fusions)", + norm_count, + ) + return vae + + +__all__ = ["FusedWanRMSNormSiLU", "maybe_optimize_wan_vae"] diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 52c2c646d..b12eb3f42 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -43,6 +43,10 @@ from vllm_omni.diffusion.offloader import get_offload_backend from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob +from vllm_omni.diffusion.vae_optimizations import ( + optimize_pipeline_vaes, + use_pipeline_vae_fast_path, +) from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.diffusion.worker.utils import ( @@ -243,6 +247,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): custom_pipeline_name=custom_pipeline_name, device=self.device, ) + # Keep precision rematerialization in the tagged weights pool + # so sleep level 1 can offload every resulting parameter. + optimize_pipeline_vaes(self.pipeline, self.od_config) time_after_load = time.perf_counter() logger.info( @@ -499,6 +506,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if require_request_batch_support and not getattr(self.pipeline, "supports_request_batch", False): raise RuntimeError(f"{type(self.pipeline).__name__} does not support request-batch forward.") + quality = getattr(reqs[0].sampling_params, "quality", None) + if any(getattr(req.sampling_params, "quality", None) != quality for req in reqs[1:]): + raise RuntimeError("A diffusion worker batch cannot mix request quality levels.") + use_fast_vae = quality == "high" + # Use no_grad() for HSDP compatibility, inference_mode() otherwise for # better perf. HSDP2's fully_shard pre-forward hooks need tensor version # counters, which inference tensors do not track. @@ -523,7 +535,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config): with record_function(record_name): - raw_outputs = self.pipeline.forward(batch) + with use_pipeline_vae_fast_path(self.pipeline, use_fast_vae): + raw_outputs = self.pipeline.forward(batch) outputs = _normalize_pipeline_outputs( raw_outputs, expected_count=len(reqs), @@ -774,7 +787,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if should_decode: clear_pipeline_stage_durations(self.pipeline) - result = self.pipeline.post_decode(req) + with use_pipeline_vae_fast_path( + self.pipeline, + getattr(req.sampling, "quality", None) == "high", + ): + result = self.pipeline.post_decode(req) if result is not None: self._attach_stepwise_metrics( req, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index a3d1c791f..404d329ca 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -576,6 +576,14 @@ class DiffusionWorker: usage_before = allocator.get_current_usage() + if self.model_runner is not None and self.model_runner.pipeline is not None: + from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches + + # Folded weights are allocated lazily on the request path rather + # than in the tagged weight pool. Drop them before sleep so they + # neither remain resident nor need a separate CPU backup. + clear_pipeline_vae_fast_path_caches(self.model_runner.pipeline) + if level == 2 and self.model_runner is not None: if hasattr(self.model_runner, "graph_runners"): self.model_runner.graph_runners.clear() diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 02ec3587b..8a06ec050 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -491,6 +491,7 @@ class OrchestratorArgs: step_execution: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False + vae_decode_precision: str | None = None enable_multithread_weight_load: bool = True num_weight_load_threads: int = 4 enable_cpu_offload: bool = False diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index fc80cfa93..e4685dd28 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1026,6 +1026,7 @@ class AsyncOmniEngine: "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0), "vae_use_slicing": kwargs.get("vae_use_slicing", False), "vae_use_tiling": kwargs.get("vae_use_tiling", False), + "vae_decode_precision": kwargs.get("vae_decode_precision", None), "cache_backend": cache_backend, "cache_config": cache_config, "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 7fb9cb8cd..c7a6eeaae 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -629,6 +629,15 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).", ) + omni_config_group.add_argument( + "--vae-decode-precision", + choices=("fp32", "fp16", "bf16"), + default=None, + help=( + "Optional Wan decode-only VAE precision. The encoder remains FP32; " + "unset preserves the pipeline's existing precision behavior." + ), + ) # Parallel weight loading (faster diffusion startup) omni_config_group.add_argument(