diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 3d2caeabf..d75316c0e 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -751,6 +751,7 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path): " - stage_id: 0", " diffusion_compile_granularity: regional", " diffusion_compile_dynamic: false", + " diffusion_compile_vae: false", ] ) ) @@ -762,13 +763,16 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path): cli_overrides={ "diffusion_compile_granularity": "full", "diffusion_compile_dynamic": True, + "diffusion_compile_vae": True, }, ).stage_by_id(0) assert configured_stage.diffusion_config.diffusion_compile_granularity == "regional" assert configured_stage.diffusion_config.diffusion_compile_dynamic is False + assert configured_stage.diffusion_config.diffusion_compile_vae is False assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full" assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True + assert overridden_stage.diffusion_config.diffusion_compile_vae is True def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic(): @@ -776,6 +780,11 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic(): omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false") +def test_structured_diffusion_config_rejects_non_boolean_compile_vae(): + with pytest.raises(ValidationError, match="diffusion_compile_vae"): + omni_config_module._DiffusionConfigProjection(diffusion_compile_vae="false") + + def test_structured_diffusion_config_rejects_invalid_compile_granularity(): with pytest.raises(ValidationError, match="diffusion_compile_granularity"): omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block") diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py index 447d8c75f..b689ed95f 100644 --- a/tests/diffusion/test_compile.py +++ b/tests/diffusion/test_compile.py @@ -5,7 +5,7 @@ import pytest import torch.nn as nn import vllm_omni.diffusion.compile as compile_module -from vllm_omni.diffusion.compile import regionally_compile +from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -74,3 +74,50 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat regionally_compile(model, dynamic=True) assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards + + +def test_compile_callable_falls_back_once_after_lazy_failure(monkeypatch, caplog): + compiled_calls = 0 + eager_calls = 0 + + def eager(value): + nonlocal eager_calls + eager_calls += 1 + return value + 1 + + def compile_fn(fn, **kwargs): + assert fn is eager + assert kwargs == {"mode": "default"} + + def compiled(value): + nonlocal compiled_calls + compiled_calls += 1 + raise RuntimeError("lazy failure") + + return compiled + + monkeypatch.setattr(compile_module.torch, "compile", compile_fn) + callable_fn, activated = compile_callable_with_fallback(eager, label="test callable", mode="default") + + assert activated + assert callable_fn(1) == 2 + assert callable_fn(2) == 3 + assert compiled_calls == 1 + assert eager_calls == 2 + assert "disabling compilation and retrying eagerly" in caplog.text + + +def test_compile_callable_keeps_eager_on_setup_failure(monkeypatch, caplog): + def eager(value): + return value + 1 + + def compile_fn(*args, **kwargs): + raise RuntimeError("setup failure") + + monkeypatch.setattr(compile_module.torch, "compile", compile_fn) + callable_fn, activated = compile_callable_with_fallback(eager, label="test callable") + + assert not activated + assert callable_fn is eager + assert callable_fn(1) == 2 + assert "continuing eagerly" in caplog.text diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index a0c6acb02..a3c1723d4 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -515,16 +515,19 @@ class TestDiffusionCompileConfig: assert config.diffusion_compile_granularity == "regional" assert config.diffusion_compile_dynamic is True + assert config.diffusion_compile_vae is True def test_from_kwargs_preserves_compile_controls(self) -> None: config = OmniDiffusionConfig.from_kwargs( model="test", diffusion_compile_granularity="full", diffusion_compile_dynamic=False, + diffusion_compile_vae=False, ) assert config.diffusion_compile_granularity == "full" assert config.diffusion_compile_dynamic is False + assert config.diffusion_compile_vae is False def test_config_rejects_invalid_compile_granularity(self) -> None: with pytest.raises(ValueError, match="diffusion_compile_granularity"): @@ -534,6 +537,10 @@ class TestDiffusionCompileConfig: with pytest.raises(TypeError, match="diffusion_compile_dynamic"): OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false") + def test_config_rejects_non_boolean_compile_vae(self) -> None: + with pytest.raises(TypeError, match="diffusion_compile_vae"): + OmniDiffusionConfig(model="test", diffusion_compile_vae="false") + @pytest.mark.parametrize( "kwargs, feature", [ diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index a1f8440e3..4b4c9891d 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -130,6 +130,23 @@ class _CompileTrackingModel: return self +class _DecodeVAE(torch.nn.Module): + def decode(self, value): + return value + 1 + + +class _PipelineWithVAE(torch.nn.Module): + _dit_modules = [] + _encoder_modules = [] + _vae_modules = ["vae", "vae_alias"] + _resident_modules = [] + + def __init__(self): + super().__init__() + self.vae = _DecodeVAE() + self.vae_alias = self.vae + + def _make_request(): sampling_params = SimpleNamespace( generator=None, @@ -320,6 +337,31 @@ def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypa assert "lazy compilation errors" in caplog.text +@pytest.mark.core_model +@pytest.mark.cpu +def test_compile_vae_decodes_uses_discovery_and_deduplicates_aliases(monkeypatch): + runner = object.__new__(DiffusionModelRunner) + runner.pipeline = _PipelineWithVAE() + compile_calls = [] + + def compile_fn(fn, **kwargs): + compile_calls.append((fn, kwargs)) + return lambda *args, **call_kwargs: fn(*args, **call_kwargs) + 10 + + monkeypatch.setattr(model_runner_module.torch, "compile", compile_fn) + + DiffusionModelRunner._compile_vae_decodes(runner) + + assert len(compile_calls) == 1 + assert compile_calls[0][1] == { + "mode": "default", + "fullgraph": False, + "dynamic": None, + } + assert runner.pipeline.vae.decode(1) == 12 + assert runner.pipeline.vae_alias.decode is runner.pipeline.vae.decode + + @pytest.mark.core_model @pytest.mark.cpu def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index ce9f0ebe1..66283b47d 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -971,6 +971,44 @@ class TestDiffusionEngine: admitted_request = engine.add_req_and_wait_for_response.call_args.args[0] assert admitted_request.request_id == "dummy_req_id" + def test_compile_warmup_repeats_through_normal_request_path(self, mocker: MockerFixture) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace( + model_class_name="mock_model", + diffusion_load_format="default", + diffusion_compile_vae=True, + enforce_eager=False, + ) + engine.add_req_and_wait_for_response = mocker.Mock( + side_effect=[DiffusionOutput(output=None), DiffusionOutput(output=None)] + ) + + engine._dummy_run() + + requests = [call.args[0] for call in engine.add_req_and_wait_for_response.call_args_list] + assert len(requests) == 2 + assert requests[0].is_dummy_run() + assert not requests[1].is_dummy_run() + assert requests[1].prompt == {"prompt": "A detailed image."} + assert requests[1].sampling_params.need_kv_receive is False + assert requests[1].sampling_params.num_inference_steps == 1 + + def test_compile_real_path_warmup_failure_is_non_fatal(self, mocker: MockerFixture, caplog) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace( + model_class_name="mock_model", + diffusion_load_format="default", + diffusion_compile_vae=True, + enforce_eager=False, + ) + engine.add_req_and_wait_for_response = mocker.Mock( + side_effect=[DiffusionOutput(output=None), RuntimeError("prewarm failure")] + ) + + engine._dummy_run() + + assert "serving will continue" in caplog.text + class TestStepScheduler: def setup_method(self) -> None: diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 64b586bf2..9614760c1 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -87,6 +87,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): "enforce_eager": None, "diffusion_compile_granularity": None, "diffusion_compile_dynamic": None, + "diffusion_compile_vae": None, } )[0] @@ -98,6 +99,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): assert stage_cfg["engine_args"]["enforce_eager"] is False assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional" assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True + assert stage_cfg["engine_args"]["diffusion_compile_vae"] is True def test_default_stage_config_propagates_ulysses_mode(): @@ -315,7 +317,7 @@ def test_serve_cli_forwards_distributed_offload_residency(): def test_serve_cli_accepts_diffusion_compile_controls(): - """Ensure both compile controls reach the diffusion stage.""" + """Ensure all compile controls reach the diffusion stage.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) @@ -328,6 +330,7 @@ def test_serve_cli_accepts_diffusion_compile_controls(): "--diffusion-compile-granularity", "full", "--no-diffusion-compile-dynamic", + "--no-diffusion-compile-vae", ] ) @@ -336,8 +339,10 @@ def test_serve_cli_accepts_diffusion_compile_controls(): assert args.diffusion_compile_granularity == "full" assert args.diffusion_compile_dynamic is False + assert args.diffusion_compile_vae is False assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full" assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False + assert stage_cfg["engine_args"]["diffusion_compile_vae"] is False def test_serve_cli_accepts_diffusion_attention_backend(): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 8997b12e1..88e453208 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -570,6 +570,7 @@ class _DiffusionConfigProjection: pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) + diffusion_compile_vae: bool = Field(default=True, strict=True) fa_deterministic: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5b6b9a01a..a66bef899 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -394,6 +394,7 @@ class StageDeployConfig: # Diffusion execution, cache, and VAE behavior. diffusion_compile_granularity: str | None = None diffusion_compile_dynamic: bool | None = None + diffusion_compile_vae: bool | None = None fa_deterministic: bool | None = None cache_backend: str | None = None cache_config: dict[str, Any] | None = None diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 227d4f330..9c4615043 100644 --- a/vllm_omni/diffusion/compile.py +++ b/vllm_omni/diffusion/compile.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable from typing import Any import torch @@ -9,6 +10,59 @@ from vllm.logger import init_logger logger = init_logger(__name__) +class _FallbackCompiledCallable: + """Run a compiled callable until it fails, then permanently use eager.""" + + def __init__( + self, + eager: Callable[..., Any], + compiled: Callable[..., Any], + *, + label: str, + ) -> None: + self._eager = eager + self._compiled: Callable[..., Any] | None = compiled + self._label = label + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + compiled = self._compiled + if compiled is None: + return self._eager(*args, **kwargs) + try: + return compiled(*args, **kwargs) + except Exception as exc: + self._compiled = None + logger.warning( + "%s failed during lazy torch.compile execution (%s); disabling compilation and retrying eagerly.", + self._label, + exc, + ) + return self._eager(*args, **kwargs) + + +def compile_callable_with_fallback( + fn: Callable[..., Any], + *, + label: str, + **compile_kwargs: Any, +) -> tuple[Callable[..., Any], bool]: + """Compile ``fn`` with a fail-closed eager fallback. + + The boolean return value reports whether lazy compilation was activated. + Synchronous setup failures return the original callable unchanged. + """ + try: + compiled = torch.compile(fn, **compile_kwargs) + except Exception as exc: + logger.warning( + "%s torch.compile setup failed (%s); continuing eagerly.", + label, + exc, + ) + return fn, False + return _FallbackCompiledCallable(fn, compiled, label=label), True + + def _matches_repeated_block( name: str, module: nn.Module, diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 73051692d..7b3a6e247 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -726,6 +726,9 @@ class OmniDiffusionConfig: # provide its own setup_compile() implementation. diffusion_compile_granularity: str = "regional" diffusion_compile_dynamic: bool = True + # Compile discovered VAE ``decode`` callables whenever generic diffusion + # torch.compile is enabled. Runtime failures permanently fall back to eager. + diffusion_compile_vae: bool = True # Parallel weight loading (for faster diffusion model startup) enable_multithread_weight_load: bool = True @@ -930,6 +933,8 @@ class OmniDiffusionConfig: ) if not isinstance(self.diffusion_compile_dynamic, bool): raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}") + if not isinstance(self.diffusion_compile_vae, bool): + raise TypeError(f"diffusion_compile_vae must be a bool, got {type(self.diffusion_compile_vae)!r}") self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) if self.omni_kv_config is None: diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index c9b4bd72f..3b6a420d1 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -56,6 +56,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) _ASYNC_OUTPUT_TIMEOUT = 30.0 # seconds +_VAE_COMPILE_WARMUP_REQUEST_ID = "vae_compile_warmup_req_id" __all__ = [ "DiffusionEngine", @@ -851,28 +852,26 @@ class DiffusionEngine: num_inference_steps = 1 height = 512 width = 512 - prompt: OmniTextPrompt = {"prompt": "dummy run"} supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config) - if supports_image_input: - # Provide a dummy image input if the model supports it - color_format = image_color_format(self.od_config.model_class_name) - dummy_image = PIL.Image.new(color_format, (width, height)) - prompt.setdefault("multi_modal_data", {})["image"] = dummy_image - - if supports_audio_input: - audio_sr = 16000 - dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32) - prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio - num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) if num_frames <= 0: logger.info("Skipping dummy warmup run (num_frames=0)") return - req = OmniDiffusionRequest( - prompt=prompt, - request_id=DUMMY_DIFFUSION_REQUEST_ID, - sampling_params=OmniDiffusionSamplingParams( + + def build_request(request_id: str, prompt_text: str, *, real_path: bool) -> OmniDiffusionRequest: + prompt: OmniTextPrompt = {"prompt": prompt_text} + if supports_image_input: + color_format = image_color_format(self.od_config.model_class_name) + prompt.setdefault("multi_modal_data", {})["image"] = PIL.Image.new( + color_format, + (width, height), + ) + if supports_audio_input: + audio_sr = 16000 + prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32) + + sampling_params = OmniDiffusionSamplingParams( height=height, width=width, num_inference_steps=num_inference_steps, @@ -885,13 +884,47 @@ class DiffusionEngine: # Disable CFG for warmup to avoid triggering CFG parallel # validation when cfg_parallel_size > 1. extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0}, - ), - ) + ) + if real_path: + # Internal pretraffic must exercise ordinary pipeline branches + # without waiting for KV produced by an upstream stage. + sampling_params.need_kv_receive = False + return OmniDiffusionRequest( + prompt=prompt, + request_id=request_id, + sampling_params=sampling_params, + ) + + req = build_request(DUMMY_DIFFUSION_REQUEST_ID, "dummy run", real_path=False) logger.info("dummy run to warm up the model") output = self.add_req_and_wait_for_response(req) if output.error: raise RuntimeError(f"Dummy run failed: {output.error}") + compile_real_path = bool( + getattr(self.od_config, "diffusion_compile_vae", False) + and not getattr(self.od_config, "enforce_eager", True) + ) + if not compile_real_path: + return + + logger.info("real request-path warmup for compiled VAE decode") + real_path_req = build_request( + _VAE_COMPILE_WARMUP_REQUEST_ID, + "A detailed image.", + real_path=True, + ) + try: + real_path_output = self.add_req_and_wait_for_response(real_path_req) + except Exception as exc: + logger.warning("Real request-path compile warmup failed (%s); serving will continue.", exc) + return + if real_path_output.error: + logger.warning( + "Real request-path compile warmup returned an error (%s); serving will continue.", + real_path_output.error, + ) + def _submit_rpc( self, method: str, diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 52c2c646d..6e09e3936 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -28,7 +28,7 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import ( resolve_prompt_embed_cache_config, ) from vllm_omni.diffusion.cache.selector import get_cache_backend -from vllm_omni.diffusion.compile import regionally_compile +from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -40,6 +40,7 @@ from vllm_omni.diffusion.models.interface import ( supports_step_execution, ) from vllm_omni.diffusion.offloader import get_offload_backend +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery 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 @@ -188,6 +189,39 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): compile_dynamic, ) + def _compile_vae_decodes(self) -> None: + """Compile declared VAE decode callables with eager fallback.""" + assert self.pipeline is not None + modules = ModuleDiscovery.discover(self.pipeline) + compiled_count = 0 + for name, vae in zip(modules.vae_names, modules.vaes, strict=True): + decode = getattr(vae, "decode", None) + if not callable(decode): + logger.debug("Model runner: VAE component %s has no decode callable; skipping compile.", name) + continue + + label = f"VAE component {name}.decode" + compiled_decode, activated = compile_callable_with_fallback( + decode, + label=label, + mode="default", + fullgraph=False, + dynamic=None, + ) + if not activated: + continue + setattr(vae, "decode", compiled_decode) + compiled_count += 1 + + if compiled_count: + logger.info( + "Model runner: configured lazy torch.compile for %d VAE decode callable(s); " + "startup warmup will compile the real request path.", + compiled_count, + ) + else: + logger.info("Model runner: no compile-compatible VAE decode callable was discovered.") + def load_model( self, memory_pool_context_fn: Callable[[str], AbstractContextManager[Any]] | None = None, @@ -277,7 +311,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: if current_omni_platform.supports_torch_inductor(): - if hasattr(self.pipeline, "setup_compile"): + pipeline_owns_compile = hasattr(self.pipeline, "setup_compile") + if pipeline_owns_compile: try: self.pipeline.setup_compile() except Exception as exc: @@ -291,6 +326,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): transformer_attrs = ("transformer", "transformer_2") for attr_name in transformer_attrs: self._compile_transformer(attr_name) + + if self.od_config.diffusion_compile_vae and not pipeline_owns_compile: + self._compile_vae_decodes() else: logger.warning( "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 02ec3587b..3cede3cf7 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -485,6 +485,7 @@ class OrchestratorArgs: diffusion_attention_config: str | None = None diffusion_compile_granularity: str | None = None diffusion_compile_dynamic: bool | None = None + diffusion_compile_vae: bool | None = None cache_backend: str = "none" cache_config: str | None = None enable_cache_dit_summary: bool = False diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index fc80cfa93..d6e09e9d2 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1043,6 +1043,9 @@ class AsyncOmniEngine: "diffusion_compile_dynamic": ( True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"] ), + "diffusion_compile_vae": ( + True if kwargs.get("diffusion_compile_vae") is None else kwargs["diffusion_compile_vae"] + ), "fa_deterministic": bool(kwargs.get("fa_deterministic", False)), "boundary_ratio": kwargs.get("boundary_ratio", None), "flow_shift": kwargs.get("flow_shift", None), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 7fb9cb8cd..534067d5b 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -446,6 +446,15 @@ class OmniServeCommand(CLISubcommand): "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic." ), ) + omni_config_group.add_argument( + "--diffusion-compile-vae", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Compile discovered VAE decode callables when diffusion torch.compile is enabled. " + "Use --no-diffusion-compile-vae to keep VAE decode eager." + ), + ) omni_config_group.add_argument( "--fa-deterministic", dest="fa_deterministic",