From 4d3eed152a697412c966d2ac97e0009b92490b5e Mon Sep 17 00:00:00 2001 From: "Y. Fisher" Date: Wed, 13 May 2026 22:22:43 +0800 Subject: [PATCH] [Feat][Config] Support additional_config for diffusion worker (#3020) Signed-off-by: KexiongYu Signed-off-by: Y. Fisher --- .../hunyuan_image3/README.md | 13 +++++ .../hunyuan_image3/end2end.py | 35 ++++++++++++ tests/diffusion/test_diffusion_worker.py | 54 +++++++++++++++++- .../test_async_omni_diffusion_config.py | 55 +++++++++++++++++++ tests/test_diffusion_config_propagation.py | 6 ++ vllm_omni/diffusion/data.py | 8 +++ .../diffusion/worker/diffusion_worker.py | 36 ++++++++++-- vllm_omni/engine/async_omni_engine.py | 6 ++ 8 files changed, 207 insertions(+), 6 deletions(-) diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md index 98908ace0..6db4cbec9 100644 --- a/examples/offline_inference/hunyuan_image3/README.md +++ b/examples/offline_inference/hunyuan_image3/README.md @@ -100,12 +100,25 @@ python examples/offline_inference/hunyuan_image3/end2end.py \ --prompts "A cute cat" ``` +## Additional Config + +You can pass diffusion worker `additional_config` from the offline example as a JSON object. +This maps to the upstream vLLM `VllmConfig.additional_config` platform extension field: +https://docs.vllm.ai/en/stable/api/vllm/config/#vllm.config.VllmConfig.additional_config + +```bash +python end2end.py --modality text2img \ + --prompts "A cute cat" \ + --additional-config '{"torchair_graph_config":{"enabled":true}}' +``` + ## Key Arguments | Argument | Description | | :--- | :--- | | `--deploy-config` | Preferred config path for unified deploy YAMLs. | | `--stage-configs-path` | Legacy stage config path, kept only for compatibility. Prefer `--deploy-config`. | +| `--additional-config` | JSON object forwarded to diffusion worker `additional_config`. | | `--modality` | Offline-only convenience flag. One of `text2img`, `img2img`, `img2text`, `text2text`. It selects prompt formatting, internal `mode`, and default deploy config for this script. Online serving uses `--deploy-config` plus the endpoint and, for chat completions, request `modalities` instead. | | `--steps` | Number of diffusion inference steps for image generation. | | `--guidance-scale` | Classifier-free guidance scale for image generation. | diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index f4eb8cffc..5232568f1 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -14,6 +14,7 @@ Usage: """ import argparse +import json import os from pathlib import Path @@ -116,6 +117,17 @@ def parse_args(): parser.add_argument("--log-stats", action="store_true", default=False) parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") + parser.add_argument( + "--additional-config", + type=str, + default=None, + help=( + "JSON object forwarded to Omni/additional_config, for example " + '\'{"torchair_graph_config":{"enabled":true}}\'. Different platforms may support different ' + "configs. Make sure the configs are valid for the platform you are using. " + "Contents must be hashable." + ), + ) from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults @@ -123,9 +135,27 @@ def parse_args(): return parser.parse_args() +def parse_additional_config(raw_value: str | None) -> dict | None: + """Parse a JSON string into an additional_config mapping.""" + if raw_value is None: + return None + + try: + additional_config = json.loads(raw_value) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid --additional-config JSON: {exc}") from exc + + if additional_config is None: + return None + if not isinstance(additional_config, dict): + raise ValueError(f"--additional-config must decode to a JSON object, got {type(additional_config).__name__}") + return additional_config + + def main(): args = parse_args() os.makedirs(args.output, exist_ok=True) + additional_config = parse_additional_config(args.additional_config) # Determine task for prompt formatting from modality + bot behavior. task = _MODALITY_TASK_MAP[args.modality] @@ -160,6 +190,9 @@ def main(): "init_timeout": args.init_timeout, "enforce_eager": args.enforce_eager, } + + if additional_config is not None: + omni_kwargs["additional_config"] = additional_config if deploy_config is not None: omni_kwargs["deploy_config"] = deploy_config else: @@ -262,6 +295,8 @@ def main(): print(f" Output size: {args.width}x{args.height}") if args.image_path: print(f" Input image: {args.image_path}") + if additional_config is not None: + print(f" Additional config: {additional_config}") print(f" Prompts: {prompts}") print(f"{'=' * 60}\n") diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 88b4d8c06..8ad6c2734 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -14,7 +14,11 @@ import pytest import torch from pytest_mock import MockerFixture -from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker, _make_diffusion_vllm_model_config +from vllm_omni.diffusion.worker.diffusion_worker import ( + DiffusionWorker, + _create_diffusion_worker_vllm_config, + _make_diffusion_vllm_model_config, +) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] @@ -348,4 +352,50 @@ class TestDiffusionWorkerWakeUp: # buffer2 should NOT be restored since it wasn't saved mock_buffer2.data.copy_.assert_not_called() - assert bool(result) is True + assert result is True + + +class TestWorkerVllmConfigAdditionalConfig: + """Test worker-side VllmConfig construction with additional_config.""" + + def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture): + mock_vllm_config = mocker.Mock() + mock_vllm_cls = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", + return_value=mock_vllm_config, + ) + od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}}) + + result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) + + assert result is mock_vllm_config + assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config + + def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config( + self, mocker: MockerFixture + ): + mock_vllm_config = mocker.Mock() + mock_vllm_cls = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", + side_effect=[ + TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"), + mock_vllm_config, + ], + ) + od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}}) + + result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) + + assert result is mock_vllm_config + assert mock_vllm_cls.call_count == 2 + assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config + + def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture): + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", + side_effect=TypeError("additional_config values must be hashable"), + ) + od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}}) + + with pytest.raises(TypeError, match="hashable"): + _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 79a6b376b..ab2229673 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -287,3 +289,56 @@ def test_serve_cli_accepts_diffusion_attention_backend(): assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None assert diffusion_attention_config.default.backend == "FLASH_ATTN" + + +def test_serve_cli_accepts_additional_config(): + """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config.""" + parser = FlexibleArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "Qwen/Qwen-Image", + "--omni", + "--additional-config", + '{"torchair_graph_config":{"enabled":true}}', + ] + ) + + stage_cfg = _create_default_diffusion_stage_cfg(args)[0] + engine_args = stage_cfg["engine_args"] + + assert args.additional_config == {"torchair_graph_config": {"enabled": True}} + assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}} + + +def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker): + """Ensure YAML/deploy stage resolution forwards top-level additional_config.""" + fake_diffusion_stage = SimpleNamespace( + stage_type="diffusion", + engine_args=SimpleNamespace(), + ) + fake_llm_stage = SimpleNamespace( + stage_type="llm", + engine_args=SimpleNamespace(), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage]), + ) + + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + engine._strip_single_engine_args = lambda kwargs: kwargs + + _, stage_configs = engine._resolve_stage_configs( + "dummy-model", + { + "stage_configs_path": "dummy.yaml", + "additional_config": {"torchair_graph_config": {"enabled": True}}, + }, + ) + + assert not hasattr(stage_configs[0].engine_args, "additional_config") + assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}} diff --git a/tests/test_diffusion_config_propagation.py b/tests/test_diffusion_config_propagation.py index eeb3505ef..389dc9c98 100644 --- a/tests/test_diffusion_config_propagation.py +++ b/tests/test_diffusion_config_propagation.py @@ -119,3 +119,9 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): assert od_config.supports_multimodal_inputs is True assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES + + +def test_additional_config_roundtrip(): + additional_config = {"torchair_graph_config": {"enabled": True}} + od = _roundtrip_diffusion_config(model="x", additional_config=additional_config) + assert od.additional_config == additional_config diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 72776dcc3..9cbc7b877 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -552,6 +552,7 @@ class OmniDiffusionConfig: # Omni configuration (injected from stage config) omni_kv_config: dict[str, Any] = field(default_factory=dict) + additional_config: dict[str, Any] = field(default_factory=dict) profiler_config: "ProfilerConfig | dict[str, Any] | None" = None @@ -650,6 +651,13 @@ class OmniDiffusionConfig: self.profiler_config = ProfilerConfig(**self.profiler_config) + if self.additional_config is None: + self.additional_config = {} + elif isinstance(self.additional_config, Mapping): + self.additional_config = dict(self.additional_config) + else: + raise TypeError(f"additional_config must be a mapping or None, got {type(self.additional_config)!r}") + # Convert parallel_config dict/DictConfig to DiffusionParallelConfig # Use Mapping to handle both plain dicts and OmegaConf DictConfig if isinstance(self.parallel_config, Mapping): diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 72ea62045..3909fd2d4 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -126,6 +126,37 @@ def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[No yield +def _is_unexpected_additional_config_type_error(exc: TypeError) -> bool: + """Return True only for constructor rejections of the additional_config kwarg.""" + message = str(exc) + return "unexpected keyword argument" in message and "additional_config" in message + + +def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDiffusionConfig) -> VllmConfig: + """Create a worker-local VllmConfig while preserving additional_config when supported.""" + config_kwargs: dict[str, Any] = { + "compilation_config": CompilationConfig(), + "device_config": DeviceConfig(device=device), + } + if od_config.additional_config: + config_kwargs["additional_config"] = od_config.additional_config + + try: + return VllmConfig(**config_kwargs) + except TypeError as exc: + if not _is_unexpected_additional_config_type_error(exc): + raise + + logger.debug("Worker-local VllmConfig does not accept additional_config in constructor: %s", exc) + config_kwargs.pop("additional_config", None) + vllm_config = VllmConfig(**config_kwargs) + try: + setattr(vllm_config, "additional_config", dict(od_config.additional_config)) + except Exception as set_exc: # pragma: no cover - defensive for older vLLM builds + logger.warning("Failed to attach additional_config to worker VllmConfig: %s", set_exc) + return vllm_config + + class DiffusionWorker: """ A worker that manages GPU infrastructure and delegates to the model runner. @@ -188,10 +219,7 @@ class DiffusionWorker: # Create vllm_config for parallel configuration. Pass explicit device_config # so DeviceConfig does not rely on current_platform in worker subprocesses. - vllm_config = VllmConfig( - compilation_config=CompilationConfig(), - device_config=DeviceConfig(device=self.device), - ) + vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config) vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 151671051..ed4eed808 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1490,6 +1490,7 @@ class AsyncOmniEngine: "max_num_seqs": kwargs.get("max_num_seqs") or 1, "parallel_config": parallel_config, "model_class_name": kwargs.get("model_class_name", None), + "additional_config": kwargs.get("additional_config", None), "step_execution": kwargs.get("step_execution", False), "vae_use_slicing": kwargs.get("vae_use_slicing", False), "vae_use_tiling": kwargs.get("vae_use_tiling", False), @@ -1644,6 +1645,11 @@ class AsyncOmniEngine: continue if not hasattr(cfg, "engine_args") or cfg.engine_args is None: cfg.engine_args = OmegaConf.create({}) + additional_config = kwargs.get("additional_config") + if additional_config is not None: + current_additional_config = getattr(cfg.engine_args, "additional_config", None) + if current_additional_config in (None, {}): + cfg.engine_args.additional_config = additional_config if kwargs.get("lora_path") is not None: if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None: cfg.engine_args.lora_path = kwargs["lora_path"]