[Feat][Config] Support additional_config for diffusion worker (#3020)

Signed-off-by: KexiongYu <yukexiong1@huawei.com>
Signed-off-by: Y. Fisher <yukexiong1@huawei.com>

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. |

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")
 

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)

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}}

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

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):

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

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"]